Fixed LOAD DATA LOCAL INFILE commands

* Fixes for both "buffered" & "unbuffered" cursor types
* Registered `test_load_local` to run with tests
* Refactored `test_load_local` tests to work with the tornado framework

More information here:
#29

Author: srstrickland

tornado_mysql/connections.py

@@ -105,7 +105,7 @@ def is_ascii(data):
 def _scramble(password, message):
     if not password:
         return b'\0'
-    if DEBUG: print('password=' + password)
+    if DEBUG: print('password=' + str(password))
     stage1 = sha_new(password).digest()
     stage2 = sha_new(stage1).digest()
     s = sha_new()
@@ -1064,7 +1064,7 @@ def read(self):
             if first_packet.is_ok_packet():
                 self._read_ok_packet(first_packet)
             elif first_packet.is_load_local_packet():
-                self._read_load_local_packet(first_packet)
+                yield self._read_load_local_packet(first_packet)
             else:
                 yield self._read_result_packet(first_packet)
         finally:
@@ -1079,6 +1079,10 @@ def init_unbuffered_query(self):
             self._read_ok_packet(first_packet)
             self.unbuffered_active = False
             self.connection = None
+        elif first_packet.is_load_local_packet():
+            yield self._read_load_local_packet(first_packet)
+            self.unbuffered_active = False
+            self.connection = None
         else:
             self.field_count = first_packet.read_length_encoded_integer()
             yield self._get_descriptions()
@@ -1097,12 +1101,13 @@ def _read_ok_packet(self, first_packet):
         self.message = ok_packet.message
         self.has_next = ok_packet.has_next
 
+    @gen.coroutine
     def _read_load_local_packet(self, first_packet):
         load_packet = LoadLocalPacketWrapper(first_packet)
         sender = LoadLocalFile(load_packet.filename, self.connection)
         sender.send_data()
 
-        ok_packet = self.connection._read_packet()
+        ok_packet = yield self.connection._read_packet()
         if not ok_packet.is_ok_packet():
             raise OperationalError(2014, "Commands Out of Sync")
         self._read_ok_packet(ok_packet)
@@ -1219,7 +1224,7 @@ def __init__(self, filename, connection):
 
     def send_data(self):
         """Send data packets from the local file to the server"""
-        if not self.connection.socket:
+        if not self.connection._stream:
             raise InterfaceError("(0, '')")
 
         # sequence id is 2 as we already sent a query packet

tornado_mysql/tests/__init__.py

@@ -1,5 +1,6 @@
 from tornado_mysql.tests.test_issues import *
 from tornado_mysql.tests.test_basic import *
+from tornado_mysql.tests.test_load_local import *
 from tornado_mysql.tests.test_nextset import *
 from tornado_mysql.tests.test_DictCursor import *
 from tornado_mysql.tests.test_connection import TestConnection

tornado_mysql/tests/test_load_local.py

@@ -1,5 +1,6 @@
-from pymysql import OperationalError, Warning
-from pymysql.tests import base
+from tornado.testing import gen_test
+from tornado_mysql.err import OperationalError, Warning
+from tornado_mysql.tests import base
 
 import os
 import warnings
@@ -8,59 +9,68 @@
 
 
 class TestLoadLocal(base.PyMySQLTestCase):
+    @gen_test
     def test_no_file(self):
         """Test load local infile when the file does not exist"""
         conn = self.connections[0]
         c = conn.cursor()
-        c.execute("CREATE TABLE test_load_local (a INTEGER, b INTEGER)")
+        with warnings.catch_warnings():
+            warnings.simplefilter("ignore")
+            yield c.execute("DROP TABLE IF EXISTS test_load_local")
+        yield c.execute("CREATE TABLE test_load_local (a INTEGER, b INTEGER)")
         try:
-            self.assertRaises(
-                OperationalError,
-                c.execute,
-                ("LOAD DATA LOCAL INFILE 'no_data.txt' INTO TABLE "
-                 "test_load_local fields terminated by ','")
-            )
+            with self.assertRaises(OperationalError) as cm:
+                yield c.execute(
+                    "LOAD DATA LOCAL INFILE 'no_data.txt' INTO TABLE "
+                    "test_load_local fields terminated by ','")
         finally:
-            c.execute("DROP TABLE test_load_local")
+            yield c.execute("DROP TABLE IF EXISTS test_load_local")
             c.close()
-
+    @gen_test
     def test_load_file(self):
         """Test load local infile with a valid file"""
         conn = self.connections[0]
         c = conn.cursor()
-        c.execute("CREATE TABLE test_load_local (a INTEGER, b INTEGER)")
+        with warnings.catch_warnings():
+            warnings.simplefilter("ignore")
+            yield c.execute("DROP TABLE IF EXISTS test_load_local")
+        yield c.execute("CREATE TABLE test_load_local (a INTEGER, b INTEGER)")
         filename = os.path.join(os.path.dirname(os.path.realpath(__file__)),
                                 'data',
                                 'load_local_data.txt')
         try:
-            c.execute(
+            yield c.execute(
                 ("LOAD DATA LOCAL INFILE '{0}' INTO TABLE " +
                  "test_load_local FIELDS TERMINATED BY ','").format(filename)
             )
-            c.execute("SELECT COUNT(*) FROM test_load_local")
+            yield c.execute("SELECT COUNT(*) FROM test_load_local")
             self.assertEqual(22749, c.fetchone()[0])
         finally:
-            c.execute("DROP TABLE test_load_local")
+            yield c.execute("DROP TABLE IF EXISTS test_load_local")
 
+    @gen_test
     def test_load_warnings(self):
         """Test load local infile produces the appropriate warnings"""
         conn = self.connections[0]
         c = conn.cursor()
-        c.execute("CREATE TABLE test_load_local (a INTEGER, b INTEGER)")
+        with warnings.catch_warnings():
+            warnings.simplefilter("ignore")
+            yield c.execute("DROP TABLE IF EXISTS test_load_local")
+        yield c.execute("CREATE TABLE test_load_local (a INTEGER, b INTEGER)")
         filename = os.path.join(os.path.dirname(os.path.realpath(__file__)),
                                 'data',
                                 'load_local_warn_data.txt')
         try:
             with warnings.catch_warnings(record=True) as w:
                 warnings.simplefilter('always')
-                c.execute(
+                yield c.execute(
                     ("LOAD DATA LOCAL INFILE '{0}' INTO TABLE " +
                      "test_load_local FIELDS TERMINATED BY ','").format(filename)
                 )
                 self.assertEqual(w[0].category, Warning)
                 self.assertTrue("Incorrect integer value" in str(w[-1].message))
         finally:
-            c.execute("DROP TABLE test_load_local")
+            yield c.execute("DROP TABLE IF EXISTS test_load_local")
 
 
 if __name__ == "__main__":