import asyncio import datetime import decimal import enum import functools import gc import itertools import os import sqlite3 from typing import MutableMapping from unittest.mock import MagicMock, patch import pytest import sqlalchemy from databases import Database, DatabaseURL assert "TEST_DATABASE_URLS" in os.environ, "TEST_DATABASE_URLS is not set." DATABASE_URLS = [url.strip() for url in os.environ["TEST_DATABASE_URLS"].split(",")] class AsyncMock(MagicMock): async def __call__(self, *args, **kwargs): return super(AsyncMock, self).__call__(*args, **kwargs) class MyEpochType(sqlalchemy.types.TypeDecorator): impl = sqlalchemy.Integer epoch = datetime.date(1970, 1, 1) def process_bind_param(self, value, dialect): return (value - self.epoch).days def process_result_value(self, value, dialect): return self.epoch + datetime.timedelta(days=value) metadata = sqlalchemy.MetaData() notes = sqlalchemy.Table( "notes", metadata, sqlalchemy.Column("id", sqlalchemy.Integer, primary_key=True), sqlalchemy.Column("text", sqlalchemy.String(length=100)), sqlalchemy.Column("completed", sqlalchemy.Boolean), ) # Used to test DateTime articles = sqlalchemy.Table( "articles", metadata, sqlalchemy.Column("id", sqlalchemy.Integer, primary_key=True), sqlalchemy.Column("title", sqlalchemy.String(length=100)), sqlalchemy.Column("published", sqlalchemy.DateTime), ) # Used to test Date events = sqlalchemy.Table( "events", metadata, sqlalchemy.Column("id", sqlalchemy.Integer, primary_key=True), sqlalchemy.Column("date", sqlalchemy.Date), ) # Used to test Time daily_schedule = sqlalchemy.Table( "daily_schedule", metadata, sqlalchemy.Column("id", sqlalchemy.Integer, primary_key=True), sqlalchemy.Column("time", sqlalchemy.Time), ) class TshirtSize(enum.Enum): SMALL = "SMALL" MEDIUM = "MEDIUM" LARGE = "LARGE" XL = "XL" class TshirtColor(enum.Enum): BLUE = 0 GREEN = 1 YELLOW = 2 RED = 3 # Used to test Enum tshirt_size = sqlalchemy.Table( "tshirt_size", metadata, sqlalchemy.Column("id", sqlalchemy.Integer, primary_key=True), sqlalchemy.Column("size", sqlalchemy.Enum(TshirtSize)), sqlalchemy.Column("color", sqlalchemy.Enum(TshirtColor)), ) # Used to test JSON session = sqlalchemy.Table( "session", metadata, sqlalchemy.Column("id", sqlalchemy.Integer, primary_key=True), sqlalchemy.Column("data", sqlalchemy.JSON), ) # Used to test custom column types custom_date = sqlalchemy.Table( "custom_date", metadata, sqlalchemy.Column("id", sqlalchemy.Integer, primary_key=True), sqlalchemy.Column("title", sqlalchemy.String(length=100)), sqlalchemy.Column("published", MyEpochType), ) # Used to test Numeric prices = sqlalchemy.Table( "prices", metadata, sqlalchemy.Column("id", sqlalchemy.Integer, primary_key=True), sqlalchemy.Column("price", sqlalchemy.Numeric(precision=30, scale=20)), ) @pytest.fixture(autouse=True, scope="function") def create_test_database(): # Create test databases with tables creation for url in DATABASE_URLS: database_url = DatabaseURL(url) if database_url.scheme in ["mysql", "mysql+aiomysql", "mysql+asyncmy"]: url = str(database_url.replace(driver="pymysql")) elif database_url.scheme in [ "postgresql+aiopg", "sqlite+aiosqlite", "postgresql+asyncpg", ]: url = str(database_url.replace(driver=None)) engine = sqlalchemy.create_engine(url) metadata.create_all(engine) # Run the test suite yield # Drop test databases for url in DATABASE_URLS: database_url = DatabaseURL(url) if database_url.scheme in ["mysql", "mysql+aiomysql", "mysql+asyncmy"]: url = str(database_url.replace(driver="pymysql")) elif database_url.scheme in [ "postgresql+aiopg", "sqlite+aiosqlite", "postgresql+asyncpg", ]: url = str(database_url.replace(driver=None)) engine = sqlalchemy.create_engine(url) metadata.drop_all(engine) # Run garbage collection to ensure any in-memory databases are dropped gc.collect() def async_adapter(wrapped_func): """ Decorator used to run async test cases. """ @functools.wraps(wrapped_func) def run_sync(*args, **kwargs): loop = asyncio.new_event_loop() task = wrapped_func(*args, **kwargs) return loop.run_until_complete(task) return run_sync @pytest.mark.parametrize("database_url", DATABASE_URLS) @async_adapter async def test_queries(database_url): """ Test that the basic `execute()`, `execute_many()`, `fetch_all()``, and `fetch_one()` interfaces are all supported (using SQLAlchemy core). """ async with Database(database_url) as database: async with database.transaction(force_rollback=True): # execute() query = notes.insert() values = {"text": "example1", "completed": True} await database.execute(query, values) # execute_many() query = notes.insert() values = [ {"text": "example2", "completed": False}, {"text": "example3", "completed": True}, ] await database.execute_many(query, values) # fetch_all() query = notes.select() results = await database.fetch_all(query=query) assert len(results) == 3 assert results[0]["text"] == "example1" assert results[0]["completed"] == True assert results[1]["text"] == "example2" assert results[1]["completed"] == False assert results[2]["text"] == "example3" assert results[2]["completed"] == True # fetch_one() query = notes.select() result = await database.fetch_one(query=query) assert result["text"] == "example1" assert result["completed"] == True # fetch_val() query = sqlalchemy.sql.select(*[notes.c.text]) result = await database.fetch_val(query=query) assert result == "example1" # fetch_val() with no rows query = sqlalchemy.sql.select(*[notes.c.text]).where( notes.c.text == "impossible" ) result = await database.fetch_val(query=query) assert result is None # fetch_val() with a different column query = sqlalchemy.sql.select(*[notes.c.id, notes.c.text]) result = await database.fetch_val(query=query, column=1) assert result == "example1" # row access (needed to maintain test coverage for Record.__getitem__ in postgres backend) query = sqlalchemy.sql.select(*[notes.c.text]) result = await database.fetch_one(query=query) assert result["text"] == "example1" assert result[0] == "example1" # iterate() query = notes.select() iterate_results = [] async for result in database.iterate(query=query): iterate_results.append(result) assert len(iterate_results) == 3 assert iterate_results[0]["text"] == "example1" assert iterate_results[0]["completed"] == True assert iterate_results[1]["text"] == "example2" assert iterate_results[1]["completed"] == False assert iterate_results[2]["text"] == "example3" assert iterate_results[2]["completed"] == True @pytest.mark.parametrize("database_url", DATABASE_URLS) @async_adapter async def test_queries_raw(database_url): """ Test that the basic `execute()`, `execute_many()`, `fetch_all()``, and `fetch_one()` interfaces are all supported (raw queries). """ async with Database(database_url) as database: async with database.transaction(force_rollback=True): # execute() query = "INSERT INTO notes(text, completed) VALUES (:text, :completed)" values = {"text": "example1", "completed": True} await database.execute(query, values) # execute_many() query = "INSERT INTO notes(text, completed) VALUES (:text, :completed)" values = [ {"text": "example2", "completed": False}, {"text": "example3", "completed": True}, ] await database.execute_many(query, values) # fetch_all() query = "SELECT * FROM notes WHERE completed = :completed" results = await database.fetch_all(query=query, values={"completed": True}) assert len(results) == 2 assert results[0]["text"] == "example1" assert results[0]["completed"] == True assert results[1]["text"] == "example3" assert results[1]["completed"] == True # fetch_one() query = "SELECT * FROM notes WHERE completed = :completed" result = await database.fetch_one(query=query, values={"completed": False}) assert result["text"] == "example2" assert result["completed"] == False # fetch_val() query = "SELECT completed FROM notes WHERE text = :text" result = await database.fetch_val(query=query, values={"text": "example1"}) assert result == True query = "SELECT * FROM notes WHERE text = :text" result = await database.fetch_val( query=query, values={"text": "example1"}, column="completed" ) assert result == True # iterate() query = "SELECT * FROM notes" iterate_results = [] async for result in database.iterate(query=query): iterate_results.append(result) assert len(iterate_results) == 3 assert iterate_results[0]["text"] == "example1" assert iterate_results[0]["completed"] == True assert iterate_results[1]["text"] == "example2" assert iterate_results[1]["completed"] == False assert iterate_results[2]["text"] == "example3" assert iterate_results[2]["completed"] == True @pytest.mark.parametrize("database_url", DATABASE_URLS) @async_adapter async def test_ddl_queries(database_url): """ Test that the built-in DDL elements such as `DropTable()`, `CreateTable()` are supported (using SQLAlchemy core). """ async with Database(database_url) as database: async with database.transaction(force_rollback=True): # DropTable() query = sqlalchemy.schema.DropTable(notes) await database.execute(query) # CreateTable() query = sqlalchemy.schema.CreateTable(notes) await database.execute(query) @pytest.mark.parametrize("exception", [Exception, asyncio.CancelledError]) @pytest.mark.parametrize("database_url", DATABASE_URLS) @async_adapter async def test_queries_after_error(database_url, exception): """ Test that the basic `execute()` works after a previous error. """ async with Database(database_url) as database: with patch.object( database.connection()._connection, "acquire", new=AsyncMock(side_effect=exception), ): with pytest.raises(exception): query = notes.select() await database.fetch_all(query) query = notes.select() await database.fetch_all(query) @pytest.mark.parametrize("database_url", DATABASE_URLS) @async_adapter async def test_results_support_mapping_interface(database_url): """ Casting results to a dict should work, since the interface defines them as supporting the mapping interface. """ async with Database(database_url) as database: async with database.transaction(force_rollback=True): # execute() query = notes.insert() values = {"text": "example1", "completed": True} await database.execute(query, values) # fetch_all() query = notes.select() results = await database.fetch_all(query=query) results_as_dicts = [dict(item) for item in results] assert len(results[0]) == 3 assert len(results_as_dicts[0]) == 3 assert isinstance(results_as_dicts[0]["id"], int) assert results_as_dicts[0]["text"] == "example1" assert results_as_dicts[0]["completed"] == True @pytest.mark.parametrize("database_url", DATABASE_URLS) @async_adapter async def test_results_support_column_reference(database_url): """ Casting results to a dict should work, since the interface defines them as supporting the mapping interface. """ async with Database(database_url) as database: async with database.transaction(force_rollback=True): now = datetime.datetime.now().replace(microsecond=0) today = datetime.date.today() # execute() query = articles.insert() values = {"title": "Hello, world Article", "published": now} await database.execute(query, values) query = custom_date.insert() values = {"title": "Hello, world Custom", "published": today} await database.execute(query, values) # fetch_all() query = sqlalchemy.select(*[articles, custom_date]) results = await database.fetch_all(query=query) assert len(results) == 1 assert results[0][articles.c.title] == "Hello, world Article" assert results[0][articles.c.published] == now assert results[0][custom_date.c.title] == "Hello, world Custom" assert results[0][custom_date.c.published] == today @pytest.mark.parametrize("database_url", DATABASE_URLS) @async_adapter async def test_result_values_allow_duplicate_names(database_url): """ The values of a result should respect when two columns are selected with the same name. """ async with Database(database_url) as database: async with database.transaction(force_rollback=True): query = "SELECT 1 AS id, 2 AS id" row = await database.fetch_one(query=query) assert list(row._mapping.keys()) == ["id", "id"] assert list(row._mapping.values()) == [1, 2] @pytest.mark.parametrize("database_url", DATABASE_URLS) @async_adapter async def test_fetch_one_returning_no_results(database_url): """ fetch_one should return `None` when no results match. """ async with Database(database_url) as database: async with database.transaction(force_rollback=True): # fetch_all() query = notes.select() result = await database.fetch_one(query=query) assert result is None @pytest.mark.parametrize("database_url", DATABASE_URLS) @async_adapter async def test_execute_return_val(database_url): """ Test using return value from `execute()` to get an inserted primary key. """ async with Database(database_url) as database: async with database.transaction(force_rollback=True): query = notes.insert() values = {"text": "example1", "completed": True} pk = await database.execute(query, values) assert isinstance(pk, int) # Apparently for `aiopg` it's OID that will always 0 in this case # As it's only one action within this cursor life cycle # It's recommended to use the `RETURNING` clause # For obtaining the record id if database.url.scheme == "postgresql+aiopg": assert pk == 0 else: query = notes.select().where(notes.c.id == pk) result = await database.fetch_one(query) assert result["text"] == "example1" assert result["completed"] == True @pytest.mark.parametrize("database_url", DATABASE_URLS) @async_adapter async def test_rollback_isolation(database_url): """ Ensure that `database.transaction(force_rollback=True)` provides strict isolation. """ async with Database(database_url) as database: # Perform some INSERT operations on the database. async with database.transaction(force_rollback=True): query = notes.insert().values(text="example1", completed=True) await database.execute(query) # Ensure INSERT operations have been rolled back. query = notes.select() results = await database.fetch_all(query=query) assert len(results) == 0 @pytest.mark.parametrize("database_url", DATABASE_URLS) @async_adapter async def test_rollback_isolation_with_contextmanager(database_url): """ Ensure that `database.force_rollback()` provides strict isolation. """ database = Database(database_url) with database.force_rollback(): async with database: # Perform some INSERT operations on the database. query = notes.insert().values(text="example1", completed=True) await database.execute(query) async with database: # Ensure INSERT operations have been rolled back. query = notes.select() results = await database.fetch_all(query=query) assert len(results) == 0 @pytest.mark.parametrize("database_url", DATABASE_URLS) @async_adapter async def test_transaction_commit(database_url): """ Ensure that transaction commit is supported. """ async with Database(database_url) as database: async with database.transaction(force_rollback=True): async with database.transaction(): query = notes.insert().values(text="example1", completed=True) await database.execute(query) query = notes.select() results = await database.fetch_all(query=query) assert len(results) == 1 @pytest.mark.parametrize("database_url", DATABASE_URLS) @async_adapter async def test_transaction_context_child_task_inheritance(database_url): """ Ensure that transactions are inherited by child tasks. """ async with Database(database_url) as database: async def check_transaction(transaction, active_transaction): # Should have inherited the same transaction backend from the parent task assert transaction._transaction is active_transaction async with database.transaction() as transaction: await asyncio.create_task( check_transaction(transaction, transaction._transaction) ) @pytest.mark.parametrize("database_url", DATABASE_URLS) @async_adapter async def test_transaction_context_child_task_inheritance_example(database_url): """ Ensure that child tasks may influence inherited transactions. """ # This is an practical example of the above test. async with Database(database_url) as database: async with database.transaction(): # Create a note await database.execute( notes.insert().values(id=1, text="setup", completed=True) ) # Change the note from the same task await database.execute( notes.update().where(notes.c.id == 1).values(text="prior") ) # Confirm the change result = await database.fetch_one(notes.select().where(notes.c.id == 1)) assert result.text == "prior" async def run_update_from_child_task(connection): # Change the note from a child task await connection.execute( notes.update().where(notes.c.id == 1).values(text="test") ) await asyncio.create_task(run_update_from_child_task(database.connection())) # Confirm the child's change result = await database.fetch_one(notes.select().where(notes.c.id == 1)) assert result.text == "test" @pytest.mark.parametrize("database_url", DATABASE_URLS) @async_adapter async def test_transaction_context_sibling_task_isolation(database_url): """ Ensure that transactions are isolated between sibling tasks. """ start = asyncio.Event() end = asyncio.Event() async with Database(database_url) as database: async def check_transaction(transaction): await start.wait() # Parent task is now in a transaction, we should not # see its transaction backend since this task was # _started_ in a context where no transaction was active. assert transaction._transaction is None end.set() transaction = database.transaction() assert transaction._transaction is None task = asyncio.create_task(check_transaction(transaction)) async with transaction: start.set() assert transaction._transaction is not None await end.wait() # Cleanup for "Task not awaited" warning await task @pytest.mark.parametrize("database_url", DATABASE_URLS) @async_adapter async def test_transaction_context_sibling_task_isolation_example(database_url): """ Ensure that transactions are running in sibling tasks are isolated from eachother. """ # This is an practical example of the above test. setup = asyncio.Event() done = asyncio.Event() async def tx1(connection): async with connection.transaction(): await db.execute( notes.insert(), values={"id": 1, "text": "tx1", "completed": False} ) setup.set() await done.wait() async def tx2(connection): async with connection.transaction(): await setup.wait() result = await db.fetch_all(notes.select()) assert result == [], result done.set() async with Database(database_url) as db: await asyncio.gather(tx1(db), tx2(db)) @pytest.mark.parametrize("database_url", DATABASE_URLS) @async_adapter async def test_connection_cleanup_contextmanager(database_url): """ Ensure that task connections are not persisted unecessarily. """ ready = asyncio.Event() done = asyncio.Event() async def check_child_connection(database: Database): async with database.connection(): ready.set() await done.wait() async with Database(database_url) as database: # Should have a connection in this task # .connect is lazy, it doesn't create a Connection, but .connection does connection = database.connection() assert isinstance(database._connection_map, MutableMapping) assert database._connection_map.get(asyncio.current_task()) is connection # Create a child task and see if it registers a connection task = asyncio.create_task(check_child_connection(database)) await ready.wait() assert database._connection_map.get(task) is not None assert database._connection_map.get(task) is not connection # Let the child task finish, and see if it cleaned up done.set() await task # This is normal exit logic cleanup, the WeakKeyDictionary # shouldn't have cleaned up yet since the task is still referenced assert task not in database._connection_map # Context manager closes, all open connections are removed assert isinstance(database._connection_map, MutableMapping) assert len(database._connection_map) == 0 @pytest.mark.parametrize("database_url", DATABASE_URLS) @async_adapter async def test_connection_cleanup_garbagecollector(database_url): """ Ensure that connections for tasks are not persisted unecessarily, even if exit handlers are not called. """ database = Database(database_url) await database.connect() created = asyncio.Event() async def check_child_connection(database: Database): # neither .disconnect nor .__aexit__ are called before deleting this task database.connection() created.set() task = asyncio.create_task(check_child_connection(database)) await created.wait() assert task in database._connection_map await task del task gc.collect() # Should not have a connection for the task anymore assert len(database._connection_map) == 0 @pytest.mark.parametrize("database_url", DATABASE_URLS) @async_adapter async def test_transaction_context_cleanup_contextmanager(database_url): """ Ensure that contextvar transactions are not persisted unecessarily. """ from databases.core import _ACTIVE_TRANSACTIONS assert _ACTIVE_TRANSACTIONS.get() is None async with Database(database_url) as database: async with database.transaction() as transaction: open_transactions = _ACTIVE_TRANSACTIONS.get() assert isinstance(open_transactions, MutableMapping) assert open_transactions.get(transaction) is transaction._transaction # Context manager closes, open_transactions is cleaned up open_transactions = _ACTIVE_TRANSACTIONS.get() assert isinstance(open_transactions, MutableMapping) assert open_transactions.get(transaction, None) is None @pytest.mark.parametrize("database_url", DATABASE_URLS) @async_adapter async def test_transaction_context_cleanup_garbagecollector(database_url): """ Ensure that contextvar transactions are not persisted unecessarily, even if exit handlers are not called. This test should be an XFAIL, but cannot be due to the way that is hangs during teardown. """ from databases.core import _ACTIVE_TRANSACTIONS assert _ACTIVE_TRANSACTIONS.get() is None async with Database(database_url) as database: transaction = database.transaction() await transaction.start() # Should be tracking the transaction open_transactions = _ACTIVE_TRANSACTIONS.get() assert isinstance(open_transactions, MutableMapping) assert open_transactions.get(transaction) is transaction._transaction # neither .commit, .rollback, nor .__aexit__ are called del transaction gc.collect() # TODO(zevisert,review): Could skip instead of using the logic below # A strong reference to the transaction is kept alive by the connection's # ._transaction_stack, so it is still be tracked at this point. assert len(open_transactions) == 1 # If that were magically cleared, the transaction would be cleaned up, # but as it stands this always causes a hang during teardown at # `Database(...).disconnect()` if the transaction is not closed. transaction = database.connection()._transaction_stack[-1] await transaction.rollback() del transaction # Now with the transaction rolled-back, it should be cleaned up. assert len(open_transactions) == 0 @pytest.mark.parametrize("database_url", DATABASE_URLS) @async_adapter async def test_transaction_commit_serializable(database_url): """ Ensure that serializable transaction commit via extra parameters is supported. """ database_url = DatabaseURL(database_url) if database_url.scheme not in ["postgresql", "postgresql+asyncpg"]: pytest.skip("Test (currently) only supports asyncpg") if database_url.scheme == "postgresql+asyncpg": database_url = database_url.replace(driver=None) def insert_independently(): engine = sqlalchemy.create_engine(str(database_url)) conn = engine.connect() query = notes.insert().values(text="example1", completed=True) conn.execute(query) conn.close() def delete_independently(): engine = sqlalchemy.create_engine(str(database_url)) conn = engine.connect() query = notes.delete() conn.execute(query) conn.close() async with Database(database_url) as database: async with database.transaction(force_rollback=True, isolation="serializable"): query = notes.select() results = await database.fetch_all(query=query) assert len(results) == 0 insert_independently() query = notes.select() results = await database.fetch_all(query=query) assert len(results) == 0 delete_independently() @pytest.mark.parametrize("database_url", DATABASE_URLS) @async_adapter async def test_transaction_rollback(database_url): """ Ensure that transaction rollback is supported. """ async with Database(database_url) as database: async with database.transaction(force_rollback=True): try: async with database.transaction(): query = notes.insert().values(text="example1", completed=True) await database.execute(query) raise RuntimeError() except RuntimeError: pass query = notes.select() results = await database.fetch_all(query=query) assert len(results) == 0 @pytest.mark.parametrize("database_url", DATABASE_URLS) @async_adapter async def test_transaction_commit_low_level(database_url): """ Ensure that an explicit `await transaction.commit()` is supported. """ async with Database(database_url) as database: async with database.transaction(force_rollback=True): transaction = await database.transaction() try: query = notes.insert().values(text="example1", completed=True) await database.execute(query) except: # pragma: no cover await transaction.rollback() else: await transaction.commit() query = notes.select() results = await database.fetch_all(query=query) assert len(results) == 1 @pytest.mark.parametrize("database_url", DATABASE_URLS) @async_adapter async def test_transaction_rollback_low_level(database_url): """ Ensure that an explicit `await transaction.rollback()` is supported. """ async with Database(database_url) as database: async with database.transaction(force_rollback=True): transaction = await database.transaction() try: query = notes.insert().values(text="example1", completed=True) await database.execute(query) raise RuntimeError() except: await transaction.rollback() else: # pragma: no cover await transaction.commit() query = notes.select() results = await database.fetch_all(query=query) assert len(results) == 0 @pytest.mark.parametrize("database_url", DATABASE_URLS) @async_adapter async def test_transaction_decorator(database_url): """ Ensure that @database.transaction() is supported. """ database = Database(database_url, force_rollback=True) @database.transaction() async def insert_data(raise_exception): query = notes.insert().values(text="example", completed=True) await database.execute(query) if raise_exception: raise RuntimeError() async with database: with pytest.raises(RuntimeError): await insert_data(raise_exception=True) results = await database.fetch_all(query=notes.select()) assert len(results) == 0 await insert_data(raise_exception=False) results = await database.fetch_all(query=notes.select()) assert len(results) == 1 @pytest.mark.parametrize("database_url", DATABASE_URLS) @async_adapter async def test_transaction_decorator_concurrent(database_url): """ Ensure that @database.transaction() can be called concurrently. """ database = Database(database_url) @database.transaction() async def insert_data(): await database.execute( query=notes.insert().values(text="example", completed=True) ) async with database: await asyncio.gather( insert_data(), insert_data(), insert_data(), insert_data(), insert_data(), insert_data(), ) results = await database.fetch_all(query=notes.select()) assert len(results) == 6 @pytest.mark.parametrize("database_url", DATABASE_URLS) @async_adapter async def test_datetime_field(database_url): """ Test DataTime columns, to ensure records are coerced to/from proper Python types. """ async with Database(database_url) as database: async with database.transaction(force_rollback=True): now = datetime.datetime.now().replace(microsecond=0) # execute() query = articles.insert() values = {"title": "Hello, world", "published": now} await database.execute(query, values) # fetch_all() query = articles.select() results = await database.fetch_all(query=query) assert len(results) == 1 assert results[0]["title"] == "Hello, world" assert results[0]["published"] == now @pytest.mark.parametrize("database_url", DATABASE_URLS) @async_adapter async def test_date_field(database_url): """ Test Date columns, to ensure records are coerced to/from proper Python types. """ async with Database(database_url) as database: async with database.transaction(force_rollback=True): now = datetime.date.today() # execute() query = events.insert() values = {"date": now} await database.execute(query, values) # fetch_all() query = events.select() results = await database.fetch_all(query=query) assert len(results) == 1 assert results[0]["date"] == now @pytest.mark.parametrize("database_url", DATABASE_URLS) @async_adapter async def test_time_field(database_url): """ Test Time columns, to ensure records are coerced to/from proper Python types. """ async with Database(database_url) as database: async with database.transaction(force_rollback=True): now = datetime.datetime.now().time().replace(microsecond=0) # execute() query = daily_schedule.insert() values = {"time": now} await database.execute(query, values) # fetch_all() query = daily_schedule.select() results = await database.fetch_all(query=query) assert len(results) == 1 assert results[0]["time"] == now @pytest.mark.parametrize("database_url", DATABASE_URLS) @async_adapter async def test_decimal_field(database_url): """ Test Decimal (NUMERIC) columns, to ensure records are coerced to/from proper Python types. """ async with Database(database_url) as database: async with database.transaction(force_rollback=True): price = decimal.Decimal("0.700000000000001") # execute() query = prices.insert() values = {"price": price} await database.execute(query, values) # fetch_all() query = prices.select() results = await database.fetch_all(query=query) assert len(results) == 1 if database_url.startswith("sqlite"): # aiosqlite does not support native decimals --> a roud-off error is expected assert results[0]["price"] == pytest.approx(price) else: assert results[0]["price"] == price @pytest.mark.parametrize("database_url", DATABASE_URLS) @async_adapter async def test_enum_field(database_url): """ Test enum columns, to ensure correct cross-database support. """ async with Database(database_url) as database: async with database.transaction(force_rollback=True): # execute() size = TshirtSize.SMALL color = TshirtColor.GREEN values = {"size": size, "color": color} query = tshirt_size.insert() await database.execute(query, values) # fetch_all() query = tshirt_size.select() results = await database.fetch_all(query=query) assert len(results) == 1 assert results[0]["size"] == size assert results[0]["color"] == color @pytest.mark.parametrize("database_url", DATABASE_URLS) @async_adapter async def test_json_dict_field(database_url): """ Test JSON columns, to ensure correct cross-database support. """ async with Database(database_url) as database: async with database.transaction(force_rollback=True): # execute() data = {"text": "hello", "boolean": True, "int": 1} values = {"data": data} query = session.insert() await database.execute(query, values) # fetch_all() query = session.select() results = await database.fetch_all(query=query) assert len(results) == 1 assert results[0]["data"] == {"text": "hello", "boolean": True, "int": 1} @pytest.mark.parametrize("database_url", DATABASE_URLS) @async_adapter async def test_json_list_field(database_url): """ Test JSON columns, to ensure correct cross-database support. """ async with Database(database_url) as database: async with database.transaction(force_rollback=True): # execute() data = ["lemon", "raspberry", "lime", "pumice"] values = {"data": data} query = session.insert() await database.execute(query, values) # fetch_all() query = session.select() results = await database.fetch_all(query=query) assert len(results) == 1 assert results[0]["data"] == ["lemon", "raspberry", "lime", "pumice"] @pytest.mark.parametrize("database_url", DATABASE_URLS) @async_adapter async def test_custom_field(database_url): """ Test custom column types. """ async with Database(database_url) as database: async with database.transaction(force_rollback=True): today = datetime.date.today() # execute() query = custom_date.insert() values = {"title": "Hello, world", "published": today} await database.execute(query, values) # fetch_all() query = custom_date.select() results = await database.fetch_all(query=query) assert len(results) == 1 assert results[0]["title"] == "Hello, world" assert results[0]["published"] == today @pytest.mark.parametrize("database_url", DATABASE_URLS) @async_adapter async def test_connections_isolation(database_url): """ Ensure that changes are visible between different connections. To check this we have to not create a transaction, so that each query ends up on a different connection from the pool. """ async with Database(database_url) as database: try: query = notes.insert().values(text="example1", completed=True) await database.execute(query) query = notes.select() results = await database.fetch_all(query=query) assert len(results) == 1 finally: query = notes.delete() await database.execute(query) @pytest.mark.parametrize("database_url", DATABASE_URLS) @async_adapter async def test_commit_on_root_transaction(database_url): """ Because our tests are generally wrapped in rollback-islation, they don't have coverage for commiting the root transaction. Deal with this here, and delete the records rather than rolling back. """ async with Database(database_url) as database: try: async with database.transaction(): query = notes.insert().values(text="example1", completed=True) await database.execute(query) query = notes.select() results = await database.fetch_all(query=query) assert len(results) == 1 finally: query = notes.delete() await database.execute(query) @pytest.mark.parametrize("database_url", DATABASE_URLS) @async_adapter async def test_connect_and_disconnect(database_url): """ Test explicit connect() and disconnect(). """ database = Database(database_url) assert not database.is_connected await database.connect() assert database.is_connected await database.disconnect() assert not database.is_connected # connect and disconnect idempotence await database.connect() await database.connect() assert database.is_connected await database.disconnect() await database.disconnect() assert not database.is_connected @pytest.mark.parametrize("database_url", DATABASE_URLS) @async_adapter async def test_connection_context_same_task(database_url): async with Database(database_url) as database: async with database.connection() as connection_1: async with database.connection() as connection_2: assert connection_1 is connection_2 @pytest.mark.parametrize("database_url", DATABASE_URLS) @async_adapter async def test_connection_context_multiple_sibling_tasks(database_url): async with Database(database_url) as database: connection_1 = None connection_2 = None test_complete = asyncio.Event() async def get_connection_1(): nonlocal connection_1 async with database.connection() as connection: connection_1 = connection await test_complete.wait() async def get_connection_2(): nonlocal connection_2 async with database.connection() as connection: connection_2 = connection await test_complete.wait() task_1 = asyncio.create_task(get_connection_1()) task_2 = asyncio.create_task(get_connection_2()) while connection_1 is None or connection_2 is None: await asyncio.sleep(0.000001) assert connection_1 is not connection_2 test_complete.set() await task_1 await task_2 @pytest.mark.parametrize("database_url", DATABASE_URLS) @async_adapter async def test_connection_context_multiple_tasks(database_url): async with Database(database_url) as database: parent_connection = database.connection() connection_1 = None connection_2 = None task_1_ready = asyncio.Event() task_2_ready = asyncio.Event() test_complete = asyncio.Event() async def get_connection_1(): nonlocal connection_1 async with database.connection() as connection: connection_1 = connection task_1_ready.set() await test_complete.wait() async def get_connection_2(): nonlocal connection_2 async with database.connection() as connection: connection_2 = connection task_2_ready.set() await test_complete.wait() task_1 = asyncio.create_task(get_connection_1()) task_2 = asyncio.create_task(get_connection_2()) await task_1_ready.wait() await task_2_ready.wait() assert connection_1 is not parent_connection assert connection_2 is not parent_connection assert connection_1 is not connection_2 test_complete.set() await task_1 await task_2 @pytest.mark.parametrize( "database_url1,database_url2", ( pytest.param(db1, db2, id=f"{db1} | {db2}") for (db1, db2) in itertools.combinations(DATABASE_URLS, 2) ), ) @async_adapter async def test_connection_context_multiple_databases(database_url1, database_url2): async with Database(database_url1) as database1: async with Database(database_url2) as database2: assert database1.connection() is not database2.connection() @pytest.mark.parametrize("database_url", DATABASE_URLS) @async_adapter async def test_connection_context_with_raw_connection(database_url): """ Test connection contexts with respect to the raw connection. """ async with Database(database_url) as database: async with database.connection() as connection_1: async with database.connection() as connection_2: assert connection_1 is connection_2 assert connection_1.raw_connection is connection_2.raw_connection @pytest.mark.parametrize("database_url", DATABASE_URLS) @async_adapter async def test_queries_with_expose_backend_connection(database_url): """ Replication of `execute()`, `execute_many()`, `fetch_all()``, and `fetch_one()` using the raw driver interface. """ async with Database(database_url) as database: async with database.connection() as connection: async with connection.transaction(force_rollback=True): # Get the raw connection raw_connection = connection.raw_connection # Insert query if database.url.scheme in [ "mysql", "mysql+asyncmy", "mysql+aiomysql", "postgresql+aiopg", ]: insert_query = "INSERT INTO notes (text, completed) VALUES (%s, %s)" else: insert_query = "INSERT INTO notes (text, completed) VALUES ($1, $2)" # execute() values = ("example1", True) if database.url.scheme in [ "mysql", "mysql+aiomysql", "postgresql+aiopg", ]: cursor = await raw_connection.cursor() await cursor.execute(insert_query, values) elif database.url.scheme == "mysql+asyncmy": async with raw_connection.cursor() as cursor: await cursor.execute(insert_query, values) elif database.url.scheme in ["postgresql", "postgresql+asyncpg"]: await raw_connection.execute(insert_query, *values) elif database.url.scheme in ["sqlite", "sqlite+aiosqlite"]: await raw_connection.execute(insert_query, values) # execute_many() values = [("example2", False), ("example3", True)] if database.url.scheme in ["mysql", "mysql+aiomysql"]: cursor = await raw_connection.cursor() await cursor.executemany(insert_query, values) elif database.url.scheme == "mysql+asyncmy": async with raw_connection.cursor() as cursor: await cursor.executemany(insert_query, values) elif database.url.scheme == "postgresql+aiopg": cursor = await raw_connection.cursor() # No async support for `executemany` for value in values: await cursor.execute(insert_query, value) else: await raw_connection.executemany(insert_query, values) # Select query select_query = "SELECT notes.id, notes.text, notes.completed FROM notes" # fetch_all() if database.url.scheme in [ "mysql", "mysql+aiomysql", "postgresql+aiopg", ]: cursor = await raw_connection.cursor() await cursor.execute(select_query) results = await cursor.fetchall() elif database.url.scheme == "mysql+asyncmy": async with raw_connection.cursor() as cursor: await cursor.execute(select_query) results = await cursor.fetchall() elif database.url.scheme in ["postgresql", "postgresql+asyncpg"]: results = await raw_connection.fetch(select_query) elif database.url.scheme in ["sqlite", "sqlite+aiosqlite"]: results = await raw_connection.execute_fetchall(select_query) assert len(results) == 3 # Raw output for the raw request assert results[0][1] == "example1" assert results[0][2] == True assert results[1][1] == "example2" assert results[1][2] == False assert results[2][1] == "example3" assert results[2][2] == True # fetch_one() if database.url.scheme in ["postgresql", "postgresql+asyncpg"]: result = await raw_connection.fetchrow(select_query) elif database.url.scheme == "mysql+asyncmy": async with raw_connection.cursor() as cursor: await cursor.execute(select_query) result = await cursor.fetchone() else: cursor = await raw_connection.cursor() await cursor.execute(select_query) result = await cursor.fetchone() # Raw output for the raw request assert result[1] == "example1" assert result[2] == True @pytest.mark.parametrize("database_url", DATABASE_URLS) @async_adapter async def test_database_url_interface(database_url): """ Test that Database instances expose a `.url` attribute. """ async with Database(database_url) as database: assert isinstance(database.url, DatabaseURL) assert database.url == database_url @pytest.mark.parametrize("database_url", DATABASE_URLS) @async_adapter async def test_concurrent_access_on_single_connection(database_url): async with Database(database_url, force_rollback=True) as database: async def db_lookup(): await database.fetch_one("SELECT 1 AS value") await asyncio.gather( db_lookup(), db_lookup(), ) @pytest.mark.parametrize("database_url", DATABASE_URLS) @async_adapter async def test_concurrent_transactions_on_single_connection(database_url: str): async with Database(database_url) as database: @database.transaction() async def db_lookup(): await database.fetch_one(query="SELECT 1 AS value") await asyncio.gather( db_lookup(), db_lookup(), ) @pytest.mark.parametrize("database_url", DATABASE_URLS) @async_adapter async def test_concurrent_tasks_on_single_connection(database_url: str): async with Database(database_url) as database: async def db_lookup(): await database.fetch_one(query="SELECT 1 AS value") await asyncio.gather( asyncio.create_task(db_lookup()), asyncio.create_task(db_lookup()), ) @pytest.mark.parametrize("database_url", DATABASE_URLS) @async_adapter async def test_concurrent_task_transactions_on_single_connection(database_url: str): async with Database(database_url) as database: @database.transaction() async def db_lookup(): await database.fetch_one(query="SELECT 1 AS value") await asyncio.gather( asyncio.create_task(db_lookup()), asyncio.create_task(db_lookup()), ) @pytest.mark.parametrize("database_url", DATABASE_URLS) def test_global_connection_is_initialized_lazily(database_url): """ Ensure that global connection is initialized at latest possible time so it's _query_lock will belong to same event loop that async_adapter has initialized. See https://github.com/encode/databases/issues/157 for more context. """ database_url = DatabaseURL(database_url) if database_url.dialect != "postgresql": pytest.skip("Test requires `pg_sleep()`") database = Database(database_url, force_rollback=True) @async_adapter async def run_database_queries(): async with database: async def db_lookup(): await database.fetch_one("SELECT pg_sleep(1)") await asyncio.gather(db_lookup(), db_lookup()) run_database_queries() @pytest.mark.parametrize("database_url", DATABASE_URLS) @async_adapter async def test_iterate_outside_transaction_with_values(database_url): """ Ensure `iterate()` works even without a transaction on all drivers. The asyncpg driver relies on server-side cursors without hold for iteration, which requires a transaction to be created. This is mentionned in both their documentation and their test suite. """ database_url = DatabaseURL(database_url) if database_url.dialect == "mysql": pytest.skip("MySQL does not support `FROM (VALUES ...)` (F641)") async with Database(database_url) as database: query = "SELECT * FROM (VALUES (1), (2), (3), (4), (5)) as t" iterate_results = [] async for result in database.iterate(query=query): iterate_results.append(result) assert len(iterate_results) == 5 @pytest.mark.parametrize("database_url", DATABASE_URLS) @async_adapter async def test_iterate_outside_transaction_with_temp_table(database_url): """ Same as test_iterate_outside_transaction_with_values but uses a temporary table instead of a list of values. """ database_url = DatabaseURL(database_url) if database_url.dialect == "sqlite": pytest.skip("SQLite interface does not work with temporary tables.") async with Database(database_url) as database: query = "CREATE TEMPORARY TABLE no_transac(num INTEGER)" await database.execute(query) query = "INSERT INTO no_transac(num) VALUES (1), (2), (3), (4), (5)" await database.execute(query) query = "SELECT * FROM no_transac" iterate_results = [] async for result in database.iterate(query=query): iterate_results.append(result) assert len(iterate_results) == 5 @pytest.mark.parametrize("database_url", DATABASE_URLS) @pytest.mark.parametrize("select_query", [notes.select(), "SELECT * FROM notes"]) @async_adapter async def test_column_names(database_url, select_query): """ Test that column names are exposed correctly through `._mapping.keys()` on each row. """ async with Database(database_url) as database: async with database.transaction(force_rollback=True): # insert values query = notes.insert() values = {"text": "example1", "completed": True} await database.execute(query, values) # fetch results results = await database.fetch_all(query=select_query) assert len(results) == 1 assert sorted(results[0]._mapping.keys()) == ["completed", "id", "text"] assert results[0]["text"] == "example1" assert results[0]["completed"] == True @pytest.mark.parametrize("database_url", DATABASE_URLS) @async_adapter async def test_postcompile_queries(database_url): """ Since SQLAlchemy 1.4, IN operators needs to do render_postcompile """ async with Database(database_url) as database: query = notes.insert() values = {"text": "example1", "completed": True} await database.execute(query, values) query = notes.select().where(notes.c.id.in_([2, 3])) results = await database.fetch_all(query=query) assert len(results) == 0 @pytest.mark.parametrize("database_url", DATABASE_URLS) @async_adapter async def test_result_named_access(database_url): async with Database(database_url) as database: query = notes.insert() values = {"text": "example1", "completed": True} await database.execute(query, values) query = notes.select().where(notes.c.text == "example1") result = await database.fetch_one(query=query) assert result.text == "example1" assert result.completed is True @pytest.mark.parametrize("database_url", DATABASE_URLS) @async_adapter async def test_mapping_property_interface(database_url): """ Test that all connections implement interface with `_mapping` property """ async with Database(database_url) as database: query = notes.select() single_result = await database.fetch_one(query=query) assert single_result._mapping["text"] == "example1" assert single_result._mapping["completed"] is True list_result = await database.fetch_all(query=query) assert list_result[0]._mapping["text"] == "example1" assert list_result[0]._mapping["completed"] is True @async_adapter async def test_should_not_maintain_ref_when_no_cache_param(): async with Database( "sqlite:///file::memory:", uri=True, ) as database: query = sqlalchemy.schema.CreateTable(notes) await database.execute(query) query = notes.insert() values = {"text": "example1", "completed": True} with pytest.raises(sqlite3.OperationalError): await database.execute(query, values) @async_adapter async def test_should_maintain_ref_when_cache_param(): async with Database( "sqlite:///file::memory:?cache=shared", uri=True, ) as database: query = sqlalchemy.schema.CreateTable(notes) await database.execute(query) query = notes.insert() values = {"text": "example1", "completed": True} await database.execute(query, values) query = notes.select().where(notes.c.text == "example1") result = await database.fetch_one(query=query) assert result.text == "example1" assert result.completed is True @async_adapter async def test_should_remove_ref_on_disconnect(): async with Database( "sqlite:///file::memory:?cache=shared", uri=True, ) as database: query = sqlalchemy.schema.CreateTable(notes) await database.execute(query) query = notes.insert() values = {"text": "example1", "completed": True} await database.execute(query, values) # Run garbage collection to reset the database if we dropped the reference gc.collect() async with Database( "sqlite:///file::memory:?cache=shared", uri=True, ) as database: query = notes.select() with pytest.raises(sqlite3.OperationalError): await database.fetch_all(query=query) @pytest.mark.parametrize("database_url", DATABASE_URLS) @async_adapter async def test_mapping_property_interface(database_url): """ Test that all connections implement interface with `_mapping` property """ async with Database(database_url) as database: query = notes.insert() values = {"text": "example1", "completed": True} await database.execute(query, values) query = notes.select() single_result = await database.fetch_one(query=query) assert single_result._mapping["text"] == "example1" assert single_result._mapping["completed"] is True list_result = await database.fetch_all(query=query) assert list_result[0]._mapping["text"] == "example1" assert list_result[0]._mapping["completed"] is True