smoke.py 17 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438439440441442443444445446447448449450451452453454455456457458459460461462463464
  1. # Copyright Amethyst Reese
  2. # Licensed under the MIT license
  3. import asyncio
  4. import sqlite3
  5. from pathlib import Path
  6. from sqlite3 import OperationalError
  7. from tempfile import TemporaryDirectory
  8. from threading import Thread
  9. from unittest import IsolatedAsyncioTestCase, SkipTest
  10. from unittest.mock import patch
  11. import aiosqlite
  12. from .helpers import setup_logger
  13. class SmokeTest(IsolatedAsyncioTestCase):
  14. @classmethod
  15. def setUpClass(cls):
  16. setup_logger()
  17. def setUp(self):
  18. td = TemporaryDirectory()
  19. self.addCleanup(td.cleanup)
  20. self.db = Path(td.name).resolve() / "test.db"
  21. async def test_connection_await(self):
  22. db = await aiosqlite.connect(self.db)
  23. self.assertIsInstance(db, aiosqlite.Connection)
  24. async with db.execute("select 1, 2") as cursor:
  25. rows = await cursor.fetchall()
  26. self.assertEqual(rows, [(1, 2)])
  27. await db.close()
  28. async def test_connection_context(self):
  29. async with aiosqlite.connect(self.db) as db:
  30. self.assertIsInstance(db, aiosqlite.Connection)
  31. async with db.execute("select 1, 2") as cursor:
  32. rows = await cursor.fetchall()
  33. self.assertEqual(rows, [(1, 2)])
  34. async def test_connection_locations(self):
  35. TEST_DB = self.db.as_posix()
  36. class Fake: # pylint: disable=too-few-public-methods
  37. def __str__(self):
  38. return TEST_DB
  39. locs = (Path(TEST_DB), TEST_DB, TEST_DB.encode(), Fake())
  40. async with aiosqlite.connect(locs[0]) as db:
  41. await db.execute("create table foo (i integer, k integer)")
  42. await db.execute("insert into foo (i, k) values (1, 5)")
  43. await db.commit()
  44. cursor = await db.execute("select * from foo")
  45. rows = await cursor.fetchall()
  46. for loc in locs:
  47. async with aiosqlite.connect(loc) as db:
  48. cursor = await db.execute("select * from foo")
  49. self.assertEqual(await cursor.fetchall(), rows)
  50. async def test_multiple_connections(self):
  51. async with aiosqlite.connect(self.db) as db:
  52. await db.execute(
  53. "create table multiple_connections "
  54. "(i integer primary key asc, k integer)"
  55. )
  56. async def do_one_conn(i):
  57. async with aiosqlite.connect(self.db) as db:
  58. await db.execute("insert into multiple_connections (k) values (?)", [i])
  59. await db.commit()
  60. await asyncio.gather(*[do_one_conn(i) for i in range(10)])
  61. async with aiosqlite.connect(self.db) as db:
  62. cursor = await db.execute("select * from multiple_connections")
  63. rows = await cursor.fetchall()
  64. assert len(rows) == 10
  65. async def test_multiple_queries(self):
  66. async with aiosqlite.connect(self.db) as db:
  67. await db.execute(
  68. "create table multiple_queries "
  69. "(i integer primary key asc, k integer)"
  70. )
  71. await asyncio.gather(
  72. *[
  73. db.execute("insert into multiple_queries (k) values (?)", [i])
  74. for i in range(10)
  75. ]
  76. )
  77. await db.commit()
  78. async with aiosqlite.connect(self.db) as db:
  79. cursor = await db.execute("select * from multiple_queries")
  80. rows = await cursor.fetchall()
  81. assert len(rows) == 10
  82. async def test_iterable_cursor(self):
  83. async with aiosqlite.connect(self.db) as db:
  84. cursor = await db.cursor()
  85. await cursor.execute(
  86. "create table iterable_cursor " "(i integer primary key asc, k integer)"
  87. )
  88. await cursor.executemany(
  89. "insert into iterable_cursor (k) values (?)", [[i] for i in range(10)]
  90. )
  91. await db.commit()
  92. async with aiosqlite.connect(self.db) as db:
  93. cursor = await db.execute("select * from iterable_cursor")
  94. rows = []
  95. async for row in cursor:
  96. rows.append(row)
  97. assert len(rows) == 10
  98. async def test_multi_loop_usage(self):
  99. results = {}
  100. def runner(k, conn):
  101. async def query():
  102. async with conn.execute("select * from foo") as cursor:
  103. rows = await cursor.fetchall()
  104. self.assertEqual(len(rows), 2)
  105. return rows
  106. with self.subTest(k):
  107. loop = asyncio.new_event_loop()
  108. rows = loop.run_until_complete(query())
  109. loop.close()
  110. results[k] = rows
  111. async with aiosqlite.connect(":memory:") as db:
  112. await db.execute("create table foo (id int, name varchar)")
  113. await db.execute(
  114. "insert into foo values (?, ?), (?, ?)", (1, "Sally", 2, "Janet")
  115. )
  116. await db.commit()
  117. threads = [Thread(target=runner, args=(k, db)) for k in range(4)]
  118. for thread in threads:
  119. thread.start()
  120. for thread in threads:
  121. thread.join()
  122. self.assertEqual(len(results), 4)
  123. for rows in results.values():
  124. self.assertEqual(len(rows), 2)
  125. async def test_context_cursor(self):
  126. async with aiosqlite.connect(self.db) as db:
  127. async with db.cursor() as cursor:
  128. await cursor.execute(
  129. "create table context_cursor "
  130. "(i integer primary key asc, k integer)"
  131. )
  132. await cursor.executemany(
  133. "insert into context_cursor (k) values (?)",
  134. [[i] for i in range(10)],
  135. )
  136. await db.commit()
  137. async with aiosqlite.connect(self.db) as db:
  138. async with db.execute("select * from context_cursor") as cursor:
  139. rows = []
  140. async for row in cursor:
  141. rows.append(row)
  142. assert len(rows) == 10
  143. async def test_cursor_return_self(self):
  144. async with aiosqlite.connect(self.db) as db:
  145. cursor = await db.cursor()
  146. result = await cursor.execute(
  147. "create table test_cursor_return_self (i integer, k integer)"
  148. )
  149. self.assertEqual(result, cursor, "cursor execute returns itself")
  150. result = await cursor.executemany(
  151. "insert into test_cursor_return_self values (?, ?)", [(1, 1), (2, 2)]
  152. )
  153. self.assertEqual(result, cursor)
  154. result = await cursor.executescript(
  155. "insert into test_cursor_return_self values (3, 3);"
  156. "insert into test_cursor_return_self values (4, 4);"
  157. "insert into test_cursor_return_self values (5, 5);"
  158. )
  159. self.assertEqual(result, cursor)
  160. async def test_connection_properties(self):
  161. async with aiosqlite.connect(self.db) as db:
  162. self.assertEqual(db.total_changes, 0)
  163. async with db.cursor() as cursor:
  164. self.assertFalse(db.in_transaction)
  165. await cursor.execute(
  166. "create table test_properties "
  167. "(i integer primary key asc, k integer, d text)"
  168. )
  169. await cursor.execute(
  170. "insert into test_properties (k, d) values (1, 'hi')"
  171. )
  172. self.assertTrue(db.in_transaction)
  173. await db.commit()
  174. self.assertFalse(db.in_transaction)
  175. self.assertEqual(db.total_changes, 1)
  176. self.assertIsNone(db.row_factory)
  177. self.assertEqual(db.text_factory, str)
  178. async with db.cursor() as cursor:
  179. await cursor.execute("select * from test_properties")
  180. row = await cursor.fetchone()
  181. self.assertIsInstance(row, tuple)
  182. self.assertEqual(row, (1, 1, "hi"))
  183. with self.assertRaises(TypeError):
  184. _ = row["k"]
  185. async with db.cursor() as cursor:
  186. cursor.row_factory = aiosqlite.Row
  187. self.assertEqual(cursor.row_factory, aiosqlite.Row)
  188. await cursor.execute("select * from test_properties")
  189. row = await cursor.fetchone()
  190. self.assertIsInstance(row, aiosqlite.Row)
  191. self.assertEqual(row[1], 1)
  192. self.assertEqual(row[2], "hi")
  193. self.assertEqual(row["k"], 1)
  194. self.assertEqual(row["d"], "hi")
  195. db.row_factory = aiosqlite.Row
  196. db.text_factory = bytes
  197. self.assertEqual(db.row_factory, aiosqlite.Row)
  198. self.assertEqual(db.text_factory, bytes)
  199. async with db.cursor() as cursor:
  200. await cursor.execute("select * from test_properties")
  201. row = await cursor.fetchone()
  202. self.assertIsInstance(row, aiosqlite.Row)
  203. self.assertEqual(row[1], 1)
  204. self.assertEqual(row[2], b"hi")
  205. self.assertEqual(row["k"], 1)
  206. self.assertEqual(row["d"], b"hi")
  207. async def test_fetch_all(self):
  208. async with aiosqlite.connect(self.db) as db:
  209. await db.execute(
  210. "create table test_fetch_all (i integer primary key asc, k integer)"
  211. )
  212. await db.execute(
  213. "insert into test_fetch_all (k) values (10), (24), (16), (32)"
  214. )
  215. await db.commit()
  216. async with aiosqlite.connect(self.db) as db:
  217. cursor = await db.execute("select k from test_fetch_all where k < 30")
  218. rows = await cursor.fetchall()
  219. self.assertEqual(rows, [(10,), (24,), (16,)])
  220. async def test_enable_load_extension(self):
  221. """Assert that after enabling extension loading, they can be loaded"""
  222. async with aiosqlite.connect(self.db) as db:
  223. try:
  224. await db.enable_load_extension(True)
  225. await db.load_extension("test")
  226. except OperationalError as e:
  227. assert "not authorized" not in e.args
  228. except AttributeError as e:
  229. raise SkipTest(
  230. "python was not compiled with sqlite3 "
  231. "extension support, so we can't test it"
  232. ) from e
  233. async def test_set_progress_handler(self):
  234. """
  235. Assert that after setting a progress handler returning 1, DB operations are aborted
  236. """
  237. async with aiosqlite.connect(self.db) as db:
  238. await db.set_progress_handler(lambda: 1, 1)
  239. with self.assertRaises(OperationalError):
  240. await db.execute(
  241. "create table test_progress_handler (i integer primary key asc, k integer)"
  242. )
  243. async def test_create_function(self):
  244. """Assert that after creating a custom function, it can be used"""
  245. def no_arg():
  246. return "no arg"
  247. def one_arg(num):
  248. return num * 2
  249. async with aiosqlite.connect(self.db) as db:
  250. await db.create_function("no_arg", 0, no_arg)
  251. await db.create_function("one_arg", 1, one_arg)
  252. async with db.execute("SELECT no_arg();") as res:
  253. row = await res.fetchone()
  254. self.assertEqual(row[0], "no arg")
  255. async with db.execute("SELECT one_arg(10);") as res:
  256. row = await res.fetchone()
  257. self.assertEqual(row[0], 20)
  258. async def test_create_function_deterministic(self):
  259. """Assert that after creating a deterministic custom function, it can be used.
  260. https://sqlite.org/deterministic.html
  261. """
  262. def one_arg(num):
  263. return num * 2
  264. async with aiosqlite.connect(self.db) as db:
  265. await db.create_function("one_arg", 1, one_arg, deterministic=True)
  266. await db.execute("create table foo (id int, bar int)")
  267. # Non-deterministic functions cannot be used in indexes
  268. await db.execute("create index t on foo(one_arg(bar))")
  269. async def test_set_trace_callback(self):
  270. statements = []
  271. def callback(statement: str):
  272. statements.append(statement)
  273. async with aiosqlite.connect(self.db) as db:
  274. await db.set_trace_callback(callback)
  275. await db.execute("select 10")
  276. self.assertIn("select 10", statements)
  277. async def test_connect_error(self):
  278. bad_db = Path("/something/that/shouldnt/exist.db")
  279. with self.assertRaisesRegex(OperationalError, "unable to open database"):
  280. async with aiosqlite.connect(bad_db) as db:
  281. self.assertIsNone(db) # should never be reached
  282. with self.assertRaisesRegex(OperationalError, "unable to open database"):
  283. await aiosqlite.connect(bad_db)
  284. async def test_connect_base_exception(self):
  285. # Check if connect task is cancelled, thread is properly closed.
  286. def _raise_cancelled_error(*_, **__):
  287. raise asyncio.CancelledError("I changed my mind")
  288. connection = aiosqlite.Connection(lambda: sqlite3.connect(":memory:"), 64)
  289. with (
  290. patch.object(sqlite3, "connect", side_effect=_raise_cancelled_error),
  291. self.assertRaisesRegex(asyncio.CancelledError, "I changed my mind"),
  292. ):
  293. async with connection:
  294. ...
  295. # Terminate the thread here if the test fails to have a clear error.
  296. if connection._running:
  297. connection._stop_running()
  298. raise AssertionError("connection thread was not stopped")
  299. async def test_iterdump(self):
  300. async with aiosqlite.connect(":memory:") as db:
  301. await db.execute("create table foo (i integer, k charvar(250))")
  302. await db.executemany(
  303. "insert into foo values (?, ?)", [(1, "hello"), (2, "world")]
  304. )
  305. lines = [line async for line in db.iterdump()]
  306. self.assertEqual(
  307. lines,
  308. [
  309. "BEGIN TRANSACTION;",
  310. "CREATE TABLE foo (i integer, k charvar(250));",
  311. "INSERT INTO \"foo\" VALUES(1,'hello');",
  312. "INSERT INTO \"foo\" VALUES(2,'world');",
  313. "COMMIT;",
  314. ],
  315. )
  316. async def test_cursor_on_closed_connection(self):
  317. db = await aiosqlite.connect(self.db)
  318. cursor = await db.execute("select 1, 2")
  319. await db.close()
  320. with self.assertRaisesRegex(ValueError, "Connection closed"):
  321. await cursor.fetchall()
  322. with self.assertRaisesRegex(ValueError, "Connection closed"):
  323. await cursor.fetchall()
  324. async def test_cursor_on_closed_connection_loop(self):
  325. db = await aiosqlite.connect(self.db)
  326. cursor = await db.execute("select 1, 2")
  327. tasks = []
  328. for i in range(100):
  329. if i == 50:
  330. tasks.append(asyncio.ensure_future(db.close()))
  331. tasks.append(asyncio.ensure_future(cursor.fetchall()))
  332. for task in tasks:
  333. try:
  334. await task
  335. except sqlite3.ProgrammingError:
  336. pass
  337. async def test_close_twice(self):
  338. db = await aiosqlite.connect(self.db)
  339. await db.close()
  340. # no error
  341. await db.close()
  342. async def test_backup_aiosqlite(self):
  343. def progress(a, b, c):
  344. print(a, b, c)
  345. async with (
  346. aiosqlite.connect(":memory:") as db1,
  347. aiosqlite.connect(":memory:") as db2,
  348. ):
  349. await db1.execute("create table foo (i integer, k charvar(250))")
  350. await db1.executemany(
  351. "insert into foo values (?, ?)", [(1, "hello"), (2, "world")]
  352. )
  353. await db1.commit()
  354. with self.assertRaisesRegex(OperationalError, "no such table: foo"):
  355. await db2.execute("select * from foo")
  356. await db1.backup(db2, progress=progress)
  357. async with db2.execute("select * from foo") as cursor:
  358. rows = await cursor.fetchall()
  359. self.assertEqual(rows, [(1, "hello"), (2, "world")])
  360. async def test_backup_sqlite(self):
  361. async with aiosqlite.connect(":memory:") as db1:
  362. with sqlite3.connect(":memory:") as db2:
  363. await db1.execute("create table foo (i integer, k charvar(250))")
  364. await db1.executemany(
  365. "insert into foo values (?, ?)", [(1, "hello"), (2, "world")]
  366. )
  367. await db1.commit()
  368. with self.assertRaisesRegex(OperationalError, "no such table: foo"):
  369. db2.execute("select * from foo")
  370. await db1.backup(db2)
  371. cursor = db2.execute("select * from foo")
  372. rows = cursor.fetchall()
  373. self.assertEqual(rows, [(1, "hello"), (2, "world")])