"""Workspace-scoped PostgreSQL persistence and durable projection work.""" import json from contextlib import contextmanager from uuid import uuid4 from sqlalchemy import create_engine, text from sqlalchemy.exc import IntegrityError, SQLAlchemyError from sqlalchemy.pool import NullPool from .migrate import expected_migrations from .models import Card, CardInput, CardQuery, MemoryConflict, MemoryNotFound, MemoryUnavailable class MemoryRepository: def __init__(self, database_url: str, workspace_id: str, *, engine=None, connection=None): self.workspace_id = workspace_id self.engine = engine or create_engine( database_url, poolclass=NullPool, connect_args={"connect_timeout": 5}, ) self.connection = connection def close(self): self.engine.dispose() @contextmanager def transaction(self): connection = self.connection try: connection = connection or self.engine.connect() with (connection.begin_nested() if connection.in_transaction() else connection.begin()): connection.execute(text("SET LOCAL ROLE thoth_memory_runtime")) connection.execute(text("SELECT set_config('thoth.memory_workspace', :w, true)"), {"w": self.workspace_id}) connection.execute(text("SET LOCAL statement_timeout = '15s'")) installed = dict(connection.execute(text( "SELECT version, checksum FROM thoth_memory.migrations" )).all()) if installed != expected_migrations(): raise MemoryUnavailable("Memory schema is incompatible; run installation migrations") yield connection except IntegrityError: raise MemoryConflict("Memory references conflict with the current archive") from None except SQLAlchemyError: raise MemoryUnavailable("Memory archive is unavailable; check its migrations and access") \ from None finally: if self.connection is None and connection is not None: connection.close() @contextmanager def operation(self): """Serialize each workspace across SQL commits and the bounded vector call.""" try: connection = self.engine.connect() connection.execute(text("SET statement_timeout = '15s'")) connection.execute(text("SELECT pg_advisory_lock(hashtextextended(:w, 792114204))"), {"w": self.workspace_id}) connection.commit() except SQLAlchemyError: if 'connection' in locals(): connection.close() raise MemoryUnavailable("Memory archive is busy or unavailable") from None try: yield MemoryRepository("", self.workspace_id, engine=self.engine, connection=connection) finally: # NullPool closes the physical connection, releasing the session advisory lock. connection.close() def _card(self, connection, row) -> Card: params = {"w": self.workspace_id, "id": row["id"]} links = connection.execute(text( "SELECT target_id, meaning FROM thoth_memory.links " "WHERE workspace_id=:w AND source_id=:id ORDER BY target_id" ), params).mappings().all() dependencies = connection.execute(text( 'SELECT database_id AS database, schema_name, table_name AS "table", ' 'column_name AS "column" FROM thoth_memory.dependencies ' "WHERE workspace_id=:w AND card_id=:id " "ORDER BY database_id, schema_name, table_name, column_name" ), params).mappings().all() return Card.model_validate({ **row["data"], "id": row["id"], "workspace_id": self.workspace_id, "family": row["family"], "subject": row["subject"], "origin": row["origin"], "created_at": row["created_at"], "updated_at": row["updated_at"], "revision": row["revision"], "indexed": row["indexed"], "links": [dict(v) for v in links], "dependencies": [dict(v) for v in dependencies], }) @staticmethod def _selection(): return ("SELECT c.*, (p.revision=c.revision AND NOT p.pending " "AND p.action='upsert' AND p.format=2) AS indexed FROM thoth_memory.cards c " "JOIN thoth_memory.projections p ON p.workspace_id=c.workspace_id " "AND p.card_id=c.id ") def get(self, card_id: str) -> Card: with self.transaction() as c: row = c.execute(text(self._selection()+"WHERE c.workspace_id=:w AND c.id=:id"), {"w": self.workspace_id, "id": card_id}).mappings().first() if row is None: raise MemoryNotFound("Memory card was not found in this workspace") return self._card(c, row) def list(self, query: CardQuery) -> dict: conditions = ["c.workspace_id=:w"] params = {"w": self.workspace_id} if query.q: conditions.append("(c.id ILIKE :q OR c.subject ILIKE :q OR c.data::text ILIKE :q)") params["q"] = "%" + query.q.replace("\\", "\\\\").replace("%", "\\%").replace("_", "\\_") + "%" for key in ("family", "origin"): if value := getattr(query, key): conditions.append(f"c.{key}=:{key}") params[key] = value if query.concept: conditions.append("c.data->'concepts' @> CAST(:concept AS jsonb)") params["concept"] = json.dumps([query.concept]) refs = [] for key, column in [("database", "database_id"), ("table", "table_name"), ("column", "column_name")]: if value := getattr(query, key): refs.append(f"d.{column}=:{key}") params[key] = value if refs: conditions.append("EXISTS (SELECT 1 FROM thoth_memory.dependencies d WHERE " "d.workspace_id=c.workspace_id AND d.card_id=c.id AND " + " AND ".join(refs) + ")") for key, comparison in [("updated_after", ">="), ("updated_before", "<=")]: if value := getattr(query, key): conditions.append(f"c.updated_at {comparison} :{key}") params[key] = value where = " WHERE " + " AND ".join(conditions) with self.transaction() as c: total = c.execute(text("SELECT count(*) FROM thoth_memory.cards c" + where), params).scalar_one() rows = c.execute(text(self._selection() + where + f" ORDER BY c.{query.sort} {query.direction}, c.id ASC LIMIT :limit OFFSET :offset"), {**params, "limit": query.page_size, "offset": (query.page - 1) * query.page_size}, ).mappings().all() return {"items": [self._card(c, row).model_dump(mode="json") for row in rows], "total": total, "page": query.page, "page_size": query.page_size} def exact_match(self, value: CardInput) -> Card | None: """Match authored content only; provenance and index state do not create new knowledge.""" data = value.model_dump(mode="json", exclude={"dependencies", "links"}) with self.transaction() as c: rows = c.execute(text(self._selection() + "WHERE c.workspace_id=:w AND c.family=:family AND c.subject=:subject " "AND c.data - 'session_id' - 'decision_seq'=CAST(:data AS jsonb) ORDER BY c.id"), {"w": self.workspace_id, "family": value.family, "subject": value.subject, "data": json.dumps(data)}).mappings() for row in rows: candidate = self._card(c, row) def ordered(values): return sorted(json.dumps(v.model_dump(), sort_keys=True) for v in values) if (ordered(candidate.dependencies) == ordered(value.dependencies) and ordered(candidate.links) == ordered(value.links)): return candidate return None def source(self, source_key: str): with self.transaction() as c: row = c.execute(text("SELECT card_id, action FROM thoth_memory.projections " "WHERE workspace_id=:w AND source_key=:s"), {"w": self.workspace_id, "s": source_key}).mappings().first() return dict(row) if row else None def save(self, value: CardInput, *, card_id: str | None = None, source_key: str | None = None, session_id: str | None = None, decision_seq: int | None = None, new_id: str | None = None) -> str: creating = card_id is None card_id = card_id or new_id or "mem-" + str(uuid4()) revision = str(uuid4()) data = value.model_dump(mode="json", exclude={"links", "dependencies"}) with self.transaction() as c: if not creating: old = c.execute(text("SELECT data FROM thoth_memory.cards " "WHERE workspace_id=:w AND id=:id"), {"w": self.workspace_id, "id": card_id}).scalar_one_or_none() if old is None: raise MemoryNotFound("Memory card was not found in this workspace") data.update({k: old.get(k) for k in ("session_id", "decision_seq")}) else: if c.execute(text("SELECT 1 FROM thoth_memory.projections " "WHERE workspace_id=:w AND card_id=:id"), {"w": self.workspace_id, "id": card_id}).first(): raise MemoryConflict("A proposed card identity was already used") data.update(session_id=session_id, decision_seq=decision_seq) params = {"w": self.workspace_id, "id": card_id, "r": revision, "data": json.dumps(data), "family": value.family, "subject": value.subject, "origin": "workflow" if source_key else "manual", "source": source_key} c.execute(text("INSERT INTO thoth_memory.cards " "(workspace_id,id,family,subject,origin,data,revision) " "VALUES (:w,:id,:family,:subject,:origin,CAST(:data AS jsonb),:r) " "ON CONFLICT (workspace_id,id) DO UPDATE SET family=EXCLUDED.family, " "subject=EXCLUDED.subject,data=EXCLUDED.data,revision=EXCLUDED.revision, " "updated_at=clock_timestamp()"), params) c.execute(text("DELETE FROM thoth_memory.links WHERE workspace_id=:w AND source_id=:id"), params) for link in value.links: c.execute(text("INSERT INTO thoth_memory.links VALUES (:w,:id,:target,:meaning)"), {**params, "target": link.target_id, "meaning": link.meaning}) c.execute(text("DELETE FROM thoth_memory.dependencies " "WHERE workspace_id=:w AND card_id=:id"), params) for dep in {tuple(d.model_dump().values()) for d in value.dependencies}: c.execute(text("INSERT INTO thoth_memory.dependencies VALUES " "(:w,:id,:database,:schema,:table,:column)"), {**params, **dict(zip(("database", "schema", "table", "column"), dep))}) c.execute(text("INSERT INTO thoth_memory.projections " "(workspace_id,card_id,revision,action,source_key) VALUES (:w,:id,:r,'upsert',:source) " "ON CONFLICT (workspace_id,card_id) DO UPDATE SET revision=EXCLUDED.revision, " "action='upsert',pending=true,error=NULL,updated_at=clock_timestamp()"), params) return card_id def delete(self, card_id: str): with self.transaction() as c: params = {"w": self.workspace_id, "id": card_id, "r": str(uuid4())} deleted = c.execute(text("DELETE FROM thoth_memory.cards " "WHERE workspace_id=:w AND id=:id"), params).rowcount if not deleted: raise MemoryNotFound("Memory card was not found in this workspace") c.execute(text("UPDATE thoth_memory.projections SET action='delete',revision=:r," "pending=true,error=NULL,updated_at=clock_timestamp() " "WHERE workspace_id=:w AND card_id=:id"), params) def projections(self, *, pending: bool = True): with self.transaction() as c: needs_update = "(pending OR (action='upsert' AND format<>2))" rows = c.execute(text("SELECT card_id,revision,action,"+needs_update+" AS pending," "error,updated_at " "FROM thoth_memory.projections WHERE workspace_id=:w " + ("AND "+needs_update+" " if pending else "") + "ORDER BY updated_at,card_id"), {"w": self.workspace_id}).mappings().all() return [dict(row) for row in rows] def projection_result(self, card_id: str, revision: str, error: str | None): with self.transaction() as c: c.execute(text("UPDATE thoth_memory.projections SET pending=:p,error=:error,format=2," "updated_at=clock_timestamp() WHERE workspace_id=:w AND card_id=:id AND revision=:r"), {"p": error is not None, "error": error, "w": self.workspace_id, "id": card_id, "r": revision}) def invalidate_all(self): with self.transaction() as c: c.execute(text("UPDATE thoth_memory.projections SET pending=true,error=NULL " "WHERE workspace_id=:w"), {"w": self.workspace_id})