More tests for the simple_* methods

changelog.d/16596.misc

@@ -0,0 +1 @@
+Improve tests of the SQL generator.

synapse/storage/database.py

@@ -2062,9 +2062,7 @@ def simple_update_many_txn(
             where_clause = ""
 
         # UPDATE mytable SET col1 = ?, col2 = ? WHERE col3 = ? AND col4 = ?
-        sql = f"""
-            UPDATE {table} SET {set_clause} {where_clause}
-        """
+        sql = f"UPDATE {table} SET {set_clause} {where_clause}"
 
         txn.execute_batch(sql, args)
 
@@ -2283,17 +2281,14 @@ def simple_delete_many_txn(
         if not values:
             return 0
 
-        sql = "DELETE FROM %s" % table
-
         clause, values = make_in_list_sql_clause(txn.database_engine, column, values)
         clauses = [clause]
 
         for key, value in keyvalues.items():
             clauses.append("%s = ?" % (key,))
             values.append(value)
 
-        if clauses:
-            sql = "%s WHERE %s" % (sql, " AND ".join(clauses))
+        sql = "DELETE FROM %s WHERE %s" % (table, " AND ".join(clauses))
         txn.execute(sql, values)
 
         return txn.rowcount

tests/storage/test_base.py

@@ -14,7 +14,7 @@
 
 from collections import OrderedDict
 from typing import Generator
-from unittest.mock import Mock
+from unittest.mock import Mock, call
 
 from twisted.internet import defer
 
@@ -55,6 +55,8 @@ def runWithConnection(func, *args, **kwargs):  # type: ignore[no-untyped-def]
         engine = create_engine(sqlite_config)
         fake_engine = Mock(wraps=engine)
         fake_engine.in_transaction.return_value = False
+        fake_engine.module.OperationalError = engine.module.OperationalError
+        fake_engine.module.DatabaseError = engine.module.DatabaseError
 
         db = DatabasePool(Mock(), Mock(config=sqlite_config), fake_engine)
         db._db_pool = self.db_pool
@@ -91,6 +93,33 @@ def test_insert_3cols(self) -> Generator["defer.Deferred[object]", object, None]
             "INSERT INTO tablename (colA, colB, colC) VALUES(?, ?, ?)", (1, 2, 3)
         )
 
+    @defer.inlineCallbacks
+    def test_insert_many(self) -> Generator["defer.Deferred[object]", object, None]:
+        yield defer.ensureDeferred(
+            self.datastore.db_pool.simple_insert_many(
+                table="tablename",
+                keys=(
+                    "col1",
+                    "col2",
+                ),
+                values=[
+                    (
+                        "val1",
+                        "val2",
+                    ),
+                    ("val3", "val4"),
+                ],
+                desc="",
+            )
+        )
+
+        # TODO Test postgres variant.
+
+        self.mock_txn.executemany.assert_called_with(
+            "INSERT INTO tablename (col1, col2) VALUES(?, ?)",
+            [("val1", "val2"), ("val3", "val4")],
+        )
+
     @defer.inlineCallbacks
     def test_select_one_1col(self) -> Generator["defer.Deferred[object]", object, None]:
         self.mock_txn.rowcount = 1
@@ -160,6 +189,54 @@ def test_select_list(self) -> Generator["defer.Deferred[object]", object, None]:
             "SELECT colA FROM tablename WHERE keycol = ?", ["A set"]
         )
 
+    @defer.inlineCallbacks
+    def test_select_many_batch(
+        self,
+    ) -> Generator["defer.Deferred[object]", object, None]:
+        self.mock_txn.rowcount = 3
+        self.mock_txn.fetchall.side_effect = [[(1,), (2,)], [(3,)]]
+
+        ret = yield defer.ensureDeferred(
+            self.datastore.db_pool.simple_select_many_batch(
+                table="tablename",
+                column="col1",
+                iterable=("val1", "val2", "val3"),
+                retcols=("col2",),
+                keyvalues={"col3": "val4"},
+                batch_size=2,
+            )
+        )
+
+        self.mock_txn.execute.assert_has_calls(
+            [
+                call(
+                    "SELECT col2 FROM tablename WHERE col1 = ANY(?) AND col3 = ?",
+                    [["val1", "val2"], "val4"],
+                ),
+                call(
+                    "SELECT col2 FROM tablename WHERE col1 = ANY(?) AND col3 = ?",
+                    [["val3"], "val4"],
+                ),
+            ],
+        )
+        self.assertEqual([(1,), (2,), (3,)], ret)
+
+    def test_select_many_no_iterable(self) -> None:
+        self.mock_txn.rowcount = 3
+        self.mock_txn.fetchall.side_effect = [(1,), (2,)]
+
+        ret = self.datastore.db_pool.simple_select_many_txn(
+            self.mock_txn,
+            table="tablename",
+            column="col1",
+            iterable=(),
+            retcols=("col2",),
+            keyvalues={"col3": "val4"},
+        )
+
+        self.mock_txn.execute.assert_not_called()
+        self.assertEqual([], ret)
+
     @defer.inlineCallbacks
     def test_update_one_1col(self) -> Generator["defer.Deferred[object]", object, None]:
         self.mock_txn.rowcount = 1
@@ -196,6 +273,37 @@ def test_update_one_4cols(
             [3, 4, 1, 2],
         )
 
+    @defer.inlineCallbacks
+    def test_update_many(self) -> Generator["defer.Deferred[object]", object, None]:
+        yield defer.ensureDeferred(
+            self.datastore.db_pool.simple_update_many(
+                table="tablename",
+                key_names=("col1", "col2"),
+                key_values=[("val1", "val2")],
+                value_names=("col3",),
+                value_values=[("val3",)],
+                desc="",
+            )
+        )
+
+        self.mock_txn.executemany.assert_called_with(
+            "UPDATE tablename SET col3 = ? WHERE col1 = ? AND col2 = ?",
+            [("val3", "val1", "val2"), ("val3", "val1", "val2")],
+        )
+
+        # key_values and value_values must be the same length.
+        with self.assertRaises(ValueError):
+            yield defer.ensureDeferred(
+                self.datastore.db_pool.simple_update_many(
+                    table="tablename",
+                    key_names=("col1", "col2"),
+                    key_values=[("val1", "val2")],
+                    value_names=("col3",),
+                    value_values=[],
+                    desc="",
+                )
+            )
+
     @defer.inlineCallbacks
     def test_delete_one(self) -> Generator["defer.Deferred[object]", object, None]:
         self.mock_txn.rowcount = 1
@@ -209,3 +317,163 @@ def test_delete_one(self) -> Generator["defer.Deferred[object]", object, None]:
         self.mock_txn.execute.assert_called_with(
             "DELETE FROM tablename WHERE keycol = ?", ["Go away"]
         )
+
+    @defer.inlineCallbacks
+    def test_delete_many(self) -> Generator["defer.Deferred[object]", object, None]:
+        self.mock_txn.rowcount = 2
+
+        result = yield defer.ensureDeferred(
+            self.datastore.db_pool.simple_delete_many(
+                table="tablename",
+                column="col1",
+                iterable=("val1", "val2"),
+                keyvalues={"col2": "val3"},
+                desc="",
+            )
+        )
+
+        self.mock_txn.execute.assert_called_with(
+            "DELETE FROM tablename WHERE col1 = ANY(?) AND col2 = ?",
+            [["val1", "val2"], "val3"],
+        )
+        self.assertEqual(result, 2)
+
+    @defer.inlineCallbacks
+    def test_delete_many_no_iterable(
+        self,
+    ) -> Generator["defer.Deferred[object]", object, None]:
+        result = yield defer.ensureDeferred(
+            self.datastore.db_pool.simple_delete_many(
+                table="tablename",
+                column="col1",
+                iterable=(),
+                keyvalues={"col2": "val3"},
+                desc="",
+            )
+        )
+
+        self.mock_txn.execute.assert_not_called()
+        self.assertEqual(result, 0)
+
+    @defer.inlineCallbacks
+    def test_delete_many_no_keyvalues(
+        self,
+    ) -> Generator["defer.Deferred[object]", object, None]:
+        self.mock_txn.rowcount = 2
+
+        result = yield defer.ensureDeferred(
+            self.datastore.db_pool.simple_delete_many(
+                table="tablename",
+                column="col1",
+                iterable=("val1", "val2"),
+                keyvalues={},
+                desc="",
+            )
+        )
+
+        self.mock_txn.execute.assert_called_with(
+            "DELETE FROM tablename WHERE col1 = ANY(?)", [["val1", "val2"]]
+        )
+        self.assertEqual(result, 2)
+
+    @defer.inlineCallbacks
+    def test_upsert(self) -> Generator["defer.Deferred[object]", object, None]:
+        self.mock_txn.rowcount = 1
+
+        result = yield defer.ensureDeferred(
+            self.datastore.db_pool.simple_upsert(
+                table="tablename",
+                keyvalues={"columnname": "oldvalue"},
+                values={"othercol": "newvalue"},
+            )
+        )
+
+        self.mock_txn.execute.assert_called_with(
+            "INSERT INTO tablename (columnname, othercol) VALUES (?, ?) ON CONFLICT (columnname)  DO UPDATE SET othercol=EXCLUDED.othercol",
+            ["oldvalue", "newvalue"],
+        )
+        self.assertTrue(result)
+
+    @defer.inlineCallbacks
+    def test_upsert_with_insert(
+        self,
+    ) -> Generator["defer.Deferred[object]", object, None]:
+        self.mock_txn.rowcount = 1
+
+        result = yield defer.ensureDeferred(
+            self.datastore.db_pool.simple_upsert(
+                table="tablename",
+                keyvalues={"columnname": "oldvalue"},
+                values={"othercol": "newvalue"},
+                insertion_values={"thirdcol": "insertionval"},
+            )
+        )
+
+        self.mock_txn.execute.assert_called_with(
+            "INSERT INTO tablename (columnname, thirdcol, othercol) VALUES (?, ?, ?) ON CONFLICT (columnname)  DO UPDATE SET othercol=EXCLUDED.othercol",
+            ["oldvalue", "insertionval", "newvalue"],
+        )
+        self.assertTrue(result)
+
+    @defer.inlineCallbacks
+    def test_upsert_with_where(
+        self,
+    ) -> Generator["defer.Deferred[object]", object, None]:
+        self.mock_txn.rowcount = 1
+
+        result = yield defer.ensureDeferred(
+            self.datastore.db_pool.simple_upsert(
+                table="tablename",
+                keyvalues={"columnname": "oldvalue"},
+                values={"othercol": "newvalue"},
+                where_clause="thirdcol IS NULL",
+            )
+        )
+
+        self.mock_txn.execute.assert_called_with(
+            "INSERT INTO tablename (columnname, othercol) VALUES (?, ?) ON CONFLICT (columnname) WHERE thirdcol IS NULL DO UPDATE SET othercol=EXCLUDED.othercol",
+            ["oldvalue", "newvalue"],
+        )
+        self.assertTrue(result)
+
+    @defer.inlineCallbacks
+    def test_upsert_many(self) -> Generator["defer.Deferred[object]", object, None]:
+        yield defer.ensureDeferred(
+            self.datastore.db_pool.simple_upsert_many(
+                table="tablename",
+                key_names=["columnname"],
+                key_values=[["oldvalue"]],
+                value_names=["othercol"],
+                value_values=[["newvalue"]],
+                desc="",
+            )
+        )
+
+        # TODO Test postgres variant.
+
+        self.mock_txn.executemany.assert_called_with(
+            "INSERT INTO tablename (columnname, othercol) VALUES (?, ?) ON CONFLICT (columnname) DO UPDATE SET othercol=EXCLUDED.othercol",
+            [("oldvalue", "newvalue")],
+        )
+
+    @defer.inlineCallbacks
+    def test_upsert_many_no_values(
+        self,
+    ) -> Generator["defer.Deferred[object]", object, None]:
+        yield defer.ensureDeferred(
+            self.datastore.db_pool.simple_upsert_many(
+                table="tablename",
+                key_names=["columnname"],
+                key_values=[["oldvalue"]],
+                value_names=[],
+                value_values=[],
+                desc="",
+            )
+        )
+
+        # TODO Test postgres variant.
+
+        self.mock_txn.executemany.assert_called_with(
+            "INSERT INTO tablename (columnname) VALUES (?) ON CONFLICT (columnname) DO NOTHING",
+            [("oldvalue",)],
+        )