mirror of
https://github.com/langbot-app/LangBot.git
synced 2026-10-31 21:06:18 +08:00
360 lines
15 KiB
Python
360 lines
15 KiB
Python
"""Writer admission on real pooled connections, including failure/cleanup paths."""
|
|
|
|
import asyncio
|
|
import gc
|
|
import sqlite3
|
|
from types import SimpleNamespace
|
|
|
|
import pytest
|
|
import sqlalchemy as sa
|
|
from sqlalchemy.ext.asyncio import AsyncSession, create_async_engine
|
|
from sqlalchemy.orm import DeclarativeBase, Mapped, mapped_column
|
|
|
|
from langbot.pkg.persistence.sqlite_write_gate import install
|
|
|
|
|
|
class Base(DeclarativeBase):
|
|
pass
|
|
|
|
|
|
class Row(Base):
|
|
__tablename__ = 'gate_rows'
|
|
id: Mapped[int] = mapped_column(primary_key=True)
|
|
|
|
|
|
@pytest.fixture
|
|
async def database(tmp_path):
|
|
engine = create_async_engine(f'sqlite+aiosqlite:///{tmp_path / "gate.db"}')
|
|
gate = install(engine)
|
|
async with engine.begin() as connection:
|
|
await connection.exec_driver_sql('PRAGMA journal_mode=WAL')
|
|
await connection.run_sync(Base.metadata.create_all)
|
|
yield engine, gate
|
|
await engine.dispose()
|
|
assert not gate.lock.locked()
|
|
assert gate.waiting == 0
|
|
|
|
|
|
async def queued(gate):
|
|
async with asyncio.timeout(2):
|
|
while gate.waiting == 0:
|
|
await asyncio.sleep(0)
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_session_flush_and_direct_writer_share_gate_without_blocking_reads(database):
|
|
engine, gate = database
|
|
async with AsyncSession(engine) as session, engine.connect() as direct:
|
|
session.add(Row(id=1))
|
|
await session.flush()
|
|
writer = asyncio.create_task(direct.exec_driver_sql('INSERT INTO gate_rows VALUES (2)'))
|
|
await queued(gate)
|
|
async with engine.connect() as reader:
|
|
assert await asyncio.wait_for(reader.scalar(sa.select(sa.func.count()).select_from(Row)), 1) == 0
|
|
await session.commit()
|
|
await asyncio.wait_for(writer, 1)
|
|
assert gate.lock.locked()
|
|
await direct.commit()
|
|
async with engine.connect() as reader:
|
|
assert await reader.scalar(sa.select(sa.func.count()).select_from(Row)) == 2
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_cancel_waiter_does_not_unlock_owner(database):
|
|
engine, gate = database
|
|
async with engine.connect() as owner, engine.connect() as waiting:
|
|
await owner.exec_driver_sql('INSERT INTO gate_rows VALUES (1)')
|
|
task = asyncio.create_task(waiting.exec_driver_sql('INSERT INTO gate_rows VALUES (2)'))
|
|
await queued(gate)
|
|
task.cancel()
|
|
with pytest.raises(asyncio.CancelledError):
|
|
await task
|
|
assert gate.lock.locked() and gate.waiting == 0
|
|
await owner.rollback()
|
|
await waiting.exec_driver_sql('INSERT INTO gate_rows VALUES (3)')
|
|
await waiting.commit()
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_savepoint_rollback_keeps_gate_until_outer_transaction_ends(database):
|
|
engine, gate = database
|
|
async with engine.connect() as connection:
|
|
await connection.exec_driver_sql('BEGIN')
|
|
nested = await connection.begin_nested()
|
|
await connection.execute(sa.insert(Row).values(id=1))
|
|
await nested.rollback()
|
|
assert gate.lock.locked()
|
|
await connection.rollback()
|
|
assert not gate.lock.locked()
|
|
nested = await connection.begin_nested()
|
|
await connection.execute(sa.insert(Row).values(id=2))
|
|
await nested.commit()
|
|
assert not gate.lock.locked() # SQLite released the outermost SAVEPOINT.
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_failed_statement_and_commit_release_on_rollback(database):
|
|
engine, gate = database
|
|
async with engine.connect() as connection:
|
|
await connection.exec_driver_sql('PRAGMA foreign_keys=ON')
|
|
await connection.exec_driver_sql('CREATE TABLE parent (id INTEGER PRIMARY KEY)')
|
|
await connection.exec_driver_sql(
|
|
'CREATE TABLE child (id INTEGER REFERENCES parent(id) DEFERRABLE INITIALLY DEFERRED)'
|
|
)
|
|
await connection.commit()
|
|
await connection.exec_driver_sql('INSERT INTO child VALUES (123)')
|
|
with pytest.raises(sa.exc.IntegrityError):
|
|
await connection.commit()
|
|
assert gate.lock.locked()
|
|
# A failed commit leaves SQLite's physical transaction open. Closing
|
|
# must roll it back before another connection is admitted.
|
|
assert not gate.lock.locked()
|
|
async with engine.begin() as connection:
|
|
await connection.exec_driver_sql('INSERT INTO gate_rows VALUES (1)')
|
|
with pytest.raises(sa.exc.IntegrityError):
|
|
await connection.exec_driver_sql('INSERT INTO gate_rows VALUES (1)')
|
|
assert gate.lock.locked()
|
|
assert not gate.lock.locked()
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_dispose_reconnect_and_cte_dml(database):
|
|
engine, gate = database
|
|
for _ in range(2):
|
|
async with engine.connect() as connection:
|
|
count = gate.acquisitions
|
|
await connection.exec_driver_sql('/* INSERT ignored */ WITH x(v) AS (SELECT 1) SELECT v FROM x')
|
|
assert gate.acquisitions == count
|
|
await connection.exec_driver_sql('WITH x(v) AS (SELECT 5) INSERT OR REPLACE INTO gate_rows SELECT v FROM x')
|
|
# SQLite's legacy mode autocommits a leading WITH; no retained lock.
|
|
await connection.commit()
|
|
await connection.exec_driver_sql('PRAGMA user_version=3')
|
|
assert await connection.scalar(sa.text('PRAGMA user_version')) == 3
|
|
count = gate.acquisitions
|
|
await connection.exec_driver_sql('PRAGMA user_version(4)')
|
|
assert gate.acquisitions == count + 1
|
|
assert await connection.scalar(sa.text('PRAGMA user_version')) == 4
|
|
count = gate.acquisitions
|
|
await connection.exec_driver_sql('PRAGMA table_info(gate_rows)')
|
|
assert gate.acquisitions == count
|
|
await connection.exec_driver_sql('PRAGMA incremental_vacuum(1)')
|
|
assert not gate.lock.locked()
|
|
await engine.dispose()
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_gc_force_close_releases_after_driver_close(database):
|
|
engine, gate = database
|
|
connection = await engine.connect()
|
|
await connection.exec_driver_sql('INSERT INTO gate_rows VALUES (1)')
|
|
del connection
|
|
with pytest.warns(sa.exc.SAWarning, match='garbage collector'):
|
|
gc.collect()
|
|
async with asyncio.timeout(2):
|
|
while gate.lock.locked():
|
|
await asyncio.sleep(0.01)
|
|
async with engine.begin() as connection:
|
|
await connection.exec_driver_sql('INSERT INTO gate_rows VALUES (2)')
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_cancelled_statement_drains_before_release(database):
|
|
engine, gate = database
|
|
started = asyncio.Event()
|
|
loop = asyncio.get_running_loop()
|
|
import time
|
|
|
|
def slow():
|
|
loop.call_soon_threadsafe(started.set)
|
|
time.sleep(0.1)
|
|
return 1
|
|
|
|
async with engine.connect() as connection:
|
|
raw = (await connection.get_raw_connection()).driver_connection
|
|
await raw.create_function('slow', 0, slow)
|
|
task = asyncio.create_task(connection.exec_driver_sql('INSERT INTO gate_rows VALUES (slow())'))
|
|
await asyncio.wait_for(started.wait(), 1)
|
|
task.cancel()
|
|
with pytest.raises(asyncio.CancelledError):
|
|
await task
|
|
async with engine.begin() as connection:
|
|
await asyncio.wait_for(connection.exec_driver_sql('INSERT INTO gate_rows VALUES (2)'), 2)
|
|
|
|
|
|
def test_non_sqlite_engine_is_not_touched():
|
|
# No event API or sync engine is needed on the PostgreSQL path.
|
|
assert install(SimpleNamespace(dialect=SimpleNamespace(name='postgresql'))) is None
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_cancel_commit_retains_gate_until_connection_is_closed(database):
|
|
engine, gate = database
|
|
entered = asyncio.Event()
|
|
async with engine.connect() as connection:
|
|
await connection.exec_driver_sql('INSERT INTO gate_rows VALUES (1)')
|
|
driver = (await connection.get_raw_connection()).driver_connection
|
|
|
|
async def blocked_commit():
|
|
entered.set()
|
|
await asyncio.Event().wait()
|
|
|
|
driver.commit = blocked_commit
|
|
task = asyncio.create_task(connection.commit())
|
|
await asyncio.wait_for(entered.wait(), 1)
|
|
assert gate.lock.locked()
|
|
task.cancel()
|
|
with pytest.raises(asyncio.CancelledError):
|
|
await task
|
|
async with engine.begin() as connection:
|
|
await asyncio.wait_for(connection.exec_driver_sql('INSERT INTO gate_rows VALUES (2)'), 1)
|
|
assert await connection.scalar(sa.select(sa.func.count()).select_from(Row)) == 1
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_failed_rollback_is_retried_before_pool_reuse(database):
|
|
engine, gate = database
|
|
async with engine.connect() as connection:
|
|
await connection.exec_driver_sql('INSERT INTO gate_rows VALUES (1)')
|
|
driver = (await connection.get_raw_connection()).driver_connection
|
|
original = driver.rollback
|
|
|
|
async def fail_once():
|
|
driver.rollback = original
|
|
raise sqlite3.OperationalError('injected rollback failure')
|
|
|
|
driver.rollback = fail_once
|
|
with pytest.raises(sa.exc.OperationalError, match='injected rollback failure'):
|
|
await connection.rollback()
|
|
assert gate.lock.locked()
|
|
async with engine.begin() as connection:
|
|
await asyncio.wait_for(connection.exec_driver_sql('INSERT INTO gate_rows VALUES (2)'), 1)
|
|
assert await connection.scalar(sa.select(sa.func.count()).select_from(Row)) == 1
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_pool_reset_rollback_failure_closes_before_admitting_writer(database):
|
|
engine, gate = database
|
|
connection = await engine.connect()
|
|
await connection.exec_driver_sql('INSERT INTO gate_rows VALUES (1)')
|
|
driver = (await connection.get_raw_connection()).driver_connection
|
|
|
|
async def broken_rollback():
|
|
raise sqlite3.OperationalError('persistent rollback failure')
|
|
|
|
driver.rollback = broken_rollback
|
|
with pytest.raises(sa.exc.OperationalError, match='persistent rollback failure'):
|
|
await connection.rollback()
|
|
await connection.close()
|
|
async with engine.begin() as connection:
|
|
await asyncio.wait_for(connection.exec_driver_sql('INSERT INTO gate_rows VALUES (2)'), 1)
|
|
assert await connection.scalar(sa.select(sa.func.count()).select_from(Row)) == 1
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
@pytest.mark.parametrize(
|
|
'statement,writes',
|
|
[
|
|
('WITH select1 AS (SELECT 7) INSERT INTO gate_rows SELECT * FROM select1', True),
|
|
('\ufeffINSERT INTO gate_rows VALUES (1)', True),
|
|
('WITH x AS (SELECT 2)\ufeff INSERT INTO gate_rows SELECT * FROM x', True),
|
|
('WITH 中SELECT文 AS (SELECT 8) INSERT INTO gate_rows SELECT * FROM 中SELECT文', True),
|
|
('WITH update1 AS (SELECT 9) SELECT * FROM update1', False),
|
|
('WITH a😺SELECT AS (SELECT 8) INSERT INTO gate_rows SELECT * FROM a😺SELECT', True),
|
|
('WITH a😺UPDATE AS (SELECT 9) SELECT * FROM a😺UPDATE', False),
|
|
('WITH "insert" AS (SELECT 9) SELECT * FROM "insert"', False),
|
|
('WITH replace AS (SELECT 9) SELECT * FROM replace', False),
|
|
(
|
|
'WITH RECURSIVE select1(n) AS (VALUES(1) UNION ALL SELECT n+1 FROM select1 WHERE n<3), '
|
|
'update1 AS NOT MATERIALIZED (SELECT n FROM select1) INSERT INTO gate_rows SELECT n FROM update1',
|
|
True,
|
|
),
|
|
('PRAGMA "incremental_vacuum"', True),
|
|
('PRAGMA main."wal_checkpoint"', True),
|
|
('PRAGMA [user_version](4)', True),
|
|
("PRAGMA table_info('a=b')", False),
|
|
('PRAGMA table_info(gate_rows) /* = */', False),
|
|
],
|
|
)
|
|
async def test_direct_sql_classification_matches_sqlite(database, statement, writes):
|
|
engine, gate = database
|
|
async with engine.connect() as connection:
|
|
acquired = gate.acquisitions
|
|
result = await connection.exec_driver_sql(statement)
|
|
if result.returns_rows:
|
|
result.all()
|
|
await connection.commit()
|
|
assert gate.acquisitions - acquired == int(writes)
|
|
if writes and statement.startswith('WITH'):
|
|
assert await connection.scalar(sa.select(sa.func.count()).select_from(Row)) > 0
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
@pytest.mark.parametrize('physically_closed', [False, True])
|
|
@pytest.mark.parametrize('method', ['close', 'invalidate'])
|
|
async def test_close_failure_fails_queued_and_future_writers(database, physically_closed, method):
|
|
engine, gate = database
|
|
owner = await engine.connect()
|
|
waiting = await engine.connect()
|
|
await owner.exec_driver_sql('INSERT INTO gate_rows VALUES (1)')
|
|
driver = (await owner.get_raw_connection()).driver_connection
|
|
original_close = driver.close
|
|
|
|
async def broken_close():
|
|
if physically_closed:
|
|
await original_close()
|
|
raise sqlite3.OperationalError('injected close completion error')
|
|
|
|
driver.close = broken_close
|
|
task = asyncio.create_task(waiting.exec_driver_sql('INSERT INTO gate_rows VALUES (2)'))
|
|
await queued(gate)
|
|
try:
|
|
if method == 'invalidate':
|
|
# Pool invalidation logs/swallow close errors; the gate must still fail closed.
|
|
await owner.invalidate()
|
|
else:
|
|
with pytest.raises(sqlite3.OperationalError, match='injected close completion error'):
|
|
await owner.run_sync(lambda connection: connection.connection.dbapi_connection.close())
|
|
with pytest.raises(RuntimeError, match='gate unavailable.*restart the process'):
|
|
await asyncio.wait_for(task, 1)
|
|
with pytest.raises(RuntimeError, match='gate unavailable.*restart the process'):
|
|
await waiting.exec_driver_sql('INSERT INTO gate_rows VALUES (3)')
|
|
assert gate.failed and not gate.lock.locked()
|
|
finally:
|
|
driver.close = original_close
|
|
await original_close()
|
|
await owner.invalidate()
|
|
await owner.close()
|
|
await waiting.close()
|
|
|
|
|
|
def test_engine_reuse_on_another_loop_fails_explicitly_even_after_dispose(tmp_path):
|
|
engine = create_async_engine(f'sqlite+aiosqlite:///{tmp_path / "loops.db"}')
|
|
gate = install(engine)
|
|
|
|
async def first_loop():
|
|
async with engine.begin() as connection:
|
|
await connection.exec_driver_sql('CREATE TABLE example (id INTEGER)')
|
|
# Force asyncio.Lock itself to bind through contention too.
|
|
await gate.acquire()
|
|
waiter = asyncio.create_task(gate.acquire())
|
|
await queued(gate)
|
|
gate.lock.release()
|
|
await waiter
|
|
gate.lock.release()
|
|
await engine.dispose()
|
|
|
|
async def other_loop():
|
|
with pytest.raises(RuntimeError, match='another event loop; create a new engine'):
|
|
async with engine.connect():
|
|
pass
|
|
assert not gate.lock.locked()
|
|
await engine.dispose()
|
|
|
|
for scenario in (first_loop, other_loop):
|
|
loop = asyncio.new_event_loop()
|
|
try:
|
|
loop.run_until_complete(scenario())
|
|
loop.run_until_complete(loop.shutdown_asyncgens())
|
|
finally:
|
|
loop.close()
|