mirror of
https://github.com/tiennm99/DocsGPT.git
synced 2026-10-13 04:14:42 +00:00
feat(wiki): wiki_pages table + repository + targeted chunk delete
Migration 0023_wiki_pages (source-scoped; version/content_hash/embed_status; FK source_id->sources ON DELETE CASCADE; UNIQUE(source_id,path) + prefix index). WikiPagesRepository: path-keyed CRUD with content-hash short-circuit, version bump, and move-with-reject. BaseVectorStore.delete_chunks_by_source_path (loop default) + parameterized pgvector override (DELETE WHERE source_id=%s AND metadata->>'source'=%s). Unit 1 of F-Wiki.
This commit is contained in:
1 parent
a709b7ccdc
commit
8aa0facf4a
8 files changed
+576
No files matched your search
@@ -0,0 +1,64 @@
|
||||
"""0023 wiki pages — authoritative storage for the LLM-editable wiki source.
|
||||
|
||||
Adds ``wiki_pages`` (page = unit), source-scoped (team-shareable) and the
|
||||
authoritative store for a ``config.kind="wiki"`` source. The vector store is a
|
||||
derived index re-embedded asynchronously per page; ``embed_status`` tracks that
|
||||
freshness. ``UNIQUE(source_id, path)`` makes a page addressable like a file;
|
||||
the ``text_pattern_ops`` index backs prefix listing of the virtual tree.
|
||||
|
||||
No new Postgres extensions (keeps it Neon/DBngin-portable).
|
||||
|
||||
Revision ID: 0023_wiki_pages
|
||||
Revises: 0022_source_config
|
||||
"""
|
||||
|
||||
from typing import Sequence, Union
|
||||
|
||||
from alembic import op
|
||||
|
||||
|
||||
revision: str = "0023_wiki_pages"
|
||||
down_revision: Union[str, None] = "0022_source_config"
|
||||
branch_labels: Union[str, Sequence[str], None] = None
|
||||
depends_on: Union[str, Sequence[str], None] = None
|
||||
|
||||
|
||||
def upgrade() -> None:
|
||||
op.execute(
|
||||
"""
|
||||
CREATE TABLE IF NOT EXISTS wiki_pages (
|
||||
id UUID PRIMARY KEY DEFAULT gen_random_uuid(),
|
||||
source_id UUID NOT NULL REFERENCES sources(id) ON DELETE CASCADE,
|
||||
path TEXT NOT NULL,
|
||||
title TEXT,
|
||||
content TEXT NOT NULL,
|
||||
token_count INTEGER,
|
||||
version INTEGER NOT NULL DEFAULT 1,
|
||||
content_hash TEXT,
|
||||
embed_status TEXT NOT NULL DEFAULT 'pending',
|
||||
updated_by TEXT,
|
||||
created_at TIMESTAMPTZ NOT NULL DEFAULT now(),
|
||||
updated_at TIMESTAMPTZ NOT NULL DEFAULT now()
|
||||
);
|
||||
"""
|
||||
)
|
||||
op.execute(
|
||||
"CREATE UNIQUE INDEX IF NOT EXISTS wiki_pages_source_path_uidx "
|
||||
"ON wiki_pages (source_id, path);"
|
||||
)
|
||||
op.execute(
|
||||
"CREATE INDEX IF NOT EXISTS wiki_pages_source_path_prefix_idx "
|
||||
"ON wiki_pages (source_id, path text_pattern_ops);"
|
||||
)
|
||||
op.execute("DROP TRIGGER IF EXISTS wiki_pages_set_updated_at ON wiki_pages;")
|
||||
op.execute(
|
||||
"""
|
||||
CREATE TRIGGER wiki_pages_set_updated_at
|
||||
BEFORE UPDATE ON wiki_pages
|
||||
FOR EACH ROW EXECUTE FUNCTION set_updated_at();
|
||||
"""
|
||||
)
|
||||
|
||||
|
||||
def downgrade() -> None:
|
||||
op.execute("DROP TABLE IF EXISTS wiki_pages;")
|
||||
@@ -415,6 +415,34 @@ memories_table = Table(
|
||||
UniqueConstraint("user_id", "tool_id", "path", name="memories_user_tool_path_uidx"),
|
||||
)
|
||||
|
||||
# Authoritative storage for an LLM-editable wiki source (config.kind="wiki").
|
||||
# Source-scoped (team-shareable) unlike per-user ``memories``; the vector store
|
||||
# is a derived index re-embedded per page (``embed_status`` tracks freshness).
|
||||
wiki_pages_table = Table(
|
||||
"wiki_pages",
|
||||
metadata,
|
||||
Column("id", UUID(as_uuid=True), primary_key=True, server_default=func.gen_random_uuid()),
|
||||
Column("source_id", UUID(as_uuid=True), ForeignKey("sources.id", ondelete="CASCADE"), nullable=False),
|
||||
Column("path", Text, nullable=False),
|
||||
Column("title", Text),
|
||||
Column("content", Text, nullable=False),
|
||||
Column("token_count", Integer),
|
||||
Column("version", Integer, nullable=False, server_default="1"),
|
||||
Column("content_hash", Text),
|
||||
Column("embed_status", Text, nullable=False, server_default="pending"),
|
||||
Column("updated_by", Text),
|
||||
Column("created_at", DateTime(timezone=True), nullable=False, server_default=func.now()),
|
||||
Column("updated_at", DateTime(timezone=True), nullable=False, server_default=func.now()),
|
||||
UniqueConstraint("source_id", "path", name="wiki_pages_source_path_uidx"),
|
||||
)
|
||||
|
||||
Index(
|
||||
"wiki_pages_source_path_prefix_idx",
|
||||
wiki_pages_table.c.source_id,
|
||||
wiki_pages_table.c.path,
|
||||
postgresql_ops={"path": "text_pattern_ops"},
|
||||
)
|
||||
|
||||
todos_table = Table(
|
||||
"todos",
|
||||
metadata,
|
||||
|
||||
@@ -0,0 +1,157 @@
|
||||
"""Repository for the ``wiki_pages`` table.
|
||||
|
||||
Authoritative, source-scoped storage for an LLM-editable wiki source. Mirrors
|
||||
``MemoriesRepository`` but is keyed on ``source_id`` (not ``user_id``/``tool_id``)
|
||||
and adds wiki-specific semantics: ``content_hash`` short-circuit on identical
|
||||
writes, a monotonically increasing ``version``, and an ``embed_status`` setter
|
||||
driving the async per-page re-embed.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import hashlib
|
||||
from typing import Optional
|
||||
|
||||
from sqlalchemy import Connection, text
|
||||
|
||||
from application.storage.db.base_repository import row_to_dict
|
||||
from application.utils import num_tokens_from_string
|
||||
|
||||
|
||||
def _content_hash(content: str) -> str:
|
||||
return hashlib.sha256(content.encode("utf-8")).hexdigest()
|
||||
|
||||
|
||||
class WikiPagesRepository:
|
||||
def __init__(self, conn: Connection) -> None:
|
||||
self._conn = conn
|
||||
|
||||
def upsert(
|
||||
self,
|
||||
source_id: str,
|
||||
path: str,
|
||||
content: str,
|
||||
title: Optional[str] = None,
|
||||
updated_by: Optional[str] = None,
|
||||
) -> dict:
|
||||
"""Create or overwrite a page.
|
||||
|
||||
Short-circuits (no write, returns the row unchanged) when an existing
|
||||
page already has the same ``content_hash``. Otherwise inserts, or on
|
||||
conflict updates content/title/hash/updated_by, bumps ``version``, and
|
||||
resets ``embed_status`` to ``pending``.
|
||||
"""
|
||||
content_hash = _content_hash(content)
|
||||
|
||||
existing = self.get_by_path(source_id, path)
|
||||
if existing is not None and existing.get("content_hash") == content_hash:
|
||||
return existing
|
||||
|
||||
result = self._conn.execute(
|
||||
text(
|
||||
"""
|
||||
INSERT INTO wiki_pages
|
||||
(source_id, path, title, content, token_count,
|
||||
content_hash, updated_by, embed_status)
|
||||
VALUES
|
||||
(CAST(:source_id AS uuid), :path, :title, :content, :token_count,
|
||||
:content_hash, :updated_by, 'pending')
|
||||
ON CONFLICT (source_id, path)
|
||||
DO UPDATE SET
|
||||
title = EXCLUDED.title,
|
||||
content = EXCLUDED.content,
|
||||
token_count = EXCLUDED.token_count,
|
||||
content_hash = EXCLUDED.content_hash,
|
||||
updated_by = EXCLUDED.updated_by,
|
||||
version = wiki_pages.version + 1,
|
||||
embed_status = 'pending',
|
||||
updated_at = now()
|
||||
RETURNING *
|
||||
"""
|
||||
),
|
||||
{
|
||||
"source_id": source_id,
|
||||
"path": path,
|
||||
"title": title,
|
||||
"content": content,
|
||||
"token_count": num_tokens_from_string(content),
|
||||
"content_hash": content_hash,
|
||||
"updated_by": updated_by,
|
||||
},
|
||||
)
|
||||
return row_to_dict(result.fetchone())
|
||||
|
||||
def get_by_path(self, source_id: str, path: str) -> Optional[dict]:
|
||||
result = self._conn.execute(
|
||||
text(
|
||||
"SELECT * FROM wiki_pages "
|
||||
"WHERE source_id = CAST(:source_id AS uuid) AND path = :path"
|
||||
),
|
||||
{"source_id": source_id, "path": path},
|
||||
)
|
||||
row = result.fetchone()
|
||||
return row_to_dict(row) if row is not None else None
|
||||
|
||||
def list_by_prefix(self, source_id: str, prefix: str) -> list[dict]:
|
||||
result = self._conn.execute(
|
||||
text(
|
||||
"SELECT * FROM wiki_pages "
|
||||
"WHERE source_id = CAST(:source_id AS uuid) AND path LIKE :prefix "
|
||||
"ORDER BY path"
|
||||
),
|
||||
{"source_id": source_id, "prefix": prefix + "%"},
|
||||
)
|
||||
return [row_to_dict(r) for r in result.fetchall()]
|
||||
|
||||
def list_for_source(self, source_id: str) -> list[dict]:
|
||||
result = self._conn.execute(
|
||||
text(
|
||||
"SELECT * FROM wiki_pages "
|
||||
"WHERE source_id = CAST(:source_id AS uuid) ORDER BY path"
|
||||
),
|
||||
{"source_id": source_id},
|
||||
)
|
||||
return [row_to_dict(r) for r in result.fetchall()]
|
||||
|
||||
def delete_by_path(self, source_id: str, path: str) -> int:
|
||||
result = self._conn.execute(
|
||||
text(
|
||||
"DELETE FROM wiki_pages "
|
||||
"WHERE source_id = CAST(:source_id AS uuid) AND path = :path"
|
||||
),
|
||||
{"source_id": source_id, "path": path},
|
||||
)
|
||||
return result.rowcount
|
||||
|
||||
def delete_by_prefix(self, source_id: str, prefix: str) -> int:
|
||||
result = self._conn.execute(
|
||||
text(
|
||||
"DELETE FROM wiki_pages "
|
||||
"WHERE source_id = CAST(:source_id AS uuid) AND path LIKE :prefix"
|
||||
),
|
||||
{"source_id": source_id, "prefix": prefix + "%"},
|
||||
)
|
||||
return result.rowcount
|
||||
|
||||
def update_path(self, source_id: str, old_path: str, new_path: str) -> bool:
|
||||
"""Move a page. Rejects (returns ``False``) if ``new_path`` already exists."""
|
||||
if self.get_by_path(source_id, new_path) is not None:
|
||||
return False
|
||||
result = self._conn.execute(
|
||||
text(
|
||||
"UPDATE wiki_pages SET path = :new_path, updated_at = now() "
|
||||
"WHERE source_id = CAST(:source_id AS uuid) AND path = :old_path"
|
||||
),
|
||||
{"source_id": source_id, "old_path": old_path, "new_path": new_path},
|
||||
)
|
||||
return result.rowcount > 0
|
||||
|
||||
def set_embed_status(self, source_id: str, path: str, status: str) -> bool:
|
||||
result = self._conn.execute(
|
||||
text(
|
||||
"UPDATE wiki_pages SET embed_status = :status "
|
||||
"WHERE source_id = CAST(:source_id AS uuid) AND path = :path"
|
||||
),
|
||||
{"source_id": source_id, "path": path, "status": status},
|
||||
)
|
||||
return result.rowcount > 0
|
||||
@@ -169,6 +169,21 @@ class BaseVectorStore(ABC):
|
||||
"""Delete a specific chunk from the vectorstore"""
|
||||
pass
|
||||
|
||||
def delete_chunks_by_source_path(self, path) -> int:
|
||||
"""Delete every chunk whose ``metadata.source`` equals ``path``.
|
||||
|
||||
Default implementation iterates ``get_chunks()`` and deletes the
|
||||
matches via ``delete_chunk()`` — works for any store. Override with a
|
||||
single targeted statement where the store supports it. Returns the
|
||||
number of chunks deleted.
|
||||
"""
|
||||
deleted = 0
|
||||
for chunk in self.get_chunks() or []:
|
||||
if (chunk.get("metadata") or {}).get("source") == path:
|
||||
if self.delete_chunk(chunk.get("doc_id")):
|
||||
deleted += 1
|
||||
return deleted
|
||||
|
||||
def is_azure_configured(self):
|
||||
return (
|
||||
settings.OPENAI_API_BASE
|
||||
|
||||
@@ -372,6 +372,34 @@ class PGVectorStore(BaseVectorStore):
|
||||
finally:
|
||||
cursor.close()
|
||||
|
||||
def delete_chunks_by_source_path(self, path: str) -> int:
|
||||
"""Delete this source's chunks whose ``metadata.source`` equals ``path``.
|
||||
|
||||
One targeted statement instead of the base loop+scan. The path is bound
|
||||
as a query parameter (never interpolated); only the internal table name
|
||||
is f-string interpolated. Returns the number of rows deleted.
|
||||
"""
|
||||
conn = self._get_connection()
|
||||
cursor = conn.cursor()
|
||||
|
||||
try:
|
||||
delete_query = (
|
||||
f"DELETE FROM {self._table_name} "
|
||||
f"WHERE source_id = %s AND {self._metadata_column}->>'source' = %s;"
|
||||
)
|
||||
cursor.execute(delete_query, (self._source_id, path))
|
||||
deleted_count = cursor.rowcount
|
||||
conn.commit()
|
||||
|
||||
return deleted_count
|
||||
|
||||
except Exception as e:
|
||||
conn.rollback()
|
||||
logging.error(f"Error deleting chunks by source path: {e}")
|
||||
raise
|
||||
finally:
|
||||
cursor.close()
|
||||
|
||||
def __del__(self):
|
||||
"""Close database connection when object is destroyed"""
|
||||
if hasattr(self, '_connection') and self._connection and not self._connection.closed:
|
||||
|
||||
@@ -0,0 +1,172 @@
|
||||
"""Tests for WikiPagesRepository against a real Postgres instance.
|
||||
|
||||
Wiki pages have a FK to ``sources`` (ON DELETE CASCADE), so each test creates a
|
||||
source row first.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from application.storage.db.repositories.sources import SourcesRepository
|
||||
from application.storage.db.repositories.wiki_pages import WikiPagesRepository
|
||||
|
||||
|
||||
def _repo(conn) -> WikiPagesRepository:
|
||||
return WikiPagesRepository(conn)
|
||||
|
||||
|
||||
def _make_source(conn, user_id: str = "wiki-user", name: str = "wiki-src") -> str:
|
||||
"""Insert a sources row and return its UUID as a string."""
|
||||
return str(SourcesRepository(conn).create(name, user_id=user_id, type="wiki")["id"])
|
||||
|
||||
|
||||
class TestUpsert:
|
||||
def test_creates_page(self, pg_conn):
|
||||
repo = _repo(pg_conn)
|
||||
sid = _make_source(pg_conn)
|
||||
page = repo.upsert(sid, "/docs/readme.md", "Hello world", title="Readme")
|
||||
assert page["path"] == "/docs/readme.md"
|
||||
assert page["content"] == "Hello world"
|
||||
assert page["title"] == "Readme"
|
||||
assert page["version"] == 1
|
||||
assert page["embed_status"] == "pending"
|
||||
assert page["content_hash"] is not None
|
||||
assert page["token_count"] is not None and page["token_count"] > 0
|
||||
|
||||
def test_get_by_path_after_upsert(self, pg_conn):
|
||||
repo = _repo(pg_conn)
|
||||
sid = _make_source(pg_conn)
|
||||
repo.upsert(sid, "/a.md", "content")
|
||||
got = repo.get_by_path(sid, "/a.md")
|
||||
assert got is not None
|
||||
assert got["content"] == "content"
|
||||
|
||||
def test_records_updated_by(self, pg_conn):
|
||||
repo = _repo(pg_conn)
|
||||
sid = _make_source(pg_conn)
|
||||
page = repo.upsert(sid, "/a.md", "x", updated_by="agent-1")
|
||||
assert page["updated_by"] == "agent-1"
|
||||
|
||||
|
||||
class TestContentHashShortCircuit:
|
||||
def test_identical_content_does_not_bump_version(self, pg_conn):
|
||||
repo = _repo(pg_conn)
|
||||
sid = _make_source(pg_conn)
|
||||
first = repo.upsert(sid, "/a.md", "same")
|
||||
second = repo.upsert(sid, "/a.md", "same")
|
||||
assert first["id"] == second["id"]
|
||||
assert second["version"] == 1
|
||||
assert first["content_hash"] == second["content_hash"]
|
||||
|
||||
def test_identical_content_resets_status_unchanged(self, pg_conn):
|
||||
repo = _repo(pg_conn)
|
||||
sid = _make_source(pg_conn)
|
||||
repo.upsert(sid, "/a.md", "same")
|
||||
repo.set_embed_status(sid, "/a.md", "embedded")
|
||||
# Re-upsert with identical content must not reset embed_status to pending.
|
||||
again = repo.upsert(sid, "/a.md", "same")
|
||||
assert again["embed_status"] == "embedded"
|
||||
|
||||
|
||||
class TestVersionBump:
|
||||
def test_changed_content_bumps_version_and_resets_status(self, pg_conn):
|
||||
repo = _repo(pg_conn)
|
||||
sid = _make_source(pg_conn)
|
||||
v1 = repo.upsert(sid, "/a.md", "v1")
|
||||
repo.set_embed_status(sid, "/a.md", "embedded")
|
||||
v2 = repo.upsert(sid, "/a.md", "v2")
|
||||
assert v2["version"] == 2
|
||||
assert v2["content"] == "v2"
|
||||
assert v2["embed_status"] == "pending"
|
||||
assert v2["content_hash"] != v1["content_hash"]
|
||||
|
||||
|
||||
class TestListByPrefix:
|
||||
def test_lists_matching_prefix(self, pg_conn):
|
||||
repo = _repo(pg_conn)
|
||||
sid = _make_source(pg_conn)
|
||||
repo.upsert(sid, "/docs/a.md", "a")
|
||||
repo.upsert(sid, "/docs/b.md", "b")
|
||||
repo.upsert(sid, "/other/c.md", "c")
|
||||
results = repo.list_by_prefix(sid, "/docs/")
|
||||
assert {r["path"] for r in results} == {"/docs/a.md", "/docs/b.md"}
|
||||
|
||||
def test_list_for_source_returns_all(self, pg_conn):
|
||||
repo = _repo(pg_conn)
|
||||
sid = _make_source(pg_conn)
|
||||
repo.upsert(sid, "/a.md", "a")
|
||||
repo.upsert(sid, "/b.md", "b")
|
||||
results = repo.list_for_source(sid)
|
||||
assert {r["path"] for r in results} == {"/a.md", "/b.md"}
|
||||
|
||||
def test_list_is_scoped_to_source(self, pg_conn):
|
||||
repo = _repo(pg_conn)
|
||||
sid1 = _make_source(pg_conn, name="src-1")
|
||||
sid2 = _make_source(pg_conn, name="src-2")
|
||||
repo.upsert(sid1, "/a.md", "a")
|
||||
repo.upsert(sid2, "/a.md", "b")
|
||||
assert len(repo.list_for_source(sid1)) == 1
|
||||
assert repo.get_by_path(sid1, "/a.md")["content"] == "a"
|
||||
assert repo.get_by_path(sid2, "/a.md")["content"] == "b"
|
||||
|
||||
|
||||
class TestUpdatePath:
|
||||
def test_renames_page(self, pg_conn):
|
||||
repo = _repo(pg_conn)
|
||||
sid = _make_source(pg_conn)
|
||||
repo.upsert(sid, "/old.md", "content")
|
||||
moved = repo.update_path(sid, "/old.md", "/new.md")
|
||||
assert moved is True
|
||||
assert repo.get_by_path(sid, "/old.md") is None
|
||||
assert repo.get_by_path(sid, "/new.md")["content"] == "content"
|
||||
|
||||
def test_rename_nonexistent_returns_false(self, pg_conn):
|
||||
repo = _repo(pg_conn)
|
||||
sid = _make_source(pg_conn)
|
||||
assert repo.update_path(sid, "/nope.md", "/new.md") is False
|
||||
|
||||
def test_rename_rejects_existing_target(self, pg_conn):
|
||||
repo = _repo(pg_conn)
|
||||
sid = _make_source(pg_conn)
|
||||
repo.upsert(sid, "/a.md", "a")
|
||||
repo.upsert(sid, "/b.md", "b")
|
||||
assert repo.update_path(sid, "/a.md", "/b.md") is False
|
||||
# Both pages survive untouched.
|
||||
assert repo.get_by_path(sid, "/a.md")["content"] == "a"
|
||||
assert repo.get_by_path(sid, "/b.md")["content"] == "b"
|
||||
|
||||
|
||||
class TestDelete:
|
||||
def test_delete_by_path(self, pg_conn):
|
||||
repo = _repo(pg_conn)
|
||||
sid = _make_source(pg_conn)
|
||||
repo.upsert(sid, "/x.md", "c")
|
||||
assert repo.delete_by_path(sid, "/x.md") == 1
|
||||
assert repo.get_by_path(sid, "/x.md") is None
|
||||
|
||||
def test_delete_by_path_nonexistent_returns_zero(self, pg_conn):
|
||||
repo = _repo(pg_conn)
|
||||
sid = _make_source(pg_conn)
|
||||
assert repo.delete_by_path(sid, "/nope.md") == 0
|
||||
|
||||
def test_delete_by_prefix(self, pg_conn):
|
||||
repo = _repo(pg_conn)
|
||||
sid = _make_source(pg_conn)
|
||||
repo.upsert(sid, "/dir/a.md", "a")
|
||||
repo.upsert(sid, "/dir/b.md", "b")
|
||||
repo.upsert(sid, "/other/c.md", "c")
|
||||
assert repo.delete_by_prefix(sid, "/dir/") == 2
|
||||
assert repo.get_by_path(sid, "/other/c.md") is not None
|
||||
|
||||
|
||||
class TestSetEmbedStatus:
|
||||
def test_set_embed_status(self, pg_conn):
|
||||
repo = _repo(pg_conn)
|
||||
sid = _make_source(pg_conn)
|
||||
repo.upsert(sid, "/a.md", "a")
|
||||
assert repo.set_embed_status(sid, "/a.md", "embedded") is True
|
||||
assert repo.get_by_path(sid, "/a.md")["embed_status"] == "embedded"
|
||||
|
||||
def test_set_embed_status_nonexistent_returns_false(self, pg_conn):
|
||||
repo = _repo(pg_conn)
|
||||
sid = _make_source(pg_conn)
|
||||
assert repo.set_embed_status(sid, "/nope.md", "embedded") is False
|
||||
@@ -0,0 +1,77 @@
|
||||
"""Migration round-trip test for 0023_wiki_pages."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import os
|
||||
import subprocess
|
||||
import sys
|
||||
from pathlib import Path
|
||||
|
||||
import pytest
|
||||
from sqlalchemy import text
|
||||
|
||||
|
||||
pytestmark = pytest.mark.integration
|
||||
|
||||
|
||||
def _alembic_ini() -> Path:
|
||||
return Path(__file__).resolve().parents[3] / "application" / "alembic.ini"
|
||||
|
||||
|
||||
def _run_alembic(url: str, *args: str) -> None:
|
||||
subprocess.check_call(
|
||||
[sys.executable, "-m", "alembic", "-c", str(_alembic_ini()), *args],
|
||||
timeout=60,
|
||||
env={**os.environ, "POSTGRES_URI": url},
|
||||
)
|
||||
|
||||
|
||||
def _alembic_heads(url: str) -> list[str]:
|
||||
out = subprocess.check_output(
|
||||
[sys.executable, "-m", "alembic", "-c", str(_alembic_ini()), "heads"],
|
||||
timeout=60,
|
||||
env={**os.environ, "POSTGRES_URI": url},
|
||||
text=True,
|
||||
)
|
||||
return [line for line in out.splitlines() if line.strip()]
|
||||
|
||||
|
||||
def _alembic_version(conn) -> str:
|
||||
return conn.execute(text("SELECT version_num FROM alembic_version")).scalar()
|
||||
|
||||
|
||||
def _table_exists(conn, table: str) -> bool:
|
||||
row = conn.execute(
|
||||
text(
|
||||
"SELECT 1 FROM information_schema.tables "
|
||||
"WHERE table_name = :t AND table_schema = 'public'"
|
||||
),
|
||||
{"t": table},
|
||||
).fetchone()
|
||||
return row is not None
|
||||
|
||||
|
||||
_0023 = "0023_wiki_pages"
|
||||
_0022 = "0022_source_config"
|
||||
|
||||
|
||||
class TestMigration0023RoundTrip:
|
||||
def test_single_head(self, pg_engine):
|
||||
url = pg_engine.url.render_as_string(hide_password=False)
|
||||
assert len(_alembic_heads(url)) == 1
|
||||
|
||||
def test_head_has_wiki_pages_table(self, pg_engine):
|
||||
with pg_engine.connect() as conn:
|
||||
assert _alembic_version(conn) >= _0023
|
||||
assert _table_exists(conn, "wiki_pages")
|
||||
|
||||
def test_downgrade_drops_then_upgrade_restores(self, pg_engine):
|
||||
url = pg_engine.url.render_as_string(hide_password=False)
|
||||
_run_alembic(url, "downgrade", _0022)
|
||||
with pg_engine.connect() as conn:
|
||||
assert _alembic_version(conn) == _0022
|
||||
assert not _table_exists(conn, "wiki_pages")
|
||||
_run_alembic(url, "upgrade", "head")
|
||||
with pg_engine.connect() as conn:
|
||||
assert _alembic_version(conn) >= _0023
|
||||
assert _table_exists(conn, "wiki_pages")
|
||||
@@ -331,6 +331,41 @@ class TestPGVectorStoreDeleteChunk:
|
||||
assert result is False
|
||||
|
||||
|
||||
@pytest.mark.unit
|
||||
class TestPGVectorStoreDeleteChunksBySourcePath:
|
||||
def test_targeted_delete_is_parameterized(self):
|
||||
store, mock_conn, mock_cursor, _ = _make_store(source_id="src1")
|
||||
mock_cursor.rowcount = 3
|
||||
|
||||
deleted = store.delete_chunks_by_source_path("/docs/page.md")
|
||||
|
||||
assert deleted == 3
|
||||
sql, params = mock_cursor.execute.call_args[0]
|
||||
# Single targeted DELETE; the path is a bound param, never interpolated.
|
||||
assert "/docs/page.md" not in sql
|
||||
assert "DELETE FROM" in sql
|
||||
assert "metadata->>'source' = %s" in sql
|
||||
assert "source_id = %s" in sql
|
||||
assert params == ("src1", "/docs/page.md")
|
||||
mock_conn.commit.assert_called_once()
|
||||
|
||||
def test_returns_zero_when_no_match(self):
|
||||
store, mock_conn, mock_cursor, _ = _make_store(source_id="src1")
|
||||
mock_cursor.rowcount = 0
|
||||
|
||||
assert store.delete_chunks_by_source_path("/missing.md") == 0
|
||||
mock_conn.commit.assert_called_once()
|
||||
|
||||
def test_rolls_back_and_raises_on_error(self):
|
||||
store, mock_conn, mock_cursor, _ = _make_store()
|
||||
mock_cursor.execute.side_effect = Exception("delete failed")
|
||||
|
||||
with pytest.raises(Exception, match="delete failed"):
|
||||
store.delete_chunks_by_source_path("/x.md")
|
||||
|
||||
mock_conn.rollback.assert_called_once()
|
||||
|
||||
|
||||
@pytest.mark.unit
|
||||
class TestPGVectorStoreConnection:
|
||||
def test_get_connection_creates_new_when_closed(self):
|
||||
|
||||
Reference in new issue
Block a user