Fix dropbox#50 executemany failed bug, add unit tests as well

taogeYT · May 15, 2020 · 7eaa0d6 · 7eaa0d6
1 parent 333d0d1
commit 7eaa0d6
Show file tree

Hide file tree

Showing 4 changed files with 58 additions and 20 deletions.
diff --git a/pyhive/common.py b/pyhive/common.py
@@ -19,6 +19,10 @@
 from itertools import islice
 import re
 
+MATCH_INSERT_SQL = re.compile(
+            r"\s*((?:INSERT)\b.+\bVALUES?\s*)(\(\s*(?:%s|%\(.+\)s)\s*(?:,\s*(?:%s|%\(.+\)s)\s*)*\))",
+            re.IGNORECASE | re.DOTALL)
+
 
 class DBAPICursor(with_metaclass(abc.ABCMeta, object)):
     """Base class for some common DB-API logic"""
@@ -86,26 +90,12 @@ def executemany(self, operation, seq_of_parameters):
 
         Return values are not defined.
         """
-        from .hive import _escaper
-        match_insert_sql = re.compile(
-            r"\s*((?:INSERT)\b.+\bVALUES?\s*)(\(\s*(?:%s|%\(.+\)s)\s*(?:,\s*(?:%s|%\(.+\)s)\s*)*\))",
-            re.IGNORECASE | re.DOTALL)
-        match = match_insert_sql.match(operation)
-        if match:
-            part1, part2 = match.group(1), match.group(2).rstrip()
-            values = []
-            for parameter in seq_of_parameters:
-                record = part2 % _escaper.escape_args(parameter)
-                values.append(record)
-            sql = part1 + ",".join(values)
-            self.execute(sql)
-        else:
-            for parameters in seq_of_parameters[:-1]:
-                self.execute(operation, parameters)
-                while self._state != self._STATE_FINISHED:
-                    self._fetch_more()
-            if seq_of_parameters:
-                self.execute(operation, seq_of_parameters[-1])
+        for parameters in seq_of_parameters[:-1]:
+            self.execute(operation, parameters)
+            while self._state != self._STATE_FINISHED:
+                self._fetch_more()
+        if seq_of_parameters:
+            self.execute(operation, seq_of_parameters[-1])
 
     def fetchone(self):
         """Fetch the next row of a query result set, returning a single sequence, or ``None`` when

diff --git a/pyhive/hive.py b/pyhive/hive.py
@@ -365,6 +365,26 @@ def execute(self, operation, parameters=None, **kwargs):
         _check_status(response)
         self._operationHandle = response.operationHandle
 
+    def executemany(self, operation, seq_of_parameters):
+        """Prepare a database operation (query or command) and then execute it against all parameter
+        sequences or mappings found in the sequence ``seq_of_parameters``.
+
+        Only the final result set is retained.
+
+        Return values are not defined.
+        """
+        match = common.MATCH_INSERT_SQL.match(operation)
+        if match:
+            part1, part2 = match.group(1), match.group(2).rstrip()
+            values = []
+            for parameter in seq_of_parameters:
+                record = part2 % _escaper.escape_args(parameter)
+                values.append(record)
+            sql = part1 + ",".join(values)
+            self.execute(sql)
+        else:
+            return super(Cursor, self).executemany(operation, seq_of_parameters)
+
     def cancel(self):
         req = ttypes.TCancelOperationReq(
             operationHandle=self._operationHandle,

diff --git a/pyhive/presto.py b/pyhive/presto.py
@@ -270,6 +270,26 @@ def execute(self, operation, parameters=None):
             url, data=sql.encode('utf-8'), headers=headers, **self._requests_kwargs)
         self._process_response(response)
 
+    def executemany(self, operation, seq_of_parameters):
+        """Prepare a database operation (query or command) and then execute it against all parameter
+        sequences or mappings found in the sequence ``seq_of_parameters``.
+
+        Only the final result set is retained.
+
+        Return values are not defined.
+        """
+        match = common.MATCH_INSERT_SQL.match(operation)
+        if match:
+            part1, part2 = match.group(1), match.group(2).rstrip()
+            values = []
+            for parameter in seq_of_parameters:
+                record = part2 % _escaper.escape_args(parameter)
+                values.append(record)
+            sql = part1 + ",".join(values)
+            self.execute(sql)
+        else:
+            return super(Cursor, self).executemany(operation, seq_of_parameters)
+
     def cancel(self):
         if self._state == self._STATE_NONE:
             raise ProgrammingError("No query yet")

diff --git a/pyhive/tests/dbapi_test_case.py b/pyhive/tests/dbapi_test_case.py
@@ -96,6 +96,14 @@ def test_executemany_none(self, cursor):
         self.assertIsNone(cursor.description)
         self.assertRaises(exc.ProgrammingError, cursor.fetchone)
 
+    @with_cursor
+    def test_executemany_batch_insert(self, cursor):
+        records = [("one", 1)]*10
+        cursor.execute("create temporary table test_executemany_batch_insert (foo string, bar int)")
+        cursor.executemany("INSERT INTO TABLE test_executemany_batch_insert VALUES (%s, %s)", records)
+        cursor.execute('SELECT * FROM test_executemany_batch_insert')
+        self.assertEqual(cursor.fetchall(), records)
+
     @with_cursor
     def test_fetchone_no_data(self, cursor):
         self.assertRaises(exc.ProgrammingError, cursor.fetchone)