import asyncio import contextlib import json import logging import os from collections.abc import Iterable from dataclasses import fields from datetime import UTC, datetime from email.utils import formatdate from aiohttp import web from sqlalchemy import text from sqlalchemy.ext.asyncio import AsyncEngine, AsyncSession, async_sessionmaker, create_async_engine from sqlalchemy.ext.asyncio.engine import AsyncConnection from .Events import EventBus, Events from .ItemDTO import ItemDTO from .operations import Operation, matches_condition from .Services import Services from .Singleton import ThreadSafe from .Utils import init_class LOG: logging.Logger = logging.getLogger(__name__) ITEM_DTO_FIELDS: set[str] = {f.name for f in fields(ItemDTO)} class Terminator: pass class _Op: """Queued write operation.""" __slots__ = ("item", "key", "keys", "op", "type_value") def __init__( self, op: Terminator | str, type_value: str, item: ItemDTO | None, key: str | None, keys: list[str] | None ): self.op = op self.type_value = type_value self.item = item self.key = key self.keys = keys class SqliteStore(metaclass=ThreadSafe): @staticmethod def get_instance(db_path: str | None = None) -> "SqliteStore": return SqliteStore(db_path=db_path) def attach(self, app: web.Application): Services.get_instance().add("sqlite_store", self) async def handle_event(_, __): await self.get_connection() EventBus.get_instance().subscribe(Events.STARTED, handle_event, "SqliteStore.get_connection") app.on_shutdown.append(self.on_shutdown) async def on_shutdown(self, _: web.Application): LOG.debug("Shutting down SqliteStore...") await self.close() LOG.debug("SqliteStore shut down complete.") def __init__(self, db_path: str, *, max_pending: int = 200, flush_interval: float = 0.05): self._db_path: str = db_path self._engine: AsyncEngine | None = None self._conn: AsyncConnection | None = None self._sessionmaker: async_sessionmaker[AsyncSession] | None = None self._queue: asyncio.Queue[_Op] | None = None self._task: asyncio.Task | None = None self._lock: asyncio.Lock = asyncio.Lock() self._flush_interval: float = flush_interval self._max_pending: int = max_pending async def __aenter__(self) -> "SqliteStore": await self.get_connection() return self async def __aexit__(self, exc_type, exc, tb) -> None: await self.close() def sessionmaker(self) -> async_sessionmaker[AsyncSession]: """ Return the SQLAlchemy async sessionmaker. This allows other parts of the system to create SQLAlchemy sessions that share the same database connection/engine. Returns: async_sessionmaker[AsyncSession]: The sessionmaker instance. Raises: RuntimeError: If called before connection is initialized. """ if not self._sessionmaker: msg = "Database connection not initialized. Call get_connection() first or use within async context." raise RuntimeError(msg) return self._sessionmaker async def fetch_saved(self, type_value: str) -> list[tuple[str, ItemDTO]]: await self.get_connection() result = await self._conn.execute( text( 'SELECT "id", "data", "created_at" FROM "history" WHERE "type" = :type_value ORDER BY "created_at" ASC' ), {"type_value": type_value}, ) rows = result.mappings().all() items: list[tuple[str, ItemDTO]] = [] for row in rows: row_date = datetime.strptime(row["created_at"], "%Y-%m-%d %H:%M:%S") # noqa: DTZ007 data = json.loads(row["data"]) data.pop("_id", None) item = init_class(ItemDTO, data, ITEM_DTO_FIELDS) item._id = row["id"] item.datetime = formatdate(row_date.replace(tzinfo=UTC).timestamp()) items.append((row["id"], item)) return items async def exists(self, type_value: str, key: str | None = None, url: str | None = None) -> bool: return await self.get(type_value, key=key, url=url) is not None async def get(self, type_value: str, key: str | None = None, url: str | None = None) -> bool: if not key and not url: msg = "key or url must be provided." raise KeyError(msg) await self.get_connection() clauses: list[str] = [] params: dict[str, str] = {"type_value": type_value} if key: clauses.append('"id" = :key') params["key"] = key if url: clauses.append("json_extract(data, '$.url') = :url") params["url"] = url where_clause = " OR ".join(clauses) query = ( f'SELECT "id", "data", "created_at" FROM "history" WHERE "type" = :type_value AND ({where_clause}) LIMIT 1' # noqa: S608 ) result = await self._conn.execute(text(query), params) row = result.mappings().first() if not row: return None row_date = datetime.strptime(row["created_at"], "%Y-%m-%d %H:%M:%S") # noqa: DTZ007 data = json.loads(row["data"]) data.pop("_id", None) item = init_class(ItemDTO, data, ITEM_DTO_FIELDS) item._id = row["id"] item.datetime = formatdate(row_date.replace(tzinfo=UTC).timestamp()) return item async def get_by_id(self, type_value: str, id: str) -> ItemDTO | None: await self.get_connection() result = await self._conn.execute( text('SELECT "id", "data", "created_at" FROM "history" WHERE "type" = :type_value AND "id" = :id'), {"type_value": type_value, "id": id}, ) row = result.mappings().first() if not row: return None row_date = datetime.strptime(row["created_at"], "%Y-%m-%d %H:%M:%S") # noqa: DTZ007 data = json.loads(row["data"]) data.pop("_id", None) item = init_class(ItemDTO, data, ITEM_DTO_FIELDS) item._id = row["id"] item.datetime = formatdate(row_date.replace(tzinfo=UTC).timestamp()) return item async def get_item(self, type_value: str, **kwargs) -> ItemDTO | None: """ Return first item of type matching *any* condition. Mirrors :meth:`DataStore.get_item` semantics: if any provided condition matches (OR logic) return the first row by creation time. Returns None when kwargs is empty or no row matches. """ if not kwargs: return None await self.get_connection() clauses: list[str] = [] params: dict[str, str | float | int] = {"type_value": type_value} param_counter = 0 def _safe_key(key: str) -> str | None: return key if key.replace("_", "").isalnum() else None for key, raw_value in kwargs.items(): safe_key = _safe_key(key) if not safe_key: continue if isinstance(raw_value, tuple) and len(raw_value) == 2: operation, value = raw_value else: operation, value = Operation.EQUAL, raw_value if isinstance(operation, str): try: operation = Operation(operation) except ValueError: operation = Operation.EQUAL path: str = f"$.{safe_key}" json_extract: str = f"json_extract(data, '{path}')" param_name: str = f"param_{param_counter}" param_counter += 1 if Operation.EQUAL == operation: clauses.append(f"{json_extract} = :{param_name}") params[param_name] = value elif Operation.NOT_EQUAL == operation: clauses.append(f"{json_extract} != :{param_name}") params[param_name] = value elif Operation.CONTAIN == operation: clauses.append(f"{json_extract} LIKE :{param_name} ESCAPE '\\'") params[param_name] = f"%{value}%" elif Operation.NOT_CONTAIN == operation: clauses.append(f"({json_extract} IS NULL OR {json_extract} NOT LIKE :{param_name} ESCAPE '\\')") params[param_name] = f"%{value}%" elif Operation.STARTS_WITH == operation: clauses.append(f"{json_extract} LIKE :{param_name} ESCAPE '\\'") params[param_name] = f"{value}%" elif Operation.ENDS_WITH == operation: clauses.append(f"{json_extract} LIKE :{param_name} ESCAPE '\\'") params[param_name] = f"%{value}" elif Operation.GREATER_THAN == operation: clauses.append(f"{json_extract} > :{param_name}") params[param_name] = value elif Operation.LESS_THAN == operation: clauses.append(f"{json_extract} < :{param_name}") params[param_name] = value elif Operation.GREATER_EQUAL == operation: clauses.append(f"{json_extract} >= :{param_name}") params[param_name] = value elif Operation.LESS_EQUAL == operation: clauses.append(f"{json_extract} <= :{param_name}") params[param_name] = value if not clauses: return None where_clause: str = " OR ".join(f"({clause})" for clause in clauses) query: str = f'SELECT "id", "data", "created_at" FROM "history" WHERE "type" = :type_value AND ({where_clause}) ORDER BY "created_at" ASC LIMIT 1' # noqa: S608 result = await self._conn.execute(text(query), params) row = result.mappings().first() if not row: return None row_date = datetime.strptime(row["created_at"], "%Y-%m-%d %H:%M:%S") # noqa: DTZ007 data = json.loads(row["data"]) data.pop("_id", None) item = init_class(ItemDTO, data, ITEM_DTO_FIELDS) item._id = row["id"] item.datetime = formatdate(row_date.replace(tzinfo=UTC).timestamp()) return item if any(matches_condition(k, v, item.__dict__) for k, v in kwargs.items()) else None async def count(self, type_value: str, status_filter: str | None = None) -> int: await self.get_connection() where_clauses: list[str] = ['"type" = :type_value'] params: dict[str, str] = {"type_value": type_value} if status_filter: if status_filter.startswith("!"): status_value = status_filter[1:] where_clauses.append("json_extract(data, '$.status') != :status") params["status"] = status_value else: where_clauses.append("json_extract(data, '$.status') = :status") params["status"] = status_filter where_clause: str = " AND ".join(where_clauses) count_query: str = f'SELECT COUNT(*) as count FROM "history" WHERE {where_clause}' # noqa: S608 result = await self._conn.execute(text(count_query), params) row = result.mappings().first() return row["count"] if row else 0 async def paginate( self, type_value: str, page: int, per_page: int, order: str, status_filter: str | None = None, ) -> tuple[list[tuple[str, ItemDTO]], int, int, int]: await self.get_connection() where_clauses: list[str] = ['"type" = :type_value'] params: dict[str, str | int] = {"type_value": type_value} if status_filter: if status_filter.startswith("!"): status_value = status_filter[1:] where_clauses.append("json_extract(data, '$.status') != :status") params["status"] = status_value else: where_clauses.append("json_extract(data, '$.status') = :status") params["status"] = status_filter where_clause: str = " AND ".join(where_clauses) count_query: str = f'SELECT COUNT(*) as count FROM "history" WHERE {where_clause}' # noqa: S608 result = await self._conn.execute(text(count_query), params) count_row = result.mappings().first() total_items = count_row["count"] if count_row else 0 total_pages = (total_items + per_page - 1) // per_page if total_items > 0 else 1 if page > total_pages and total_items > 0: page = total_pages offset: int = (page - 1) * per_page params["limit"] = per_page params["offset"] = offset query: str = f'SELECT "id", "data", "created_at" FROM "history" WHERE {where_clause} ORDER BY "created_at" {order} LIMIT :limit OFFSET :offset' # noqa: S608 result = await self._conn.execute(text(query), params) rows = result.mappings().all() items: list[tuple[str, ItemDTO]] = [] for row in rows: row_date = datetime.strptime(row["created_at"], "%Y-%m-%d %H:%M:%S") # noqa: DTZ007 data = json.loads(row["data"]) data.pop("_id", None) item = init_class(ItemDTO, data, ITEM_DTO_FIELDS) item._id = row["id"] item.datetime = formatdate(row_date.replace(tzinfo=UTC).timestamp()) items.append((row["id"], item)) return items, total_items, page, total_pages async def upsert(self, type_value: str, item: ItemDTO) -> None: await self.get_connection() await self._upsert_now(type_value, item) async def delete(self, type_value: str, key: str) -> None: await self.get_connection() await self._conn.execute( text('DELETE FROM "history" WHERE "type" = :type_value AND "id" = :key'), {"type_value": type_value, "key": key}, ) await self._conn.commit() async def bulk_delete(self, type_value: str, keys: Iterable[str]) -> int: await self.get_connection() keys_list: list[str] = list(keys) if not keys_list: return 0 placeholders = ",".join(f":key_{i}" for i in range(len(keys_list))) params: dict[str, str] = {"type_value": type_value} params.update({f"key_{i}": key for i, key in enumerate(keys_list)}) result = await self._conn.execute( text(f'DELETE FROM "history" WHERE "type" = :type_value AND "id" IN ({placeholders})'), # noqa: S608 params, ) await self._conn.commit() return result.rowcount if result else 0 async def enqueue_upsert(self, type_value: str, item: ItemDTO) -> None: await self._enqueue(_Op("upsert", type_value, item, None, None)) async def enqueue_delete(self, type_value: str, key: str) -> None: await self._enqueue(_Op("delete", type_value, None, key, None)) async def enqueue_bulk_delete(self, type_value: str, keys: Iterable[str]) -> None: await self._enqueue(_Op("bulk_delete", type_value, None, None, list(keys))) async def flush(self) -> None: if self._queue: await self._queue.join() async def shutdown(self) -> None: """Flush pending writes and stop writer task.""" if self._queue: try: await self._queue.put(_Op(Terminator(), "", None, None, None)) LOG.debug("Waiting for SqliteStore queue to drain...") await asyncio.wait_for(self._queue.join(), timeout=2) except TimeoutError: LOG.warning("SqliteStore queue did not drain within timeout; forcing writer shutdown.") if self._task: LOG.debug("Waiting for SqliteStore writer task to finish...") if not self._task.done(): self._task.cancel() with contextlib.suppress(asyncio.CancelledError): await asyncio.wait_for(self._task, timeout=2) self._task = None self._queue = None async def close(self) -> None: await self.shutdown() if self._conn: await self._conn.close() self._conn = None if self._engine: await self._engine.dispose() self._engine = None self._sessionmaker = None async def _enqueue(self, op: _Op) -> None: self._ensure_worker() await self._queue.put(op) def _ensure_worker(self) -> None: if self._queue is None: self._queue = asyncio.Queue(maxsize=self._max_pending) self._task = asyncio.create_task(self._writer(), name="sqlite-store-writer") async def _writer(self) -> None: while True: op = await self._queue.get() if isinstance(op.op, Terminator): self._queue.task_done() break try: async with self._lock: await self._apply(op) except Exception as ex: LOG.exception(ex) finally: self._queue.task_done() await asyncio.sleep(self._flush_interval) async def _apply(self, op: _Op) -> None: await self.get_connection() if op.op == "upsert" and op.item: await self._upsert_now_conn(self._conn, op.type_value, op.item) await self._conn.commit() elif op.op == "delete" and op.key: await self._conn.execute( text('DELETE FROM "history" WHERE "type" = :type_value AND "id" = :key'), {"type_value": op.type_value, "key": op.key}, ) await self._conn.commit() elif op.op == "bulk_delete" and op.keys: placeholders = ",".join(f":key_{i}" for i in range(len(op.keys))) params: dict[str, str] = {"type_value": op.type_value} params.update({f"key_{i}": key for i, key in enumerate(op.keys)}) await self._conn.execute( text(f'DELETE FROM "history" WHERE "type" = :type_value AND "id" IN ({placeholders})'), # noqa: S608 params, ) await self._conn.commit() async def _upsert_now(self, type_value: str, item: ItemDTO) -> None: await self.get_connection() await self._upsert_now_conn(self._conn, type_value, item) await self._conn.commit() async def _upsert_now_conn(self, conn, type_value: str, item: ItemDTO) -> None: """Helper to upsert using an existing connection.""" sql = """ INSERT INTO "history" ("id", "type", "url", "data") VALUES (:id, :type, :url, :data) ON CONFLICT DO UPDATE SET "type" = :type2, "url" = :url2, "data" = :data2, created_at = :created_at """ encoded = item.json() await conn.execute( text(sql.strip()), { "id": item._id, "type": type_value, "url": item.url, "data": encoded, "type2": type_value, "url2": item.url, "data2": encoded, "created_at": datetime.now(UTC).strftime("%Y-%m-%d %H:%M:%S"), }, ) async def execute_raw(self, query: str, params: dict | tuple | None = None) -> None: """Execute raw SQL query (for testing purposes).""" await self.get_connection() if isinstance(params, tuple): # Convert positional params to dict for SQLAlchemy # Assuming queries use ? placeholders, we need to count them placeholders = query.count("?") if placeholders != len(params): msg = ( f"Parameter count mismatch: query has {placeholders} placeholders but {len(params)} params provided" ) raise ValueError(msg) # Create numbered parameters param_dict = {f"p{i}": params[i] for i in range(len(params))} # Replace ? with :p0, :p1, etc. for i in range(len(params)): query = query.replace("?", f":p{i}", 1) await self._conn.execute(text(query), param_dict) elif isinstance(params, dict): await self._conn.execute(text(query), params) else: await self._conn.execute(text(query)) await self._conn.commit() async def fetch_raw(self, query: str, params: dict | tuple | None = None): """Fetch results from raw SQL query (for testing purposes).""" await self.get_connection() if isinstance(params, tuple): # Convert positional params to dict for SQLAlchemy placeholders = query.count("?") if placeholders != len(params): msg = ( f"Parameter count mismatch: query has {placeholders} placeholders but {len(params)} params provided" ) raise ValueError(msg) param_dict = {f"p{i}": params[i] for i in range(len(params))} for i in range(len(params)): query = query.replace("?", f":p{i}", 1) result = await self._conn.execute(text(query), param_dict) elif isinstance(params, dict): result = await self._conn.execute(text(query), params) else: result = await self._conn.execute(text(query)) return result.mappings().all() async def get_connection(self) -> AsyncConnection: if self._conn: return self._conn if not self._db_path: msg = "No database path specified for SqliteStore." raise RuntimeError(msg) from app.library import migrate from app.main import ROOT_PATH if self._db_path.startswith(":memory"): db_url: str = "sqlite+aiosqlite:///file::memory:?cache=shared&uri=true" else: os.makedirs(os.path.dirname(self._db_path) or ".", exist_ok=True) db_url: str = f"sqlite+aiosqlite:///{self._db_path}" self._engine = create_async_engine( db_url, echo=False, connect_args={"check_same_thread": False, "uri": self._db_path.startswith(":memory")}, ) self._conn = await self._engine.connect() if version := await migrate.get_version(self._conn): LOG.debug(f"DB Version: '{version}'.") await migrate.upgrade(self._conn, ROOT_PATH / "migrations") if not version: LOG.debug(f"DB Version after initial migration: '{await migrate.get_version(self._conn)}'.") await self._conn.execute(text("PRAGMA journal_mode=wal")) await self._conn.execute(text("PRAGMA busy_timeout=5000")) await self._conn.execute(text("PRAGMA foreign_keys=ON")) await self._conn.commit() self._sessionmaker = async_sessionmaker(bind=self._engine, class_=AsyncSession, expire_on_commit=False) return self._conn