Fixing LOAD DATA LOCAL INFILE for SSCursor* types (PyMySQL#473)

srstrickland · methane · commit 30422b2511f3 · 2016-06-08T15:11:39.000+09:00
diff --git a/pymysql/connections.py b/pymysql/connections.py
@@ -1319,6 +1319,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():
+            self._read_load_local_packet(first_packet)
+            self.unbuffered_active = False
+            self.connection = None
         else:
             self.field_count = first_packet.read_length_encoded_integer()
             self._get_descriptions()
diff --git a/pymysql/tests/test_load_local.py b/pymysql/tests/test_load_local.py
@@ -1,4 +1,4 @@
-from pymysql import OperationalError, Warning
+from pymysql import cursors, OperationalError, Warning
 from pymysql.tests import base
 
 import os
@@ -42,6 +42,28 @@ def test_load_file(self):
         finally:
             c.execute("DROP TABLE test_load_local")
 
+    def test_unbuffered_load_file(self):
+        """Test unbuffered load local infile with a valid file"""
+        conn = self.connections[0]
+        c = conn.cursor(cursors.SSCursor)
+        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(
+                ("LOAD DATA LOCAL INFILE '{0}' INTO TABLE " +
+                 "test_load_local FIELDS TERMINATED BY ','").format(filename)
+            )
+            c.execute("SELECT COUNT(*) FROM test_load_local")
+            self.assertEqual(22749, c.fetchone()[0])
+        finally:
+            c.close()
+            conn.close()
+            conn.connect()
+            c = conn.cursor()
+            c.execute("DROP TABLE test_load_local")
+
     def test_load_warnings(self):
         """Test load local infile produces the appropriate warnings"""
         conn = self.connections[0]
