Merge pull request #167 from AlexLisovoy/bulk_operations

Update regular expressions to make up to date with pymysql.

aiomysql/cursors.py

@@ -11,21 +11,22 @@
 from .utils import PY_35, create_future
 
 
-# https://github.com/PyMySQL/PyMySQL/blob/master/pymysql/cursors.py#L11-L15
+# https://github.com/PyMySQL/PyMySQL/blob/master/pymysql/cursors.py#L11-L18
 
 #: 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*(?:ON DUPLICATE.*)?)\Z""",
+    r"\s*((?:INSERT|REPLACE)\s.+\sVALUES?\s+)" +
+    r"(\(\s*(?:%s|%\(.+\)s)\s*(?:,\s*(?:%s|%\(.+\)s)\s*)*\))" +
+    r"(\s*(?:ON DUPLICATE.*)?);?\s*\Z",
     re.IGNORECASE | re.DOTALL)
 
 
 class Cursor:
     """Cursor is used to interact with the database."""
 
-    #: Max stetement size which :meth:`executemany` generates.
+    #: Max statement size which :meth:`executemany` generates.
     #:
     #: Max size of allowed statement is max_allowed_packet -
     # packet_header_size.
@@ -260,8 +261,8 @@ def executemany(self, query, args):
             stmt = "INSERT INTO employees (name, phone) VALUES ('%s','%s')"
             yield from cursor.executemany(stmt, data)
 
-        INSERT statements are optimized by batching the data, that is
-        using the MySQL multiple rows syntax.
+        INSERT or REPLACE statements are optimized by batching the data,
+        that is using the MySQL multiple rows syntax.
 
         :param query: `str`, sql statement
         :param args: ``tuple`` or ``list`` of arguments for sql query

tests/test_bulk_inserts.py

@@ -28,6 +28,20 @@ def f(data):
     return f
 
 
+@pytest.fixture
+def assert_dict_records(connection):
+    @asyncio.coroutine
+    def f(data):
+        cursor = yield from connection.cursor(DictCursor)
+        yield from cursor.execute(
+            "SELECT id, name, age, height FROM bulkinsert")
+        result = yield from cursor.fetchall()
+        yield from cursor.execute('COMMIT')
+        assert sorted(data, key=lambda k: k['id']) == \
+            sorted(result, key=lambda k: k['id'])
+    return f
+
+
 @pytest.mark.run_loop
 def test_bulk_insert(cursor, table, assert_records):
     data = [(0, "bob", 21, 123), (1, "jim", 56, 45), (2, "fred", 100, 180)]
@@ -101,3 +115,78 @@ def test_insert_on_duplicate_key_update(cursor, table, assert_records):
         age = values(age)""")
     yield from cursor.execute('COMMIT')
     yield from assert_records(data)
+
+
+@pytest.mark.run_loop
+def test_bulk_insert_with_params_as_dict(cursor, table, assert_dict_records):
+    data = [
+        {
+            'id': 0,
+            'name': "bob",
+            'age': 21,
+            'height': 123
+        },
+        {
+            'id': 1,
+            'name': "jim",
+            'age': 56,
+            'height': 45
+        },
+        {
+            'id': 2,
+            'name': "fred",
+            'age': 100,
+            'height': 180
+        },
+    ]
+    yield from cursor.executemany(
+        "INSERT INTO bulkinsert (id, name, age, height) "
+        "VALUES (%(id)s,%(name)s,%(age)s,%(height)s)", data)
+    expected = bytearray(b"INSERT INTO bulkinsert (id, name, age, height) "
+                         b"VALUES (0,'bob',21,123),(1,'jim',56,45),"
+                         b"(2,'fred',100,180)")
+    assert cursor._last_executed == expected
+    yield from cursor.execute('commit')
+    yield from assert_dict_records(data)
+
+
+@pytest.mark.run_loop
+def test_bulk_insert_with_precedence_spaces(cursor, table, assert_records):
+    data = [(0, "bob", 21, 123), (1, "jim", 56, 45)]
+    yield from cursor.executemany("""
+        INSERT INTO bulkinsert (id, name, age, height)
+        VALUES (%s,%s,%s,%s)
+    """, data)
+    expected = bytearray(b"INSERT INTO bulkinsert (id, name, age, height)"
+                         b"\n        VALUES (0,\'bob\',21,123),"
+                         b"(1,\'jim\',56,45)\n    ")
+    assert cursor._last_executed == expected
+    yield from cursor.execute('commit')
+    yield from assert_records(data)
+
+
+@pytest.mark.run_loop
+def test_bulk_replace(cursor, table, assert_records):
+    data = [(0, "bob", 21, 123), (0, "jim", 56, 45)]
+    sql = ("REPLACE INTO bulkinsert (id, name, age, height) " +
+           "VALUES (%s,%s,%s,%s)")
+    yield from cursor.executemany(sql, data)
+    assert cursor._last_executed.strip() == bytearray(
+        b"REPLACE INTO bulkinsert (id, name, age, height) " +
+        b"VALUES (0,'bob',21,123),(0,'jim',56,45)"
+    )
+    yield from cursor.execute('COMMIT')
+    yield from assert_records([(0, "jim", 56, 45)])
+
+
+@pytest.mark.run_loop
+def test_bulk_insert_with_semicolon_at_the_end(cursor, table, assert_records):
+    data = [(0, "bob", 21, 123), (1, "jim", 56, 45)]
+    yield from cursor.executemany(
+        "INSERT INTO bulkinsert (id, name, age, height) "
+        "VALUES (%s,%s,%s,%s);", data)
+    expected = bytearray(b"INSERT INTO bulkinsert (id, name, age, height) "
+                         b"VALUES (0,'bob',21,123),(1,'jim',56,45)")
+    assert cursor._last_executed == expected
+    yield from cursor.execute('commit')
+    yield from assert_records(data)