fix(persistence): serialize SQLite runtime write transactions

This commit is contained in:
huanghuoguoguo
2026-10-27 14:42:35 +08:00
parent a89ddf12e1
commit 8e21495577
6 changed files with 735 additions and 8 deletions
+2 -2
View File
@@ -272,13 +272,13 @@ Maintenance rule: when adding, removing, or changing an HTTP endpoint that shoul
Persistence is centered on `pkg/persistence/mgr.py`.
- SQLite is the default database; PostgreSQL is supported.
- SQLite is the default database; PostgreSQL is supported. The shared SQLite runtime engine admits one writer at a time from the first write through physical commit/rollback/close, including ORM and direct connection writes. Readers do not acquire the gate. It is process-local and does not coordinate other engines or processes; PostgreSQL is unchanged. An engine must stay on its original event loop, even after disposal. A connection-close failure makes the gate reject subsequent SQL explicitly until the process is restarted; queued callers are woken to fail rather than left blocked.
- Every SQLite connection sets a 64 MiB `journal_size_limit`, including connections recreated after migration or pool disposal. SQLite's default automatic checkpoint policy and synchronous setting remain unchanged. The limit reduces retained WAL space when the WAL resets; it does not cap active transactions or WAL pinned by readers. Online maintenance does not periodically force checkpoints or truncate WAL files. Diagnose persistent excess growth by checking checkpoint progress and actual transaction lifetimes, then shorten the responsible transaction rather than forcing readers to finish.
- After schema initialization and migration, every SQLite startup completes `wal_checkpoint(TRUNCATE)` before runtime setup and producers start. Lock waits allow up to 60 seconds; checkpoint I/O may take longer. An incomplete checkpoint fails startup without undoing successful migrations, and the next startup retries even at schema head. The connection busy timeout is restored afterward. PostgreSQL is unaffected.
- SQLite files use incremental auto-vacuum. Migration `0042_sqlite_incremental_vacuum` converts existing files before producers start, with a verified backup and restoration on conversion failure when a full rebuild is required. Fresh files enable the format before creating tables. Monitoring cleanup reclaims free pages in independent adaptive batches (1–512 pages, starting at 8, targeting 25 ms), pausing 250 ms between batches and starting no new batches after a two-second round budget. Busy writers take priority. No full `VACUUM` or forced WAL checkpoint runs during online cleanup.
- Models live under `pkg/entity/persistence/`.
- `pkg/persistence/media.py::MediaReferenceText` validates bound conversation payloads on ORM/Core writes while preserving the existing TEXT schema and string API. It rejects recognized structured inline images: retained images go through `MediaCache.externalize_chain_dump()` before serialization, while metadata-only snapshots and transcripts use `redact_inline_images()`. Reads and comparisons preserve historical values for incremental cleanup. Runtime/adapter messages, arbitrary diagnostic text, and opaque strings are outside this contract.
- `pkg/persistence/resource_urls.py` interns structured resource URL fields in retained messages, history and monitoring at the typed SQL boundary, in the source transaction. `MediaReferenceJSON` preserves native JSON monitoring details. Payloads store opaque references to the global `resource_urls` catalog; typed SQL reads and comparisons expand them to current URLs, preserving API shapes and source-record authorization. Configuration URLs and URLs in free text are excluded. Catalog lookups use only the primary-key and URL-hash unique indexes. Successful refreshes overwrite the current URL and persist the check time even when standard `get_msg` returns the same URL. Converging URLs retain only ID redirects, with no old URL. PostgreSQL row locks and SQLite's immediate write transaction serialize concurrent refreshes; SQLite holds its write lock during the upstream call. Bounded in-memory aliases let reused Runner payloads resolve old URLs without retaining those URLs in the database. Migration `0041_resource_urls` normalizes existing payloads using their source creation timestamps before producers start; downgrade expands references before dropping the catalog.
- `pkg/persistence/resource_urls.py` interns structured resource URL fields in retained messages, history and monitoring at the typed SQL boundary, in the source transaction. `MediaReferenceJSON` preserves native JSON monitoring details. Payloads store opaque references to the global `resource_urls` catalog; typed SQL reads and comparisons expand them to current URLs, preserving API shapes and source-record authorization. Configuration URLs and URLs in free text are excluded. Catalog lookups use only the primary-key and URL-hash unique indexes. Successful refreshes overwrite the current URL and persist the check time even when standard `get_msg` returns the same URL. Converging URLs retain only ID redirects, with no old URL. PostgreSQL row locks serialize concurrent refreshes. SQLite serializes refreshes within the engine, performs upstream calls outside its write transaction, and rechecks catalog rows before applying results in a short immediate transaction. Bounded in-memory aliases let reused Runner payloads resolve old URLs without retaining those URLs in the database. Migration `0041_resource_urls` normalizes existing payloads using their source creation timestamps before producers start; downgrade expands references before dropping the catalog.
- Timezone-less SQL `DateTime` columns store UTC-naive values on both backends. Normalize aware values with `pkg/persistence/datetime_utils.py::as_naive_utc` at the bind boundary, including deadlines, leases, event times, and query/retention cutoffs; do not merely strip an offset. Existing naive rows represent UTC. Restore UTC awareness for application comparisons and epoch serialization. Runner ledger/event/transcript stores follow this existing schema contract without a data migration or changes to journal/tenant authorization.
- Fresh schemas are created from current metadata, then Alembic migrations run to head. LangBot 4.x does not upgrade 3.x databases.
- New schema changes use Alembic under `pkg/persistence/alembic/versions/`; there is no legacy migration chain in 4.x.
@@ -3,7 +3,7 @@ from __future__ import annotations
import sqlalchemy
import sqlalchemy.ext.asyncio as sqlalchemy_asyncio
from .. import database
from .. import database, sqlite_write_gate
WAL_RETAIN_BYTES = 64 * 1024 * 1024
@@ -27,6 +27,7 @@ class SQLiteDatabaseManager(database.BaseDatabaseManager):
db_file_path = self.ap.instance_config.data.get('database', {}).get('sqlite', {}).get('path', 'data/langbot.db')
engine_url = f'sqlite+aiosqlite:///{db_file_path}'
self.engine = sqlalchemy_asyncio.create_async_engine(engine_url)
self.write_gate = sqlite_write_gate.install(self.engine)
sqlalchemy.event.listen(self.engine.sync_engine, 'connect', _configure_connection)
# Reader transactions must not block independent plugin callback commits.
# Configure the persistent journal before serving requests, not per write.
@@ -35,3 +36,6 @@ class SQLiteDatabaseManager(database.BaseDatabaseManager):
await connection.execute(sqlalchemy.text('PRAGMA auto_vacuum=INCREMENTAL'))
await connection.execute(sqlalchemy.text('PRAGMA journal_mode=WAL'))
await connection.commit()
def resource_stats(self) -> dict[str, int]:
return self.write_gate.resource_stats()
+44 -5
View File
@@ -5,6 +5,7 @@ ORM and Core callers retain their existing string/JSON contracts. The catalog
is shared; authorization continues to belong to the enclosing business record.
"""
import asyncio
import datetime
import hashlib
import json
@@ -255,12 +256,23 @@ def create_postgres_expansion(connection):
async def refresh_images(engine, peer, payload, *, is_current=None):
"""Serialize refreshes on catalog identities and overwrite current URLs atomically."""
"""Serialize refreshes without holding SQLite's writer during upstream calls."""
from ..platform.adapters.aiocqhttp.image_urls import image_nodes
nodes = list(image_nodes(payload))
if not nodes:
return
if engine.dialect.name == 'sqlite':
sync_engine = engine.sync_engine
if not hasattr(sync_engine, '_langbot_image_refresh_lock'):
sync_engine._langbot_image_refresh_lock = asyncio.Lock()
async with sync_engine._langbot_image_refresh_lock:
return await _refresh_images(engine, peer, payload, nodes, is_current=is_current)
return await _refresh_images(engine, peer, payload, nodes, is_current=is_current)
async def _refresh_images(engine, peer, payload, nodes, *, is_current=None):
from ..platform.adapters.aiocqhttp.image_urls import nt_image_type
def connected():
return not peer.closed and (is_current is None or is_current())
@@ -273,10 +285,7 @@ async def refresh_images(engine, peer, payload, *, is_current=None):
await db.run_sync(lambda sync: intern(sync, json.dumps(payload), identities_out=references))
identities = set(references.values())
async with engine.connect() as db:
if engine.dialect.name == 'sqlite':
await db.exec_driver_sql('BEGIN IMMEDIATE')
else:
await db.begin()
await db.begin()
try:
rows = (
(
@@ -301,6 +310,8 @@ async def refresh_images(engine, peer, payload, *, is_current=None):
# Capture origins even when renewal is not due or identification fails.
for node, message_id, image_id in nodes:
peer.image_urls.remember(by_url[node['url']]['url'], message_id, image_id)
if engine.dialect.name == 'sqlite':
await db.rollback()
# Identification and upstream credential handling remain connection-local.
if not connected() or not await peer.image_urls.identify():
return
@@ -327,11 +338,39 @@ async def refresh_images(engine, peer, payload, *, is_current=None):
return
if not connected():
return
if engine.dialect.name == 'sqlite':
cached = {}
for previous, entry in zip(old, batch):
origin = peer.image_urls.references.get(previous)
entries = [
(peer.image_urls.urls, previous),
(peer.image_urls.urls, entry['image']['url']),
(peer.image_urls.messages, str(origin[0]) if origin else None),
(peer.image_urls.keys, nt_image_type(previous)),
]
cached[previous] = [(cache, key, cache[key]) for cache, key in entries if key in cache]
await db.exec_driver_sql('BEGIN IMMEDIATE')
stamp = now()
updated = set()
for (node, _, _, row), entry, previous in zip(due, batch, old):
if previous not in completed:
continue
if engine.dialect.name == 'sqlite' and row['id'] not in updated:
current = (await db.execute(sa.select(table).where(table.c.id == row['id']))).mappings().one()
if current != row:
# Another engine may have renewed or redirected this identity.
while current['canonical_id'] is not None:
current = (
(await db.execute(sa.select(table).where(table.c.id == current['canonical_id'])))
.mappings()
.one()
)
node['url'] = current['url']
# A rejected result must not return from the peer cache on the next refresh.
for cache, key, value in cached[previous]:
if cache.get(key) is value:
del cache[key]
continue
fresh = entry['image']['url']
node['url'] = fresh
if napcat and fresh == previous:
@@ -0,0 +1,252 @@
"""Serialize SQLite writers on one application engine, without locking readers."""
from __future__ import annotations
import asyncio
import re
import time
import sqlalchemy as sa
from sqlalchemy.dialects.sqlite.aiosqlite import AsyncAdapt_aiosqlite_connection
from sqlalchemy.util.concurrency import await_only
_TOKENS = re.compile(
r"\ufeff|--[^\n]*(?:\n|$)|/\*.*?\*/|'(?:''|[^'])*'|\"(?:\"\"|[^\"])*\"|`[^`]*`|\[[^\]]*\]|[A-Za-z_\x80-\U0010ffff][A-Za-z_0-9$\x80-\U0010ffff]*|[(),.=]",
re.S,
)
_WRITES = {'INSERT', 'UPDATE', 'DELETE', 'REPLACE', 'CREATE', 'ALTER', 'DROP', 'VACUUM', 'REINDEX', 'ATTACH', 'DETACH'}
def _writes(statement: str) -> bool:
words = [
m.group()
for m in _TOKENS.finditer(statement)
if m.group() != '\ufeff' and not m.group().startswith(('--', '/*'))
]
if not words:
return False
first = words[0].upper()
if first in _WRITES:
return True
if first == 'BEGIN':
return any(word.upper() in {'IMMEDIATE', 'EXCLUSIVE'} for word in words[1:])
if first == 'PRAGMA':
# SQLite accepts both name=value and name(value) for writable PRAGMAs.
read_arguments = {
'table_info',
'table_xinfo',
'table_list',
'index_info',
'index_xinfo',
'index_list',
'foreign_key_list',
'foreign_key_check',
'integrity_check',
'quick_check',
}
# Only the pragma name is an identifier; arguments/comments may contain
# '=', keywords or parentheses without changing a read into a write.
name_index = 3 if len(words) > 3 and words[2] == '.' else 1
if len(words) <= name_index:
return False
name = words[name_index].strip('"\'`[]').lower()
argument = words[name_index + 1 : name_index + 2]
return name in {'incremental_vacuum', 'wal_checkpoint', 'optimize'} or (
argument in (['='], ['(']) and name not in read_arguments
)
if first == 'WITH':
index = 2 if len(words) > 1 and words[1].upper() == 'RECURSIVE' else 1
def group_end(start):
depth = 0
for position in range(start, len(words)):
if words[position] == '(':
depth += 1
elif words[position] == ')':
depth -= 1
if depth == 0:
return position + 1
return len(words)
while index < len(words):
index += 1 # CTE name, including identifiers that resemble keywords.
if index < len(words) and words[index] == '(':
index = group_end(index) # Optional column names.
if index >= len(words) or words[index].upper() != 'AS':
return False
index += 1
while index < len(words) and words[index].upper() in {'NOT', 'MATERIALIZED'}:
index += 1
if index >= len(words) or words[index] != '(':
return False
index = group_end(index)
if index < len(words) and words[index] == ',':
index += 1
continue
return index < len(words) and words[index].upper() in _WRITES
return False
class SQLiteWriteGate:
def __init__(self):
self.lock = asyncio.Lock()
self.loop = None
self.failed = False
self.waiting = 0
self.acquisitions = 0
self.wait_seconds = 0.0
self.max_wait_seconds = 0.0
def check(self):
loop = asyncio.get_running_loop()
if self.loop is None:
self.loop = loop
elif self.loop is not loop:
raise RuntimeError('SQLite writer gate belongs to another event loop; create a new engine for this loop')
if self.failed:
raise RuntimeError(
'SQLite writer gate unavailable after connection close failure; restart the process before writing'
)
async def acquire(self):
self.check()
start = time.monotonic()
self.waiting += 1
try:
await self.lock.acquire()
finally:
self.waiting -= 1
try:
self.check()
except BaseException:
self.lock.release()
raise
elapsed = time.monotonic() - start
self.acquisitions += 1
self.wait_seconds += elapsed
self.max_wait_seconds = max(self.max_wait_seconds, elapsed)
def resource_stats(self):
return {
'sqlite_writer_active': int(self.lock.locked()),
'sqlite_writer_waiting': self.waiting,
'sqlite_writer_acquisitions': self.acquisitions,
'sqlite_writer_wait_ms_total': int(self.wait_seconds * 1000),
'sqlite_writer_wait_ms_max': int(self.max_wait_seconds * 1000),
}
class _GatedConnection(AsyncAdapt_aiosqlite_connection):
"""Release after DBAPI completion, not SQLAlchemy's pre-commit event."""
__slots__ = ('gate', 'writer')
def __init__(self, dbapi, connection, gate):
super().__init__(dbapi, connection)
self.gate = gate
self.writer = False
def release(self):
if self.writer:
self.writer = False
self.gate.lock.release()
def acquire(self):
self.gate.check()
if not self.writer:
await_only(self.gate.acquire())
self.writer = True
def fail_closed(self):
# Wake queued callers to fail explicitly, never admit a new writer when
# physical close completion is uncertain. Recovery requires process restart.
self.gate.failed = True
self.release()
def _terminate_force_close(self):
try:
stop = getattr(self.driver_connection, 'stop', None)
closed = (
stop() if stop is not None else asyncio.get_running_loop().create_task(self._terminate_graceful_close())
)
if closed is None:
self.fail_closed()
else:
closed.add_done_callback(self._closed)
except BaseException:
self.fail_closed()
raise
def _closed(self, future):
if future.cancelled() or future.exception() is not None:
self.fail_closed()
else:
self.release()
def commit(self):
super().commit()
self.release()
def rollback(self):
super().rollback()
self.release()
def close(self):
try:
super().close()
except BaseException:
self.fail_closed()
raise
self.release()
async def _terminate_graceful_close(self):
try:
await self.driver_connection.close()
except BaseException:
self.fail_closed()
raise
self.release()
def install(engine):
"""Cover ORM flushes, Core statements and direct SQL on this async engine."""
if engine.dialect.name != 'sqlite':
return None
# shortcut: one runtime engine/process; independent writers need PostgreSQL.
gate = SQLiteWriteGate()
def connect(dialect, connection_record, args, kwargs):
gate.check()
adapted = dialect.loaded_dbapi.connect(*args, **kwargs)
return _GatedConnection(dialect.loaded_dbapi, adapted.driver_connection, gate)
def before(connection, cursor, statement, parameters, context, executemany):
gate.check()
if context.compiled is not None and context.compiled.statement.is_select:
return
if context.isinsert or context.isupdate or context.isdelete or _writes(statement):
connection.connection.dbapi_connection.acquire()
def after(connection, cursor, statement, parameters, context, executemany):
adapted = connection.connection.dbapi_connection
# Autocommit DDL/DML and releasing the outermost SAVEPOINT end the
# physical transaction even if SQLAlchemy still has an autobegin object.
if not adapted.driver_connection.in_transaction:
adapted.release()
def reset(connection, record, state):
# A failed SQLite COMMIT can retain its transaction after SQLAlchemy
# marks it inactive. Pool reset must not skip that physical rollback.
if connection.writer:
if state.asyncio_safe:
connection.rollback()
elif not engine.dialect.has_terminate:
# Older SQLAlchemy/aiosqlite cannot terminate from GC themselves.
connection._terminate_force_close()
sa.event.listen(engine.sync_engine, 'reset', reset)
sa.event.listen(engine.sync_engine, 'do_connect', connect)
sa.event.listen(engine.sync_engine, 'before_cursor_execute', before)
sa.event.listen(engine.sync_engine, 'after_cursor_execute', after)
return gate
@@ -20,6 +20,7 @@ from langbot.pkg.entity.persistence.event_log import EventLog
from langbot.pkg.entity.persistence.monitoring import MonitoringLLMCall
from langbot.pkg.persistence.media import MediaReferenceText, MediaReferenceJSON
from langbot.pkg.persistence.resource_urls import digest, refresh_images
from langbot.pkg.persistence.sqlite_write_gate import install as install_write_gate
from langbot.pkg.platform.adapters.aiocqhttp.image_urls import OneBotImageURLs
@@ -46,6 +47,7 @@ async def database(tmp_path, request):
engine = create_async_engine(url, connect_args={'server_settings': {'search_path': schema}})
else:
engine = create_async_engine(f'sqlite+aiosqlite:///{tmp_path / "catalog.db"}')
install_write_gate(engine)
try:
async with engine.begin() as db:
await db.run_sync(ResourceURL.__table__.create)
@@ -398,3 +400,74 @@ async def test_migration_preserves_source_creation_and_downgrade(database):
# The same engine can write and read literal payloads after downgrade.
await append(db, 'after', NEW)
assert NEW in await db.scalar(sa.select(EventLog.input_json).where(EventLog.event_id == 'after'))
@pytest.mark.asyncio
@pytest.mark.parametrize('napcat', [False, True])
@pytest.mark.parametrize('blocked_action', ['get_version_info', 'get_msg'])
@pytest.mark.parametrize('converge', [False, True])
async def test_sqlite_refresh_releases_writer_and_rejects_stale_result(database, blocked_action, converge, napcat):
if database.dialect.name != 'sqlite':
pytest.skip('SQLite transaction boundary')
async with database.begin() as db:
await append(db, 'a')
if converge:
await append(db, 'b', NEW)
await db.execute(sa.update(ResourceURL).values(created_at=dt.datetime(2000, 1, 1)))
source_id = await db.scalar(sa.select(ResourceURL.id).where(ResourceURL.url == OLD))
destination_id = await db.scalar(sa.select(ResourceURL.id).where(ResourceURL.url == NEW))
entered = asyncio.Event()
release = asyncio.Event()
upstream_url = OLD.replace('rkey=old', 'rkey=stale')
renewal_action = 'nc_get_rkey' if napcat else 'get_msg'
if blocked_action == 'get_msg':
blocked_action = renewal_action
async def action(name, **params):
if name == blocked_action:
entered.set()
await release.wait()
if name == 'get_version_info':
return {'app_name': 'NapCat.Onebot' if napcat else 'standard'}
if napcat:
return [{'type': 20, 'time': 1, 'ttl': 10**12, 'rkey': upstream_url.split('rkey=')[1]}]
return {'message': [{'type': 'image', 'data': {'file': 'one.jpg', 'url': upstream_url}}]}
peer = SimpleNamespace(closed=False, call_action=AsyncMock(side_effect=action))
peer.image_urls = OneBotImageURLs(peer)
payload = chain()
task = asyncio.create_task(refresh_images(database, peer, payload))
try:
await asyncio.wait_for(entered.wait(), 5)
# Plain text must retain its fast path while an image renewal waits.
await asyncio.wait_for(refresh_images(database, peer, [{'type': 'Plain', 'text': 'hello'}]), 1)
async def concurrent_writer():
async with database.begin() as db:
await append(db, 'unrelated', 'https://example.com/unrelated.png')
values = dict(url=NEW, url_hash=digest(NEW))
if converge:
values = dict(url=None, url_hash=None, canonical_id=destination_id)
await db.execute(sa.update(ResourceURL).where(ResourceURL.id == source_id).values(**values))
# It must commit while the external call is still blocked.
await asyncio.wait_for(concurrent_writer(), 2)
finally:
release.set()
await asyncio.wait_for(task, 5)
assert payload[1]['url'] == NEW
async with database.connect() as db:
assert NEW in await db.scalar(sa.select(EventLog.input_json).where(EventLog.event_id == 'a'))
assert await db.scalar(sa.select(EventLog.event_id).where(EventLog.event_id == 'unrelated')) == 'unrelated'
assert await db.scalar(sa.select(ResourceURL.last_checked_at).where(ResourceURL.id == source_id)) is None
# The rejected response also populated OneBot's connection-local caches.
# A later due refresh must ask upstream again instead of restoring that stale URL.
upstream_url = NEW
second_payload = chain(NEW)
await refresh_images(database, peer, second_payload)
assert second_payload[1]['url'] == NEW
assert sum(call.args[0] == renewal_action for call in peer.call_action.call_args_list) == 2
async with database.connect() as db:
assert NEW in await db.scalar(sa.select(EventLog.input_json).where(EventLog.event_id == 'a'))
assert not any('rkey=stale' in (url or '') for url in (await db.execute(sa.select(ResourceURL.url))).scalars())
@@ -0,0 +1,359 @@
"""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()