Merge pull request #300 from mathieulongtin/issue_288_executemany

fix issue 288: executemany now works with "insert ... on duplicate update"

Author: methane

pymysql/cursors.py

@@ -11,7 +11,7 @@
 #: Regular expression for :meth:`Cursor.executemany`.
 #: executemany only suports simple bulk insert.
 #: You can use it to load large dataset.
-RE_INSERT_VALUES = re.compile(r"""INSERT\s.+\sVALUES\s+(\(\s*%s\s*(,\s*%s\s*)*\))\s*\Z""",
+RE_INSERT_VALUES = re.compile(r"""(INSERT\s.+\sVALUES\s+)(\(\s*%s\s*(?:,\s*%s\s*)*\))(\s*(?:ON DUPLICATE.*)?)\Z""",
                               re.IGNORECASE | re.DOTALL)
 
 
@@ -145,21 +145,24 @@ def executemany(self, query, args):
 
         m = RE_INSERT_VALUES.match(query)
         if m:
-            q_values = m.group(1).rstrip()
+            q_prefix = m.group(1)
+            q_values = m.group(2).rstrip()
+            q_postfix = m.group(3) or ''
             assert q_values[0] == '(' and q_values[-1] == ')'
-            q_prefix = query[:m.start(1)]
-            return self._do_execute_many(q_prefix, q_values, args,
+            return self._do_execute_many(q_prefix, q_values, q_postfix, args,
                                          self.max_stmt_length,
                                          self._get_db().encoding)
 
         self.rowcount = sum(self.execute(query, arg) for arg in args)
         return self.rowcount
 
-    def _do_execute_many(self, prefix, values, args, max_stmt_length, encoding):
+    def _do_execute_many(self, prefix, values, postfix, args, max_stmt_length, encoding):
         conn = self._get_db()
         escape = self._escape_args
         if isinstance(prefix, text_type):
             prefix = prefix.encode(encoding)
+        if isinstance(postfix, text_type):
+            postfix = postfix.encode(encoding)
         sql = bytearray(prefix)
         args = iter(args)
         v = values % escape(next(args), conn)
@@ -171,13 +174,13 @@ def _do_execute_many(self, prefix, values, args, max_stmt_length, encoding):
             v = values % escape(arg, conn)
             if isinstance(v, text_type):
                 v = v.encode(encoding)
-            if len(sql) + len(v) + 1 > max_stmt_length:
-                rows += self.execute(sql)
+            if len(sql) + len(v) + len(postfix) + 1 > max_stmt_length:
+                rows += self.execute(sql + postfix)
                 sql = bytearray(prefix)
             else:
                 sql += b','
             sql += v
-        rows += self.execute(sql)
+        rows += self.execute(sql + postfix)
         self.rowcount = rows
         return rows
 

pymysql/tests/test_basic.py

@@ -296,7 +296,7 @@ def test_bulk_insert_multiline_statement(self):
 %s , %s,
 %s )
  """, data)
-        self.assertEqual(cursor._last_executed, bytearray(b"""insert
+        self.assertEqual(cursor._last_executed.strip(), bytearray(b"""insert
 into bulkinsert (id, name,
 age, height)
 values (0,
@@ -318,6 +318,33 @@ def test_bulk_insert_single_record(self):
         cursor.execute('commit')
         self._verify_records(data)
 
+    def test_issue_288(self):
+        """executemany should work with "insert ... on update" """
+        conn = self.connections[0]
+        cursor = conn.cursor()
+        data = [(0, "bob", 21, 123), (1, "jim", 56, 45), (2, "fred", 100, 180)]
+        cursor.executemany("""insert
+into bulkinsert (id, name,
+age, height)
+values (%s,
+%s , %s,
+%s ) on duplicate key update
+age = values(age)
+ """, data)
+        self.assertEqual(cursor._last_executed.strip(), bytearray(b"""insert
+into bulkinsert (id, name,
+age, height)
+values (0,
+'bob' , 21,
+123 ),(1,
+'jim' , 56,
+45 ),(2,
+'fred' , 100,
+180 ) on duplicate key update
+age = values(age)"""))
+        cursor.execute('commit')
+        self._verify_records(data)
+
     def test_warnings(self):
         con = self.connections[0]
         cur = con.cursor()