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"