Merge pull request #446 from sodabrew/local_infile_handler

Local local infile handler

Author: sodabrew

ext/mysql2/client.c

@@ -1,5 +1,5 @@
 #include <mysql2_ext.h>
-#include <client.h>
+
 #include <errno.h>
 #ifndef _WIN32
 #include <sys/socket.h>
@@ -146,9 +146,13 @@ static VALUE rb_raise_mysql2_error(mysql_client_wrapper *wrapper) {
 
 static void *nogvl_init(void *ptr) {
   MYSQL *client;
+  mysql_client_wrapper *wrapper = (mysql_client_wrapper *)ptr;
 
   /* may initialize embedded server and read /etc/services off disk */
-  client = mysql_init((MYSQL *)ptr);
+  client = mysql_init(wrapper->client);
+
+  if (client) mysql2_set_local_infile(client, wrapper);
+
   return (void*)(client ? Qtrue : Qfalse);
 }
 
@@ -1124,7 +1128,7 @@ static VALUE set_read_default_group(VALUE self, VALUE value) {
 static VALUE initialize_ext(VALUE self) {
   GET_CLIENT(self);
 
-  if ((VALUE)rb_thread_call_without_gvl(nogvl_init, wrapper->client, RUBY_UBF_IO, 0) == Qfalse) {
+  if ((VALUE)rb_thread_call_without_gvl(nogvl_init, wrapper, RUBY_UBF_IO, 0) == Qfalse) {
     /* TODO: warning - not enough memory? */
     return rb_raise_mysql2_error(wrapper);
   }
@@ -1139,7 +1143,7 @@ void init_mysql2_client() {
   int i;
   int dots = 0;
   const char *lib = mysql_get_client_info();
-  
+
   for (i = 0; lib[i] != 0 && MYSQL_LINK_VERSION[i] != 0; i++) {
     if (lib[i] == '.') {
       dots++;

ext/mysql2/infile.c

@@ -0,0 +1,119 @@
+#include <mysql2_ext.h>
+
+#include <errno.h>
+#include <unistd.h>
+
+#define ERROR_LEN 1024
+typedef struct
+{
+  int fd;
+  char *filename;
+  char error[ERROR_LEN];
+  mysql_client_wrapper *wrapper;
+} mysql2_local_infile_data;
+
+/* MySQL calls this function when a user begins a LOAD DATA LOCAL INFILE query.
+ *
+ * Allocate a data struct and pass it back through the data pointer.
+ *
+ * Returns:
+ * 0 on success
+ * 1 on error
+ */
+static int
+mysql2_local_infile_init(void **ptr, const char *filename, void *userdata)
+{
+  mysql2_local_infile_data *data = malloc(sizeof(mysql2_local_infile_data));
+  if (!data) return 1;
+
+  *ptr = data;
+  data->error[0] = 0;
+  data->wrapper = userdata;
+
+  data->filename = strdup(filename);
+  if (!data->filename) {
+    snprintf(data->error, ERROR_LEN, "%s: %s", strerror(errno), filename);
+    return 1;
+  }
+
+  data->fd = open(filename, O_RDONLY);
+  if (data->fd < 0) {
+    snprintf(data->error, ERROR_LEN, "%s: %s", strerror(errno), filename);
+    return 1;
+  }
+
+  return 0;
+}
+
+/* MySQL calls this function to read data from the local file.
+ *
+ * Returns:
+ * > 0   number of bytes read
+ * == 0  end of file
+ * < 0   error
+ */
+static int
+mysql2_local_infile_read(void *ptr, char *buf, uint buf_len)
+{
+  int count;
+  mysql2_local_infile_data *data = (mysql2_local_infile_data *)ptr;
+
+  count = (int)read(data->fd, buf, buf_len);
+  if (count < 0) {
+    snprintf(data->error, ERROR_LEN, "%s: %s", strerror(errno), data->filename);
+  }
+
+  return count;
+}
+
+/* MySQL calls this function when we're done with the LOCAL INFILE query.
+ *
+ * ptr will be null if the init function failed.
+ */
+static void
+mysql2_local_infile_end(void *ptr)
+{
+  mysql2_local_infile_data *data = (mysql2_local_infile_data *)ptr;
+  if (data) {
+    if (data->fd >= 0)
+      close(data->fd);
+    if (data->filename)
+      free(data->filename);
+    free(data);
+  }
+}
+
+/* MySQL calls this function if any of the functions above returned an error.
+ *
+ * This function is called even if init failed, with whatever ptr value
+ * init has set, regardless of the return value of the init function.
+ *
+ * Returns:
+ * Error message number (see http://dev.mysql.com/doc/refman/5.0/en/error-messages-client.html)
+ */
+static int
+mysql2_local_infile_error(void *ptr, char *error_msg, uint error_msg_len)
+{
+  mysql2_local_infile_data *data = (mysql2_local_infile_data *) ptr;
+
+  if (data) {
+    snprintf(error_msg, error_msg_len, "%s", data->error);
+    return CR_UNKNOWN_ERROR;
+  }
+
+  snprintf(error_msg, error_msg_len, "Out of memory");
+  return CR_OUT_OF_MEMORY;
+}
+
+/* Tell MySQL Client to use our own local_infile functions.
+ * This is both due to bugginess in the default handlers,
+ * and to improve the Rubyness of the handlers here.
+ */
+void mysql2_set_local_infile(MYSQL *mysql, void *userdata)
+{
+  mysql_set_local_infile_handler(mysql,
+                                 mysql2_local_infile_init,
+                                 mysql2_local_infile_read,
+                                 mysql2_local_infile_end,
+                                 mysql2_local_infile_error, userdata);
+}

ext/mysql2/infile.h

@@ -0,0 +1 @@
+void mysql2_set_local_infile(MYSQL *mysql, void *userdata);

ext/mysql2/mysql2_ext.h

@@ -41,5 +41,6 @@ typedef unsigned int    uint;
 
 #include <client.h>
 #include <result.h>
+#include <infile.h>
 
 #endif

ext/mysql2/result.c

@@ -1,4 +1,5 @@
 #include <mysql2_ext.h>
+
 #include <stdint.h>
 
 #include "mysql_enc_to_ruby.h"

spec/mysql2/client_spec.rb

@@ -78,7 +78,10 @@ def connect *args
   end
 
   it "should be able to connect via SSL options" do
-    pending("DON'T WORRY, THIS TEST PASSES :) - but is machine-specific. You need to have MySQL running with SSL configured and enabled. Then update the paths in this test to your needs and remove the pending state.")
+    ssl = @client.query "SHOW VARIABLES LIKE 'have_%ssl'"
+    ssl_enabled = ssl.any? {|x| x['Value'] == 'ENABLED'}
+    pending("DON'T WORRY, THIS TEST PASSES - but SSL is not enabled in your MySQL daemon.") unless ssl_enabled
+    pending("DON'T WORRY, THIS TEST PASSES - but you must update the SSL cert paths in this test and remove this pending state.")
     ssl_client = nil
     lambda {
       ssl_client = Mysql2::Client.new(
@@ -197,6 +200,49 @@ def connect *args
     end
   end
 
+  context ":local_infile" do
+    before(:all) do
+      @client_i = Mysql2::Client.new DatabaseCredentials['root'].merge(:local_infile => true)
+      local = @client_i.query "SHOW VARIABLES LIKE 'local_infile'"
+      local_enabled = local.any? {|x| x['Value'] == 'ON'}
+      pending("DON'T WORRY, THIS TEST PASSES - but LOCAL INFILE is not enabled in your MySQL daemon.") unless local_enabled
+
+      @client_i.query %[
+        CREATE TABLE IF NOT EXISTS infileTest (
+          id MEDIUMINT NOT NULL AUTO_INCREMENT PRIMARY KEY,
+          foo VARCHAR(10),
+          bar MEDIUMTEXT
+        )
+      ]
+    end
+
+    after(:all) do
+      @client_i.query "DROP TABLE infileTest"
+    end
+
+    it "should raise an error when local_infile is disabled" do
+      client = Mysql2::Client.new DatabaseCredentials['root'].merge(:local_infile => false)
+      lambda {
+        client.query "LOAD DATA LOCAL INFILE 'spec/test_data' INTO TABLE infileTest"
+      }.should raise_error(Mysql2::Error, %r{command is not allowed})
+    end
+
+    it "should raise an error when a non-existent file is loaded" do
+      lambda {
+        @client_i.query "LOAD DATA LOCAL INFILE 'this/file/is/not/here' INTO TABLE infileTest"
+      }.should_not raise_error(Mysql2::Error, %r{file not found: this/file/is/not/here})
+    end
+
+    it "should LOAD DATA LOCAL INFILE" do
+      @client_i.query "LOAD DATA LOCAL INFILE 'spec/test_data' INTO TABLE infileTest"
+      info = @client_i.query_info
+      info.should eql({:records => 1, :deleted => 0, :skipped => 0, :warnings => 0})
+
+      result = @client_i.query "SELECT * FROM infileTest"
+      result.first.should eql({'id' => 1, 'foo' => 'Hello', 'bar' => 'World'})
+    end
+  end
+
   it "should expect connect_timeout to be a positive integer" do
     lambda {
       Mysql2::Client.new(:connect_timeout => -1)

spec/test_data

@@ -0,0 +1 @@
+\N	Hello	World