"""SQLAlchemy-backed generated-file library index.""" from __future__ import annotations import hashlib from datetime import UTC, datetime, timedelta from typing import Any from uuid import uuid4 from sqlalchemy import and_, delete, func, or_, select, update from sqlalchemy.exc import IntegrityError from sqlalchemy.ext.asyncio import AsyncSession, async_sessionmaker from deerflow.persistence.artifact_library.base import ArtifactLibraryStore from deerflow.persistence.artifact_library.model import ( ArtifactLibraryItemRow, ArtifactLibrarySyncStateRow, ) _EPOCH = datetime(1970, 1, 1, tzinfo=UTC) def _artifact_id(user_id: str, item: dict[str, Any]) -> str: identity = "\0".join((user_id, str(item.get("thread_id") or ""), str(item.get("path") or ""))) return hashlib.sha256(identity.encode("utf-8")).hexdigest() def _parse_datetime(value: object) -> datetime: if isinstance(value, datetime): parsed = value else: text = str(value or "").strip() if not text: return _EPOCH try: parsed = datetime.fromisoformat(text.replace("Z", "+00:00")) except ValueError: return _EPOCH if parsed.tzinfo is None: return parsed.replace(tzinfo=UTC) return parsed.astimezone(UTC) def _iso(value: datetime | None) -> str | None: return value.isoformat() if isinstance(value, datetime) else None def _state_dict(row: ArtifactLibrarySyncStateRow | None) -> dict[str, Any]: if row is None: return { "status": "idle", "generation": None, "started_at": None, "completed_at": None, "last_error": None, } return { "status": row.status, "generation": row.generation, "started_at": _iso(row.started_at), "completed_at": _iso(row.completed_at), "last_error": row.last_error, } def _item_dict(row: ArtifactLibraryItemRow) -> dict[str, Any]: return { "thread_id": row.thread_id, "thread_title": row.thread_title, "agent_id": row.agent_id, "agent_name": row.agent_name, "thread_type": row.thread_type, "source_type": row.source_type, "strategy_project_id": row.strategy_project_id, "strategy_name": row.strategy_name, "strategy_stale": row.strategy_stale, "strategy_source_kind": row.strategy_source_kind, "download_url": row.download_url, "path": row.path, "name": row.name, "size_bytes": row.size_bytes, "mime_type": row.mime_type, "modified_at": row.modified_at.isoformat(), "file_kind": row.file_kind, } def _apply_item( row: ArtifactLibraryItemRow, *, user_id: str, generation: str, item: dict[str, Any], indexed_at: datetime, ) -> None: row.user_id = user_id row.thread_id = str(item.get("thread_id") or "")[:128] row.thread_title = str(item.get("thread_title") or "")[:512] row.agent_id = str(item.get("agent_id"))[:128] if item.get("agent_id") else None row.agent_name = str(item.get("agent_name"))[:512] if item.get("agent_name") else None row.thread_type = str(item.get("thread_type"))[:64] if item.get("thread_type") else None row.source_type = str(item.get("source_type") or "general_qa")[:32] row.strategy_project_id = str(item.get("strategy_project_id"))[:128] if item.get("strategy_project_id") else None row.strategy_name = str(item.get("strategy_name"))[:512] if item.get("strategy_name") else None row.strategy_stale = item.get("strategy_stale") is True row.strategy_source_kind = str(item.get("strategy_source_kind"))[:32] if item.get("strategy_source_kind") else None row.download_url = str(item.get("download_url")) if item.get("download_url") else None row.path = str(item.get("path") or "") row.name = str(item.get("name") or "")[:512] row.size_bytes = max(0, int(item.get("size_bytes") or 0)) row.mime_type = str(item.get("mime_type"))[:255] if item.get("mime_type") else None row.modified_at = _parse_datetime(item.get("modified_at")) row.file_kind = str(item.get("file_kind") or "other")[:32] row.sync_generation = generation row.indexed_at = indexed_at class SqlArtifactLibraryStore(ArtifactLibraryStore): def __init__(self, session_factory: async_sessionmaker[AsyncSession]) -> None: self._sf = session_factory async def claim_sync( self, user_id: str, *, force: bool = False, stale_after_seconds: float = 15.0, lease_seconds: float = 300.0, ) -> str | None: now = datetime.now(UTC) generation = str(uuid4()) try: async with self._sf() as session: async with session.begin(): claim_filters = [ ArtifactLibrarySyncStateRow.user_id == user_id, # An expired (or legacy missing) lease may be reclaimed. # A live lease remains exclusive even for force=True. or_( ArtifactLibrarySyncStateRow.status != "syncing", ArtifactLibrarySyncStateRow.lease_expires_at.is_(None), ArtifactLibrarySyncStateRow.lease_expires_at <= now, ), ] if not force: stale_before = now - timedelta(seconds=stale_after_seconds) claim_filters.append( or_( and_( ArtifactLibrarySyncStateRow.status == "error", or_( ArtifactLibrarySyncStateRow.started_at.is_(None), ArtifactLibrarySyncStateRow.started_at <= stale_before, ), ), and_( ArtifactLibrarySyncStateRow.status != "error", or_( ArtifactLibrarySyncStateRow.completed_at.is_(None), ArtifactLibrarySyncStateRow.completed_at <= stale_before, ), ), ) ) # A conditional UPDATE is the portable compare-and-swap. # SQLite ignores SELECT ... FOR UPDATE, whereas this write # serializes claimers there and also works on MySQL, # PostgreSQL, and openGauss. result = await session.execute( update(ArtifactLibrarySyncStateRow) .where(*claim_filters) .values( status="syncing", generation=generation, started_at=now, lease_expires_at=now + timedelta(seconds=lease_seconds), last_error=None, ) ) if result.rowcount == 1: return generation row = await session.get(ArtifactLibrarySyncStateRow, user_id) if row is not None: return None session.add( ArtifactLibrarySyncStateRow( user_id=user_id, status="syncing", generation=generation, started_at=now, lease_expires_at=now + timedelta(seconds=lease_seconds), last_error=None, ) ) return generation except IntegrityError: # Two workers can race while inserting a user's first state row. # The winner owns the lease; the loser simply serves the same # persisted index without starting a duplicate filesystem scan. async with self._sf() as session: if await session.get(ArtifactLibrarySyncStateRow, user_id) is not None: return None raise async def finish_sync( self, user_id: str, generation: str, items: list[dict[str, Any]], ) -> bool: now = datetime.now(UTC) deduped = {_artifact_id(user_id, item): item for item in items} async with self._sf() as session: async with session.begin(): # Write first so SQLite also serializes this publisher against # a lease takeover. Requiring status=syncing makes completion # one-shot even if the same callback is accidentally invoked # twice with an old generation token. fence = await session.execute( update(ArtifactLibrarySyncStateRow) .where( ArtifactLibrarySyncStateRow.user_id == user_id, ArtifactLibrarySyncStateRow.generation == generation, ArtifactLibrarySyncStateRow.status == "syncing", ) .values(lease_expires_at=now + timedelta(seconds=300)) ) if fence.rowcount != 1: return False existing_result = await session.execute(select(ArtifactLibraryItemRow).where(ArtifactLibraryItemRow.user_id == user_id)) existing = {row.artifact_id: row for row in existing_result.scalars()} for artifact_id, item in deduped.items(): row = existing.get(artifact_id) if row is None: row = ArtifactLibraryItemRow( artifact_id=artifact_id, user_id=user_id, thread_id="", thread_title="", source_type="general_qa", path="", name="", size_bytes=0, modified_at=_EPOCH, file_kind="other", sync_generation=generation, indexed_at=now, ) session.add(row) _apply_item( row, user_id=user_id, generation=generation, item=item, indexed_at=now, ) await session.flush() await session.execute( delete(ArtifactLibraryItemRow).where( ArtifactLibraryItemRow.user_id == user_id, ArtifactLibraryItemRow.sync_generation != generation, ) ) await session.execute( update(ArtifactLibrarySyncStateRow) .where( ArtifactLibrarySyncStateRow.user_id == user_id, ArtifactLibrarySyncStateRow.generation == generation, ArtifactLibrarySyncStateRow.status == "syncing", ) .values( status="ready", completed_at=now, lease_expires_at=None, last_error=None, ) ) return True async def fail_sync(self, user_id: str, generation: str, error: str) -> bool: async with self._sf() as session: async with session.begin(): result = await session.execute( update(ArtifactLibrarySyncStateRow) .where( ArtifactLibrarySyncStateRow.user_id == user_id, ArtifactLibrarySyncStateRow.generation == generation, ArtifactLibrarySyncStateRow.status == "syncing", ) .values( status="error", lease_expires_at=None, last_error=(error or "Artifact library refresh failed")[:1024], ) ) return result.rowcount == 1 async def get_sync_state(self, user_id: str) -> dict[str, Any]: async with self._sf() as session: row = await session.get(ArtifactLibrarySyncStateRow, user_id) return _state_dict(row) async def list_items( self, user_id: str, *, query: str | None = None, file_kind: str = "all", source_type: str = "all", offset: int = 0, limit: int = 50, ) -> tuple[int, list[dict[str, Any]]]: filters = [ArtifactLibraryItemRow.user_id == user_id] if file_kind != "all": filters.append(ArtifactLibraryItemRow.file_kind == file_kind) if source_type != "all": filters.append(ArtifactLibraryItemRow.source_type == source_type) normalized = (query or "").strip().casefold() if normalized: searchable = ( ArtifactLibraryItemRow.name, ArtifactLibraryItemRow.thread_title, ArtifactLibraryItemRow.agent_name, ArtifactLibraryItemRow.strategy_name, ArtifactLibraryItemRow.strategy_project_id, ) filters.append(or_(*(func.lower(func.coalesce(column, "")).contains(normalized, autoescape=True) for column in searchable))) count_stmt = select(func.count()).select_from(ArtifactLibraryItemRow).where(*filters) page_stmt = ( select(ArtifactLibraryItemRow) .where(*filters) .order_by( ArtifactLibraryItemRow.modified_at.desc(), ArtifactLibraryItemRow.artifact_id.desc(), ) .offset(offset) .limit(limit) ) async with self._sf() as session: total = int((await session.execute(count_stmt)).scalar_one()) rows = (await session.execute(page_stmt)).scalars().all() return total, [_item_dict(row) for row in rows]