Merge pull request brianmario#501 from marcinbunsch/init-command

Allow setting the init command on the connection.

Author: sodabrew

README.md

@@ -156,10 +156,12 @@ Mysql2::Client.new(
   :local_infile = true/false,
   :secure_auth = true/false,
   :default_file = '/path/to/my.cfg',
-  :default_group = 'my.cfg section'
+  :default_group = 'my.cfg section',
+  :init_command => sql
   )
 ```
 
+
 ### SSL options
 
 Setting any of the following options will enable an SSL connection, but only if
@@ -251,6 +253,15 @@ the `:default_file` and `:default_group` paramters. For example:
 Mysql2::Client.new(:default_file => '/user/.my.cnf', :default_group => 'client')
 ```
 
+### init_command
+
+If you specify the init_command option, the SQL string you provide will be executed after the connection is established.
+If `:reconnect` is set to `true`, init_command will also be executed after a successful reconnect.
+It is useful if you want to provide session options which survive reconnection.
+
+``` ruby
+Mysql2::Client.new(:init_command => "SET @@SESSION.sql_mode = 'STRICT_ALL_TABLES'")
+```
 
 ## Cascading config
 

ext/mysql2/client.c

@@ -67,6 +67,7 @@ struct nogvl_connect_args {
   const char *user;
   const char *passwd;
   const char *db;
+  const char *init_command;
   unsigned int port;
   const char *unix_socket;
   unsigned long client_flag;
@@ -156,6 +157,10 @@ static void *nogvl_connect(void *ptr) {
   struct nogvl_connect_args *args = ptr;
   MYSQL *client;
 
+  if (args->init_command != NULL) {
+    mysql_options(args->mysql, MYSQL_INIT_COMMAND, args->init_command);
+  }
+
   client = mysql_real_connect(args->mysql, args->host,
                               args->user, args->passwd,
                               args->db, args->port, args->unix_socket,
@@ -322,7 +327,7 @@ static VALUE rb_mysql_info(VALUE self) {
   return rb_str;
 }
 
-static VALUE rb_connect(VALUE self, VALUE user, VALUE pass, VALUE host, VALUE port, VALUE database, VALUE socket, VALUE flags) {
+static VALUE rb_connect(VALUE self, VALUE user, VALUE pass, VALUE host, VALUE port, VALUE database, VALUE socket, VALUE flags, VALUE init_command) {
   struct nogvl_connect_args args;
   VALUE rv;
   GET_CLIENT(self);
@@ -335,6 +340,7 @@ static VALUE rb_connect(VALUE self, VALUE user, VALUE pass, VALUE host, VALUE po
   args.db = NIL_P(database) ? NULL : StringValuePtr(database);
   args.mysql = wrapper->client;
   args.client_flag = NUM2ULONG(flags);
+  args.init_command = NIL_P(init_command) ? NULL : StringValuePtr(init_command);
 
   rv = (VALUE) rb_thread_call_without_gvl(nogvl_connect, &args, RUBY_UBF_IO, 0);
   if (rv == Qfalse) {
@@ -1237,7 +1243,7 @@ void init_mysql2_client() {
   rb_define_private_method(cMysql2Client, "default_group=", set_read_default_group, 1);
   rb_define_private_method(cMysql2Client, "ssl_set", set_ssl_options, 5);
   rb_define_private_method(cMysql2Client, "initialize_ext", initialize_ext, 0);
-  rb_define_private_method(cMysql2Client, "connect", rb_connect, 7);
+  rb_define_private_method(cMysql2Client, "connect", rb_connect, 8);
 
   sym_id              = ID2SYM(rb_intern("id"));
   sym_version         = ID2SYM(rb_intern("version"));

lib/mysql2/client.rb

@@ -55,6 +55,7 @@ def initialize(opts = {})
       database = opts[:database] || opts[:dbname] || opts[:db]
       socket   = opts[:socket] || opts[:sock]
       flags    = opts[:flags] ? opts[:flags] | @query_options[:connect_flags] : @query_options[:connect_flags]
+      init_command = opts[:init_command]
 
       # Correct the data types before passing these values down to the C level
       user = user.to_s unless user.nil?
@@ -63,8 +64,9 @@ def initialize(opts = {})
       port = port.to_i unless port.nil?
       database = database.to_s unless database.nil?
       socket = socket.to_s unless socket.nil?
+      init_command = init_command.to_s unless init_command.nil?
 
-      connect user, pass, host, port, database, socket, flags
+      connect user, pass, host, port, database, socket, flags, init_command
     end
 
     def self.default_query_options

spec/mysql2/client_spec.rb

@@ -54,7 +54,7 @@ def connect *args
       end
     end
     client = klient.new :flags => Mysql2::Client::FOUND_ROWS
-    (client.connect_args.last.last & Mysql2::Client::FOUND_ROWS).should be_true
+    (client.connect_args.last[6] & Mysql2::Client::FOUND_ROWS).should be_true
   end
 
   it "should default flags to (REMEMBER_OPTIONS, LONG_PASSWORD, LONG_FLAG, TRANSACTIONS, PROTOCOL_41, SECURE_CONNECTION)" do
@@ -66,14 +66,58 @@ def connect *args
       end
     end
     client = klient.new
-    (client.connect_args.last.last & (Mysql2::Client::REMEMBER_OPTIONS |
+    (client.connect_args.last[6] & (Mysql2::Client::REMEMBER_OPTIONS |
                                      Mysql2::Client::LONG_PASSWORD |
                                      Mysql2::Client::LONG_FLAG |
                                      Mysql2::Client::TRANSACTIONS |
                                      Mysql2::Client::PROTOCOL_41 |
                                      Mysql2::Client::SECURE_CONNECTION)).should be_true
   end
 
+  it "should accept init_command" do
+    klient = Class.new(Mysql2::Client) do
+      attr_reader :connect_args
+      def connect *args
+        @connect_args ||= []
+        @connect_args << args
+      end
+    end
+    command = "SET @@session.something = 1"
+    client = klient.new :init_command => command
+    client.connect_args.last[7].should eq(command)
+  end
+
+  it "should execute init command" do
+    options = DatabaseCredentials['root'].dup
+    options[:init_command] = "SET @something = 'setting_value';"
+    client = Mysql2::Client.new(options)
+    result = client.query("SELECT @something;")
+    result.first['@something'].should eq('setting_value')
+  end
+
+  it "should send init_command after reconnect" do
+    pending "Ruby 2.1 has changed Timeout behavior." if RUBY_VERSION =~ /2.1/
+    options = DatabaseCredentials['root'].dup
+    options[:init_command] = "SET @something = 'setting_value';"
+    options[:reconnect] = true
+    client = Mysql2::Client.new(options)
+
+    result = client.query("SELECT @something;")
+    result.first['@something'].should eq('setting_value')
+
+    # simulate a broken connection
+    begin
+      Timeout.timeout(1) do
+        client.query("SELECT sleep(2)")
+      end
+    rescue Timeout::Error
+    end
+    client.ping.should be_false
+
+    result = client.query("SELECT @something;")
+    result.first['@something'].should eq('setting_value')
+  end
+
   it "should have a global default_query_options hash" do
     Mysql2::Client.should respond_to(:default_query_options)
   end