Support connect attributes and set the script name to program_name attribute if it is not specified. (brianmario#760)

Author: kubo

README.md

@@ -203,6 +203,7 @@ Mysql2::Client.new(
   :read_timeout = seconds,
   :write_timeout = seconds,
   :connect_timeout = seconds,
+  :connect_attrs = {:program_name => $PROGRAM_NAME, ...},
   :reconnect = true/false,
   :local_infile = true/false,
   :secure_auth = true/false,

ext/mysql2/client.c

@@ -400,7 +400,20 @@ static VALUE rb_mysql_get_ssl_cipher(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) {
+#ifdef HAVE_CONST_MYSQL_OPT_CONNECT_ATTR_ADD
+static int opt_connect_attr_add_i(VALUE key, VALUE value, VALUE arg)
+{
+  mysql_client_wrapper *wrapper = (mysql_client_wrapper *)arg;
+  rb_encoding *enc = rb_to_encoding(wrapper->encoding);
+  key = rb_str_export_to_enc(key, enc);
+  value = rb_str_export_to_enc(value, enc);
+
+  mysql_options4(wrapper->client, MYSQL_OPT_CONNECT_ATTR_ADD, StringValueCStr(key), StringValueCStr(value));
+  return ST_CONTINUE;
+}
+#endif
+
+static VALUE rb_connect(VALUE self, VALUE user, VALUE pass, VALUE host, VALUE port, VALUE database, VALUE socket, VALUE flags, VALUE conn_attrs) {
   struct nogvl_connect_args args;
   time_t start_time, end_time, elapsed_time, connect_timeout;
   VALUE rv;
@@ -415,6 +428,11 @@ static VALUE rb_connect(VALUE self, VALUE user, VALUE pass, VALUE host, VALUE po
   args.mysql       = wrapper->client;
   args.client_flag = NUM2ULONG(flags);
 
+#ifdef HAVE_CONST_MYSQL_OPT_CONNECT_ATTR_ADD
+  mysql_options(wrapper->client, MYSQL_OPT_CONNECT_ATTR_RESET, 0);
+  rb_hash_foreach(conn_attrs, opt_connect_attr_add_i, (VALUE)wrapper);
+#endif
+
   if (wrapper->connect_timeout)
     time(&start_time);
   rv = (VALUE) rb_thread_call_without_gvl(nogvl_connect, &args, RUBY_UBF_IO, 0);
@@ -1403,7 +1421,7 @@ void init_mysql2_client() {
   rb_define_private_method(cMysql2Client, "ssl_mode=", rb_set_ssl_mode_option, 1);
   rb_define_private_method(cMysql2Client, "enable_cleartext_plugin=", set_enable_cleartext_plugin, 1);
   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);
   rb_define_private_method(cMysql2Client, "_query", rb_query, 2);
 
   sym_id              = ID2SYM(rb_intern("id"));

ext/mysql2/extconf.rb

@@ -111,6 +111,7 @@ def add_ssl_defines(header)
 have_struct_member('MYSQL', 'net.vio', mysql_h)
 have_struct_member('MYSQL', 'net.pvio', mysql_h)
 have_const('MYSQL_ENABLE_CLEARTEXT_PLUGIN', mysql_h)
+have_const('MYSQL_OPT_CONNECT_ATTR_ADD', mysql_h) # for mysql_options4
 
 # This is our wishlist. We use whichever flags work on the host.
 # -Wall and -Wextra are included by default.

lib/mysql2/client.rb

@@ -64,6 +64,11 @@ def initialize(opts = {})
       # SSL verify is a connection flag rather than a mysql_ssl_set option
       flags |= SSL_VERIFY_SERVER_CERT if opts[:sslverify]
 
+      # Set default program_name in performance_schema.session_connect_attrs
+      # and performance_schema.session_account_connect_attrs
+      conn_attrs = opts[:connect_attrs] || {}
+      conn_attrs[:program_name] = $PROGRAM_NAME unless conn_attrs.key?(:program_name)
+
       if [:user, :pass, :hostname, :dbname, :db, :sock].any? { |k| @query_options.key?(k) }
         warn "============= WARNING FROM mysql2 ============="
         warn "The options :user, :pass, :hostname, :dbname, :db, and :sock are deprecated and will be removed at some point in the future."
@@ -85,8 +90,11 @@ def initialize(opts = {})
       port = port.to_i unless port.nil?
       database = database.to_s unless database.nil?
       socket = socket.to_s unless socket.nil?
+      conn_attrs = conn_attrs.each_with_object({}) do |(key, value), hash|
+        hash[key.to_s] = value.to_s
+      end
 
-      connect user, pass, host, port, database, socket, flags
+      connect user, pass, host, port, database, socket, flags, conn_attrs
     end
 
     def parse_ssl_mode(mode)

spec/mysql2/client_spec.rb

@@ -433,6 +433,26 @@ def run_gc
     expect(client.read_timeout).to be_nil
   end
 
+  it "should set default program_name in connect_attrs" do
+    client = new_client
+    if Mysql2::Client.info[:version] < '5.6' or client.info[:version] < '5.6'
+      pending('Both client and server versions must be MySQL 5.6 or later.')
+    end
+    result = client.query("SELECT attr_value FROM performance_schema.session_account_connect_attrs WHERE processlist_id = connection_id() AND attr_name = 'program_name'")
+    expect(result.first['attr_value']).to eq($0)
+  end
+
+  it "should set custom connect_attrs" do
+    client = new_client(:connect_attrs => {:program_name => 'my_program_name', :foo => 'fooval', :bar => 'barval'})
+    if Mysql2::Client.info[:version] < '5.6' or client.info[:version] < '5.6'
+      pending('Both client and server versions must be MySQL 5.6 or later.')
+    end
+    results = Hash[client.query("SELECT * FROM performance_schema.session_account_connect_attrs WHERE processlist_id = connection_id()").map { |x| x.values_at('ATTR_NAME', 'ATTR_VALUE') }]
+    expect(results['program_name']).to eq('my_program_name')
+    expect(results['foo']).to eq('fooval')
+    expect(results['bar']).to eq('barval')
+  end
+
   context "#query" do
     it "should let you query again if iterating is finished when streaming" do
       @client.query("SELECT 1 UNION SELECT 2", :stream => true, :cache_rows => false).each.to_a