You can not select more than 25 topics Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.
 
 
 
 
 
 

807 lines
28 KiB

  1. # Copyright 2014-2016 OpenMarket Ltd
  2. #
  3. # Licensed under the Apache License, Version 2.0 (the "License");
  4. # you may not use this file except in compliance with the License.
  5. # You may obtain a copy of the License at
  6. #
  7. # http://www.apache.org/licenses/LICENSE-2.0
  8. #
  9. # Unless required by applicable law or agreed to in writing, software
  10. # distributed under the License is distributed on an "AS IS" BASIS,
  11. # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
  12. # See the License for the specific language governing permissions and
  13. # limitations under the License.
  14. from collections import OrderedDict
  15. from typing import Generator
  16. from unittest.mock import Mock, call, patch
  17. from twisted.internet import defer
  18. from synapse.storage._base import SQLBaseStore
  19. from synapse.storage.database import DatabasePool
  20. from synapse.storage.engines import create_engine
  21. from tests import unittest
  22. from tests.server import TestHomeServer
  23. from tests.utils import USE_POSTGRES_FOR_TESTS, default_config
  24. class SQLBaseStoreTestCase(unittest.TestCase):
  25. """Test the "simple" SQL generating methods in SQLBaseStore."""
  26. def setUp(self) -> None:
  27. # This is the Twisted connection pool.
  28. conn_pool = Mock(spec=["runInteraction", "runWithConnection"])
  29. self.mock_txn = Mock()
  30. if USE_POSTGRES_FOR_TESTS:
  31. # To avoid testing psycopg2 itself, patch execute_batch/execute_values
  32. # to assert how it is called.
  33. from psycopg2 import extras
  34. self.mock_execute_batch = Mock()
  35. self.execute_batch_patcher = patch.object(
  36. extras, "execute_batch", new=self.mock_execute_batch
  37. )
  38. self.execute_batch_patcher.start()
  39. self.mock_execute_values = Mock()
  40. self.execute_values_patcher = patch.object(
  41. extras, "execute_values", new=self.mock_execute_values
  42. )
  43. self.execute_values_patcher.start()
  44. self.mock_conn = Mock(
  45. spec_set=[
  46. "cursor",
  47. "rollback",
  48. "commit",
  49. "closed",
  50. "reconnect",
  51. "set_session",
  52. "encoding",
  53. ]
  54. )
  55. self.mock_conn.encoding = "UNICODE"
  56. else:
  57. self.mock_conn = Mock(spec_set=["cursor", "rollback", "commit"])
  58. self.mock_conn.cursor.return_value = self.mock_txn
  59. self.mock_txn.connection = self.mock_conn
  60. self.mock_conn.rollback.return_value = None
  61. # Our fake runInteraction just runs synchronously inline
  62. def runInteraction(func, *args, **kwargs) -> defer.Deferred: # type: ignore[no-untyped-def]
  63. return defer.succeed(func(self.mock_txn, *args, **kwargs))
  64. conn_pool.runInteraction = runInteraction
  65. def runWithConnection(func, *args, **kwargs): # type: ignore[no-untyped-def]
  66. return defer.succeed(func(self.mock_conn, *args, **kwargs))
  67. conn_pool.runWithConnection = runWithConnection
  68. config = default_config(name="test", parse=True)
  69. hs = TestHomeServer("test", config=config)
  70. if USE_POSTGRES_FOR_TESTS:
  71. db_config = {"name": "psycopg2", "args": {}}
  72. else:
  73. db_config = {"name": "sqlite3"}
  74. engine = create_engine(db_config)
  75. fake_engine = Mock(wraps=engine)
  76. fake_engine.in_transaction.return_value = False
  77. fake_engine.module.OperationalError = engine.module.OperationalError
  78. fake_engine.module.DatabaseError = engine.module.DatabaseError
  79. fake_engine.module.IntegrityError = engine.module.IntegrityError
  80. # Don't convert param style to make assertions easier.
  81. fake_engine.convert_param_style = lambda sql: sql
  82. # To fix isinstance(...) checks.
  83. fake_engine.__class__ = engine.__class__ # type: ignore[assignment]
  84. db = DatabasePool(Mock(), Mock(config=db_config), fake_engine)
  85. db._db_pool = conn_pool
  86. self.datastore = SQLBaseStore(db, None, hs) # type: ignore[arg-type]
  87. def tearDown(self) -> None:
  88. if USE_POSTGRES_FOR_TESTS:
  89. self.execute_batch_patcher.stop()
  90. self.execute_values_patcher.stop()
  91. @defer.inlineCallbacks
  92. def test_insert_1col(self) -> Generator["defer.Deferred[object]", object, None]:
  93. self.mock_txn.rowcount = 1
  94. yield defer.ensureDeferred(
  95. self.datastore.db_pool.simple_insert(
  96. table="tablename", values={"columname": "Value"}
  97. )
  98. )
  99. self.mock_txn.execute.assert_called_once_with(
  100. "INSERT INTO tablename (columname) VALUES(?)", ("Value",)
  101. )
  102. @defer.inlineCallbacks
  103. def test_insert_3cols(self) -> Generator["defer.Deferred[object]", object, None]:
  104. self.mock_txn.rowcount = 1
  105. yield defer.ensureDeferred(
  106. self.datastore.db_pool.simple_insert(
  107. table="tablename",
  108. # Use OrderedDict() so we can assert on the SQL generated
  109. values=OrderedDict([("colA", 1), ("colB", 2), ("colC", 3)]),
  110. )
  111. )
  112. self.mock_txn.execute.assert_called_once_with(
  113. "INSERT INTO tablename (colA, colB, colC) VALUES(?, ?, ?)", (1, 2, 3)
  114. )
  115. @defer.inlineCallbacks
  116. def test_insert_many(self) -> Generator["defer.Deferred[object]", object, None]:
  117. yield defer.ensureDeferred(
  118. self.datastore.db_pool.simple_insert_many(
  119. table="tablename",
  120. keys=(
  121. "col1",
  122. "col2",
  123. ),
  124. values=[
  125. (
  126. "val1",
  127. "val2",
  128. ),
  129. ("val3", "val4"),
  130. ],
  131. desc="",
  132. )
  133. )
  134. if USE_POSTGRES_FOR_TESTS:
  135. self.mock_execute_values.assert_called_once_with(
  136. self.mock_txn,
  137. "INSERT INTO tablename (col1, col2) VALUES ?",
  138. [("val1", "val2"), ("val3", "val4")],
  139. template=None,
  140. fetch=False,
  141. )
  142. else:
  143. self.mock_txn.executemany.assert_called_once_with(
  144. "INSERT INTO tablename (col1, col2) VALUES(?, ?)",
  145. [("val1", "val2"), ("val3", "val4")],
  146. )
  147. @defer.inlineCallbacks
  148. def test_insert_many_no_iterable(
  149. self,
  150. ) -> Generator["defer.Deferred[object]", object, None]:
  151. yield defer.ensureDeferred(
  152. self.datastore.db_pool.simple_insert_many(
  153. table="tablename",
  154. keys=(
  155. "col1",
  156. "col2",
  157. ),
  158. values=[],
  159. desc="",
  160. )
  161. )
  162. if USE_POSTGRES_FOR_TESTS:
  163. self.mock_execute_values.assert_not_called()
  164. else:
  165. self.mock_txn.executemany.assert_not_called()
  166. @defer.inlineCallbacks
  167. def test_select_one_1col(self) -> Generator["defer.Deferred[object]", object, None]:
  168. self.mock_txn.rowcount = 1
  169. self.mock_txn.__iter__ = Mock(return_value=iter([("Value",)]))
  170. value = yield defer.ensureDeferred(
  171. self.datastore.db_pool.simple_select_one_onecol(
  172. table="tablename", keyvalues={"keycol": "TheKey"}, retcol="retcol"
  173. )
  174. )
  175. self.assertEqual("Value", value)
  176. self.mock_txn.execute.assert_called_once_with(
  177. "SELECT retcol FROM tablename WHERE keycol = ?", ["TheKey"]
  178. )
  179. @defer.inlineCallbacks
  180. def test_select_one_3col(self) -> Generator["defer.Deferred[object]", object, None]:
  181. self.mock_txn.rowcount = 1
  182. self.mock_txn.fetchone.return_value = (1, 2, 3)
  183. ret = yield defer.ensureDeferred(
  184. self.datastore.db_pool.simple_select_one(
  185. table="tablename",
  186. keyvalues={"keycol": "TheKey"},
  187. retcols=["colA", "colB", "colC"],
  188. )
  189. )
  190. self.assertEqual((1, 2, 3), ret)
  191. self.mock_txn.execute.assert_called_once_with(
  192. "SELECT colA, colB, colC FROM tablename WHERE keycol = ?", ["TheKey"]
  193. )
  194. @defer.inlineCallbacks
  195. def test_select_one_missing(
  196. self,
  197. ) -> Generator["defer.Deferred[object]", object, None]:
  198. self.mock_txn.rowcount = 0
  199. self.mock_txn.fetchone.return_value = None
  200. ret = yield defer.ensureDeferred(
  201. self.datastore.db_pool.simple_select_one(
  202. table="tablename",
  203. keyvalues={"keycol": "Not here"},
  204. retcols=["colA"],
  205. allow_none=True,
  206. )
  207. )
  208. self.assertIsNone(ret)
  209. @defer.inlineCallbacks
  210. def test_select_list(self) -> Generator["defer.Deferred[object]", object, None]:
  211. self.mock_txn.rowcount = 3
  212. self.mock_txn.fetchall.return_value = [(1,), (2,), (3,)]
  213. self.mock_txn.description = (("colA", None, None, None, None, None, None),)
  214. ret = yield defer.ensureDeferred(
  215. self.datastore.db_pool.simple_select_list(
  216. table="tablename", keyvalues={"keycol": "A set"}, retcols=["colA"]
  217. )
  218. )
  219. self.assertEqual([(1,), (2,), (3,)], ret)
  220. self.mock_txn.execute.assert_called_once_with(
  221. "SELECT colA FROM tablename WHERE keycol = ?", ["A set"]
  222. )
  223. @defer.inlineCallbacks
  224. def test_select_many_batch(
  225. self,
  226. ) -> Generator["defer.Deferred[object]", object, None]:
  227. self.mock_txn.rowcount = 3
  228. self.mock_txn.fetchall.side_effect = [[(1,), (2,)], [(3,)]]
  229. ret = yield defer.ensureDeferred(
  230. self.datastore.db_pool.simple_select_many_batch(
  231. table="tablename",
  232. column="col1",
  233. iterable=("val1", "val2", "val3"),
  234. retcols=("col2",),
  235. keyvalues={"col3": "val4"},
  236. batch_size=2,
  237. )
  238. )
  239. self.mock_txn.execute.assert_has_calls(
  240. [
  241. call(
  242. "SELECT col2 FROM tablename WHERE col1 = ANY(?) AND col3 = ?",
  243. [["val1", "val2"], "val4"],
  244. ),
  245. call(
  246. "SELECT col2 FROM tablename WHERE col1 = ANY(?) AND col3 = ?",
  247. [["val3"], "val4"],
  248. ),
  249. ],
  250. )
  251. self.assertEqual([(1,), (2,), (3,)], ret)
  252. def test_select_many_no_iterable(self) -> None:
  253. self.mock_txn.rowcount = 3
  254. self.mock_txn.fetchall.side_effect = [(1,), (2,)]
  255. ret = self.datastore.db_pool.simple_select_many_txn(
  256. self.mock_txn,
  257. table="tablename",
  258. column="col1",
  259. iterable=(),
  260. retcols=("col2",),
  261. keyvalues={"col3": "val4"},
  262. )
  263. self.mock_txn.execute.assert_not_called()
  264. self.assertEqual([], ret)
  265. @defer.inlineCallbacks
  266. def test_update_one_1col(self) -> Generator["defer.Deferred[object]", object, None]:
  267. self.mock_txn.rowcount = 1
  268. yield defer.ensureDeferred(
  269. self.datastore.db_pool.simple_update_one(
  270. table="tablename",
  271. keyvalues={"keycol": "TheKey"},
  272. updatevalues={"columnname": "New Value"},
  273. )
  274. )
  275. self.mock_txn.execute.assert_called_once_with(
  276. "UPDATE tablename SET columnname = ? WHERE keycol = ?",
  277. ["New Value", "TheKey"],
  278. )
  279. @defer.inlineCallbacks
  280. def test_update_one_4cols(
  281. self,
  282. ) -> Generator["defer.Deferred[object]", object, None]:
  283. self.mock_txn.rowcount = 1
  284. yield defer.ensureDeferred(
  285. self.datastore.db_pool.simple_update_one(
  286. table="tablename",
  287. keyvalues=OrderedDict([("colA", 1), ("colB", 2)]),
  288. updatevalues=OrderedDict([("colC", 3), ("colD", 4)]),
  289. )
  290. )
  291. self.mock_txn.execute.assert_called_once_with(
  292. "UPDATE tablename SET colC = ?, colD = ? WHERE" " colA = ? AND colB = ?",
  293. [3, 4, 1, 2],
  294. )
  295. @defer.inlineCallbacks
  296. def test_update_many(self) -> Generator["defer.Deferred[object]", object, None]:
  297. yield defer.ensureDeferred(
  298. self.datastore.db_pool.simple_update_many(
  299. table="tablename",
  300. key_names=("col1", "col2"),
  301. key_values=[("val1", "val2")],
  302. value_names=("col3",),
  303. value_values=[("val3",)],
  304. desc="",
  305. )
  306. )
  307. if USE_POSTGRES_FOR_TESTS:
  308. self.mock_execute_batch.assert_called_once_with(
  309. self.mock_txn,
  310. "UPDATE tablename SET col3 = ? WHERE col1 = ? AND col2 = ?",
  311. [("val3", "val1", "val2")],
  312. )
  313. else:
  314. self.mock_txn.executemany.assert_called_once_with(
  315. "UPDATE tablename SET col3 = ? WHERE col1 = ? AND col2 = ?",
  316. [("val3", "val1", "val2")],
  317. )
  318. # key_values and value_values must be the same length.
  319. with self.assertRaises(ValueError):
  320. yield defer.ensureDeferred(
  321. self.datastore.db_pool.simple_update_many(
  322. table="tablename",
  323. key_names=("col1", "col2"),
  324. key_values=[("val1", "val2")],
  325. value_names=("col3",),
  326. value_values=[],
  327. desc="",
  328. )
  329. )
  330. @defer.inlineCallbacks
  331. def test_update_many_no_iterable(
  332. self,
  333. ) -> Generator["defer.Deferred[object]", object, None]:
  334. yield defer.ensureDeferred(
  335. self.datastore.db_pool.simple_update_many(
  336. table="tablename",
  337. key_names=("col1", "col2"),
  338. key_values=[],
  339. value_names=("col3",),
  340. value_values=[],
  341. desc="",
  342. )
  343. )
  344. if USE_POSTGRES_FOR_TESTS:
  345. self.mock_execute_batch.assert_not_called()
  346. else:
  347. self.mock_txn.executemany.assert_not_called()
  348. @defer.inlineCallbacks
  349. def test_delete_one(self) -> Generator["defer.Deferred[object]", object, None]:
  350. self.mock_txn.rowcount = 1
  351. yield defer.ensureDeferred(
  352. self.datastore.db_pool.simple_delete_one(
  353. table="tablename", keyvalues={"keycol": "Go away"}
  354. )
  355. )
  356. self.mock_txn.execute.assert_called_once_with(
  357. "DELETE FROM tablename WHERE keycol = ?", ["Go away"]
  358. )
  359. @defer.inlineCallbacks
  360. def test_delete_many(self) -> Generator["defer.Deferred[object]", object, None]:
  361. self.mock_txn.rowcount = 2
  362. result = yield defer.ensureDeferred(
  363. self.datastore.db_pool.simple_delete_many(
  364. table="tablename",
  365. column="col1",
  366. iterable=("val1", "val2"),
  367. keyvalues={"col2": "val3"},
  368. desc="",
  369. )
  370. )
  371. self.mock_txn.execute.assert_called_once_with(
  372. "DELETE FROM tablename WHERE col1 = ANY(?) AND col2 = ?",
  373. [["val1", "val2"], "val3"],
  374. )
  375. self.assertEqual(result, 2)
  376. @defer.inlineCallbacks
  377. def test_delete_many_no_iterable(
  378. self,
  379. ) -> Generator["defer.Deferred[object]", object, None]:
  380. result = yield defer.ensureDeferred(
  381. self.datastore.db_pool.simple_delete_many(
  382. table="tablename",
  383. column="col1",
  384. iterable=(),
  385. keyvalues={"col2": "val3"},
  386. desc="",
  387. )
  388. )
  389. self.mock_txn.execute.assert_not_called()
  390. self.assertEqual(result, 0)
  391. @defer.inlineCallbacks
  392. def test_delete_many_no_keyvalues(
  393. self,
  394. ) -> Generator["defer.Deferred[object]", object, None]:
  395. self.mock_txn.rowcount = 2
  396. result = yield defer.ensureDeferred(
  397. self.datastore.db_pool.simple_delete_many(
  398. table="tablename",
  399. column="col1",
  400. iterable=("val1", "val2"),
  401. keyvalues={},
  402. desc="",
  403. )
  404. )
  405. self.mock_txn.execute.assert_called_once_with(
  406. "DELETE FROM tablename WHERE col1 = ANY(?)", [["val1", "val2"]]
  407. )
  408. self.assertEqual(result, 2)
  409. @defer.inlineCallbacks
  410. def test_upsert(self) -> Generator["defer.Deferred[object]", object, None]:
  411. self.mock_txn.rowcount = 1
  412. result = yield defer.ensureDeferred(
  413. self.datastore.db_pool.simple_upsert(
  414. table="tablename",
  415. keyvalues={"columnname": "oldvalue"},
  416. values={"othercol": "newvalue"},
  417. )
  418. )
  419. self.mock_txn.execute.assert_called_once_with(
  420. "INSERT INTO tablename (columnname, othercol) VALUES (?, ?) ON CONFLICT (columnname) DO UPDATE SET othercol=EXCLUDED.othercol",
  421. ["oldvalue", "newvalue"],
  422. )
  423. self.assertTrue(result)
  424. @defer.inlineCallbacks
  425. def test_upsert_no_values(
  426. self,
  427. ) -> Generator["defer.Deferred[object]", object, None]:
  428. self.mock_txn.rowcount = 1
  429. result = yield defer.ensureDeferred(
  430. self.datastore.db_pool.simple_upsert(
  431. table="tablename",
  432. keyvalues={"columnname": "value"},
  433. values={},
  434. insertion_values={"columnname": "value"},
  435. )
  436. )
  437. self.mock_txn.execute.assert_called_once_with(
  438. "INSERT INTO tablename (columnname) VALUES (?) ON CONFLICT (columnname) DO NOTHING",
  439. ["value"],
  440. )
  441. self.assertTrue(result)
  442. @defer.inlineCallbacks
  443. def test_upsert_with_insertion(
  444. self,
  445. ) -> Generator["defer.Deferred[object]", object, None]:
  446. self.mock_txn.rowcount = 1
  447. result = yield defer.ensureDeferred(
  448. self.datastore.db_pool.simple_upsert(
  449. table="tablename",
  450. keyvalues={"columnname": "oldvalue"},
  451. values={"othercol": "newvalue"},
  452. insertion_values={"thirdcol": "insertionval"},
  453. )
  454. )
  455. self.mock_txn.execute.assert_called_once_with(
  456. "INSERT INTO tablename (columnname, thirdcol, othercol) VALUES (?, ?, ?) ON CONFLICT (columnname) DO UPDATE SET othercol=EXCLUDED.othercol",
  457. ["oldvalue", "insertionval", "newvalue"],
  458. )
  459. self.assertTrue(result)
  460. @defer.inlineCallbacks
  461. def test_upsert_with_where(
  462. self,
  463. ) -> Generator["defer.Deferred[object]", object, None]:
  464. self.mock_txn.rowcount = 1
  465. result = yield defer.ensureDeferred(
  466. self.datastore.db_pool.simple_upsert(
  467. table="tablename",
  468. keyvalues={"columnname": "oldvalue"},
  469. values={"othercol": "newvalue"},
  470. where_clause="thirdcol IS NULL",
  471. )
  472. )
  473. self.mock_txn.execute.assert_called_once_with(
  474. "INSERT INTO tablename (columnname, othercol) VALUES (?, ?) ON CONFLICT (columnname) WHERE thirdcol IS NULL DO UPDATE SET othercol=EXCLUDED.othercol",
  475. ["oldvalue", "newvalue"],
  476. )
  477. self.assertTrue(result)
  478. @defer.inlineCallbacks
  479. def test_upsert_many(self) -> Generator["defer.Deferred[object]", object, None]:
  480. yield defer.ensureDeferred(
  481. self.datastore.db_pool.simple_upsert_many(
  482. table="tablename",
  483. key_names=["keycol1", "keycol2"],
  484. key_values=[["keyval1", "keyval2"], ["keyval3", "keyval4"]],
  485. value_names=["valuecol3"],
  486. value_values=[["val5"], ["val6"]],
  487. desc="",
  488. )
  489. )
  490. if USE_POSTGRES_FOR_TESTS:
  491. self.mock_execute_values.assert_called_once_with(
  492. self.mock_txn,
  493. "INSERT INTO tablename (keycol1, keycol2, valuecol3) VALUES ? ON CONFLICT (keycol1, keycol2) DO UPDATE SET valuecol3=EXCLUDED.valuecol3",
  494. [("keyval1", "keyval2", "val5"), ("keyval3", "keyval4", "val6")],
  495. template=None,
  496. fetch=False,
  497. )
  498. else:
  499. self.mock_txn.executemany.assert_called_once_with(
  500. "INSERT INTO tablename (keycol1, keycol2, valuecol3) VALUES (?, ?, ?) ON CONFLICT (keycol1, keycol2) DO UPDATE SET valuecol3=EXCLUDED.valuecol3",
  501. [("keyval1", "keyval2", "val5"), ("keyval3", "keyval4", "val6")],
  502. )
  503. @defer.inlineCallbacks
  504. def test_upsert_many_no_values(
  505. self,
  506. ) -> Generator["defer.Deferred[object]", object, None]:
  507. yield defer.ensureDeferred(
  508. self.datastore.db_pool.simple_upsert_many(
  509. table="tablename",
  510. key_names=["columnname"],
  511. key_values=[["oldvalue"]],
  512. value_names=[],
  513. value_values=[],
  514. desc="",
  515. )
  516. )
  517. if USE_POSTGRES_FOR_TESTS:
  518. self.mock_execute_values.assert_called_once_with(
  519. self.mock_txn,
  520. "INSERT INTO tablename (columnname) VALUES ? ON CONFLICT (columnname) DO NOTHING",
  521. [("oldvalue",)],
  522. template=None,
  523. fetch=False,
  524. )
  525. else:
  526. self.mock_txn.executemany.assert_called_once_with(
  527. "INSERT INTO tablename (columnname) VALUES (?) ON CONFLICT (columnname) DO NOTHING",
  528. [("oldvalue",)],
  529. )
  530. @defer.inlineCallbacks
  531. def test_upsert_emulated_no_values_exists(
  532. self,
  533. ) -> Generator["defer.Deferred[object]", object, None]:
  534. self.datastore.db_pool._unsafe_to_upsert_tables.add("tablename")
  535. self.mock_txn.fetchall.return_value = [(1,)]
  536. result = yield defer.ensureDeferred(
  537. self.datastore.db_pool.simple_upsert(
  538. table="tablename",
  539. keyvalues={"columnname": "value"},
  540. values={},
  541. insertion_values={"columnname": "value"},
  542. )
  543. )
  544. if USE_POSTGRES_FOR_TESTS:
  545. self.mock_txn.execute.assert_has_calls(
  546. [
  547. call("LOCK TABLE tablename in EXCLUSIVE MODE", ()),
  548. call("SELECT 1 FROM tablename WHERE columnname = ?", ["value"]),
  549. ]
  550. )
  551. else:
  552. self.mock_txn.execute.assert_called_once_with(
  553. "SELECT 1 FROM tablename WHERE columnname = ?", ["value"]
  554. )
  555. self.assertFalse(result)
  556. @defer.inlineCallbacks
  557. def test_upsert_emulated_no_values_not_exists(
  558. self,
  559. ) -> Generator["defer.Deferred[object]", object, None]:
  560. self.datastore.db_pool._unsafe_to_upsert_tables.add("tablename")
  561. self.mock_txn.fetchall.return_value = []
  562. self.mock_txn.rowcount = 1
  563. result = yield defer.ensureDeferred(
  564. self.datastore.db_pool.simple_upsert(
  565. table="tablename",
  566. keyvalues={"columnname": "value"},
  567. values={},
  568. insertion_values={"columnname": "value"},
  569. )
  570. )
  571. self.mock_txn.execute.assert_has_calls(
  572. [
  573. call(
  574. "SELECT 1 FROM tablename WHERE columnname = ?",
  575. ["value"],
  576. ),
  577. call("INSERT INTO tablename (columnname) VALUES (?)", ["value"]),
  578. ],
  579. )
  580. self.assertTrue(result)
  581. @defer.inlineCallbacks
  582. def test_upsert_emulated_with_insertion_exists(
  583. self,
  584. ) -> Generator["defer.Deferred[object]", object, None]:
  585. self.datastore.db_pool._unsafe_to_upsert_tables.add("tablename")
  586. self.mock_txn.rowcount = 1
  587. result = yield defer.ensureDeferred(
  588. self.datastore.db_pool.simple_upsert(
  589. table="tablename",
  590. keyvalues={"columnname": "oldvalue"},
  591. values={"othercol": "newvalue"},
  592. insertion_values={"thirdcol": "insertionval"},
  593. )
  594. )
  595. if USE_POSTGRES_FOR_TESTS:
  596. self.mock_txn.execute.assert_has_calls(
  597. [
  598. call("LOCK TABLE tablename in EXCLUSIVE MODE", ()),
  599. call(
  600. "UPDATE tablename SET othercol = ? WHERE columnname = ?",
  601. ["newvalue", "oldvalue"],
  602. ),
  603. ]
  604. )
  605. else:
  606. self.mock_txn.execute.assert_called_once_with(
  607. "UPDATE tablename SET othercol = ? WHERE columnname = ?",
  608. ["newvalue", "oldvalue"],
  609. )
  610. self.assertTrue(result)
  611. @defer.inlineCallbacks
  612. def test_upsert_emulated_with_insertion_not_exists(
  613. self,
  614. ) -> Generator["defer.Deferred[object]", object, None]:
  615. self.datastore.db_pool._unsafe_to_upsert_tables.add("tablename")
  616. self.mock_txn.rowcount = 0
  617. result = yield defer.ensureDeferred(
  618. self.datastore.db_pool.simple_upsert(
  619. table="tablename",
  620. keyvalues={"columnname": "oldvalue"},
  621. values={"othercol": "newvalue"},
  622. insertion_values={"thirdcol": "insertionval"},
  623. )
  624. )
  625. self.mock_txn.execute.assert_has_calls(
  626. [
  627. call(
  628. "UPDATE tablename SET othercol = ? WHERE columnname = ?",
  629. ["newvalue", "oldvalue"],
  630. ),
  631. call(
  632. "INSERT INTO tablename (columnname, othercol, thirdcol) VALUES (?, ?, ?)",
  633. ["oldvalue", "newvalue", "insertionval"],
  634. ),
  635. ]
  636. )
  637. self.assertTrue(result)
  638. @defer.inlineCallbacks
  639. def test_upsert_emulated_with_where(
  640. self,
  641. ) -> Generator["defer.Deferred[object]", object, None]:
  642. self.datastore.db_pool._unsafe_to_upsert_tables.add("tablename")
  643. self.mock_txn.rowcount = 1
  644. result = yield defer.ensureDeferred(
  645. self.datastore.db_pool.simple_upsert(
  646. table="tablename",
  647. keyvalues={"columnname": "oldvalue"},
  648. values={"othercol": "newvalue"},
  649. where_clause="thirdcol IS NULL",
  650. )
  651. )
  652. if USE_POSTGRES_FOR_TESTS:
  653. self.mock_txn.execute.assert_has_calls(
  654. [
  655. call("LOCK TABLE tablename in EXCLUSIVE MODE", ()),
  656. call(
  657. "UPDATE tablename SET othercol = ? WHERE columnname = ? AND thirdcol IS NULL",
  658. ["newvalue", "oldvalue"],
  659. ),
  660. ]
  661. )
  662. else:
  663. self.mock_txn.execute.assert_called_once_with(
  664. "UPDATE tablename SET othercol = ? WHERE columnname = ? AND thirdcol IS NULL",
  665. ["newvalue", "oldvalue"],
  666. )
  667. self.assertTrue(result)
  668. @defer.inlineCallbacks
  669. def test_upsert_emulated_with_where_no_values(
  670. self,
  671. ) -> Generator["defer.Deferred[object]", object, None]:
  672. self.datastore.db_pool._unsafe_to_upsert_tables.add("tablename")
  673. self.mock_txn.rowcount = 1
  674. result = yield defer.ensureDeferred(
  675. self.datastore.db_pool.simple_upsert(
  676. table="tablename",
  677. keyvalues={"columnname": "oldvalue"},
  678. values={},
  679. where_clause="thirdcol IS NULL",
  680. )
  681. )
  682. if USE_POSTGRES_FOR_TESTS:
  683. self.mock_txn.execute.assert_has_calls(
  684. [
  685. call("LOCK TABLE tablename in EXCLUSIVE MODE", ()),
  686. call(
  687. "SELECT 1 FROM tablename WHERE columnname = ? AND thirdcol IS NULL",
  688. ["oldvalue"],
  689. ),
  690. ]
  691. )
  692. else:
  693. self.mock_txn.execute.assert_called_once_with(
  694. "SELECT 1 FROM tablename WHERE columnname = ? AND thirdcol IS NULL",
  695. ["oldvalue"],
  696. )
  697. self.assertFalse(result)