Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
Original file line number Diff line number Diff line change
Expand Up @@ -38,6 +38,10 @@ def rebuild_sqlite_search_index(*, include_script_ngrams: bool) -> None:
return

script_definition = "script_ngrams," if include_script_ngrams else ""
# A killed mid-migration run can leave an empty search_index_rebuilt while
# alembic_version stays behind (issue #1853). Drop the leftover so a retry
# can CREATE again instead of failing with "table already exists".
op.execute("DROP TABLE IF EXISTS search_index_rebuilt")
op.execute(f"""
CREATE VIRTUAL TABLE search_index_rebuilt USING fts5(
id UNINDEXED,
Expand Down
75 changes: 44 additions & 31 deletions src/basic_memory/repository/sqlite_vec_index.py
Original file line number Diff line number Diff line change
Expand Up @@ -252,45 +252,58 @@ async def upsert(self, project_id: int, records: Sequence[VectorRecord]) -> None
)
await session.commit()

# SQLite caps expression-tree depth at SQLITE_MAX_EXPR_DEPTH (1000). Each
# pending-row predicate is a 4-term AND under an OR chain, so a single
# statement fails above ~995 deletions (#1850). Batch below that ceiling.
_DELETE_BATCH_SIZE = 200

async def delete(self, project_id: int, records: Sequence[VectorDeletion]) -> None:
if not records:
return
await self.initialize()
async with db.scoped_session(self._session_maker) as session:
await self._ensure_loaded(session)
params: dict[str, object] = {"project_id": project_id}
predicates: list[str] = []
for index, record in enumerate(records):
params[f"entity_id_{index}"] = record.key.entity_id
params[f"chunk_key_{index}"] = record.key.chunk_key
params[f"source_hash_{index}"] = record.source_hash
predicates.append(
f"(entity_id = :entity_id_{index} "
f"AND chunk_key = :chunk_key_{index} "
f"AND source_hash = :source_hash_{index} "
"AND embedding_status = 'pending')"
)
result = await session.execute(
text(
"UPDATE search_vector_chunks SET source_hash = source_hash "
"WHERE project_id = :project_id AND ("
+ " OR ".join(predicates)
+ ") RETURNING id"
),
params,
)
rowids = [int(row_id) for row_id in result.scalars().all()]
if rowids:
params = {f"rowid_{index}": rowid for index, rowid in enumerate(rowids)}
placeholders = ", ".join(f":rowid_{index}" for index in range(len(rowids)))
await session.execute(
text(f"DELETE FROM search_vector_embeddings WHERE rowid IN ({placeholders})"),
params,
)
await session.execute(
text(f"DELETE FROM search_vector_chunks WHERE id IN ({placeholders})"),
all_rowids: list[int] = []
for start in range(0, len(records), self._DELETE_BATCH_SIZE):
batch = records[start : start + self._DELETE_BATCH_SIZE]
params: dict[str, object] = {"project_id": project_id}
predicates: list[str] = []
for index, record in enumerate(batch):
params[f"entity_id_{index}"] = record.key.entity_id
params[f"chunk_key_{index}"] = record.key.chunk_key
params[f"source_hash_{index}"] = record.source_hash
predicates.append(
f"(entity_id = :entity_id_{index} "
f"AND chunk_key = :chunk_key_{index} "
f"AND source_hash = :source_hash_{index} "
"AND embedding_status = 'pending')"
)
result = await session.execute(
text(
"UPDATE search_vector_chunks SET source_hash = source_hash "
"WHERE project_id = :project_id AND ("
+ " OR ".join(predicates)
+ ") RETURNING id"
),
params,
)
all_rowids.extend(int(row_id) for row_id in result.scalars().all())
if all_rowids:
# IN-lists are shallow; still batch to keep statement size bounded.
for start in range(0, len(all_rowids), self._DELETE_BATCH_SIZE):
batch_ids = all_rowids[start : start + self._DELETE_BATCH_SIZE]
params = {f"rowid_{index}": rowid for index, rowid in enumerate(batch_ids)}
placeholders = ", ".join(f":rowid_{index}" for index in range(len(batch_ids)))
await session.execute(
text(
f"DELETE FROM search_vector_embeddings WHERE rowid IN ({placeholders})"
),
params,
)
await session.execute(
text(f"DELETE FROM search_vector_chunks WHERE id IN ({placeholders})"),
params,
)
await session.commit()

async def delete_entity(self, project_id: int, entity_id: int) -> None:
Expand Down
147 changes: 147 additions & 0 deletions test-int/test_sqlite_index_recovery.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,147 @@
"""SQLite index recovery against real files and sqlite-vec, without application mocks."""

import json
from importlib import import_module
from pathlib import Path

import pytest
from alembic.migration import MigrationContext
from alembic.operations import Operations
from sqlalchemy import create_engine, text

from basic_memory import db
from basic_memory.repository.semantic_vector_index import VectorDeletion, VectorKey
from basic_memory.repository.sqlite_search_repository import SQLiteSearchRepository
from basic_memory.repository.sqlite_vec_index import SQLiteVecIndex
from basic_memory.services.search_service import SearchService


pytestmark = pytest.mark.real_embedder


@pytest.mark.asyncio
async def test_large_vector_delete_preserves_other_generations_and_projects(
search_service: SearchService,
) -> None:
repository = search_service.repository
if not isinstance(repository, SQLiteSearchRepository):
pytest.skip("SQLite expression-depth limits are backend-specific.")
index = repository._semantic_vector_index
assert isinstance(index, SQLiteVecIndex)
await index.initialize()
project_id = repository.project_id
embedding = json.dumps([1.0] + [0.0] * (index.scope.dimensions - 1))
count = 1201
# Put protected rows at batch boundaries as well as in the final partial batch.
ready = {199, 400}
newer = {200, 1200}
other_project = {399}
protected = ready | newer | other_project
async with db.scoped_session(repository.session_maker) as session:
await index._ensure_loaded(session)
await session.execute(
text(
"INSERT INTO search_vector_chunks "
"(id, entity_id, project_id, chunk_key, source_type, source_row_id, "
"chunk_index, chunk_text, source_hash, entity_fingerprint, embedding_model, "
"vector_index, embedding_status) VALUES "
"(:id, 1850, :project_id, :chunk_key, 'entity', 1850, :id, 'text', "
":source_hash, 'fingerprint', :model, 'sqlite-vec', :status)"
),
[
{
"id": i + 1,
"project_id": project_id + 1 if i in other_project else project_id,
"chunk_key": f"chunk:{i}",
"source_hash": "new" if i in newer else "old",
"model": index.scope.embedding_identity,
"status": "ready" if i in ready else "pending",
}
for i in range(count)
],
)
await session.execute(
text(
"INSERT INTO search_vector_embeddings "
"(rowid, project_id, embedding, source_hash) "
"VALUES (:id, :project_id, :embedding, :source_hash)"
),
[
{
"id": i + 1,
"project_id": project_id + 1 if i in other_project else project_id,
"embedding": embedding,
"source_hash": "new" if i in newer else "old",
}
for i in range(count)
],
)
await session.commit()

deletions = [
VectorDeletion(key=VectorKey(entity_id=1850, chunk_key=f"chunk:{i}"), source_hash="old")
for i in range(count)
]
await index.delete(project_id, deletions)
await index.delete(project_id, deletions)

async with db.scoped_session(repository.session_maker) as session:
await index._ensure_loaded(session)
chunks = await session.scalars(text("SELECT id FROM search_vector_chunks ORDER BY id"))
vectors = await session.scalars(
text("SELECT rowid FROM search_vector_embeddings ORDER BY rowid")
)
expected = sorted(i + 1 for i in protected)
assert chunks.all() == expected
assert vectors.all() == expected


@pytest.mark.parametrize("upgrade", [True, False])
@pytest.mark.parametrize("scratch_populated", [True, False])
def test_fts_rebuild_retries_with_committed_scratch_table(
tmp_path: Path, upgrade: bool, scratch_populated: bool
) -> None:
migration = import_module(
"basic_memory.alembic.versions.d2e3f4a5b6c7_add_script_ngrams_to_full_text_search"
)
engine = create_engine(f"sqlite:///{tmp_path / 'recovery.db'}")
try:
# Commit the interrupted state first, then reopen it as a later startup would.
with engine.begin() as connection:
columns = migration.SQLITE_COLUMNS.format(
script_column="" if upgrade else "script_ngrams,"
)
connection.exec_driver_sql(f"CREATE VIRTUAL TABLE search_index USING fts5({columns})")
connection.exec_driver_sql(
"INSERT INTO search_index (id, title, content_stems, project_id) "
"VALUES (1, 'Preserved note', 'surviving words', 7)"
)
connection.exec_driver_sql(
"CREATE VIRTUAL TABLE search_index_rebuilt USING fts5(stale)"
)
if scratch_populated:
connection.exec_driver_sql(
"INSERT INTO search_index_rebuilt (stale) VALUES ('obsolete scratch data')"
)

with engine.begin() as connection:
with Operations.context(MigrationContext.configure(connection)):
if upgrade:
migration.upgrade()
else:
migration.downgrade()

with engine.connect() as connection:
assert connection.exec_driver_sql(
"SELECT id, title, project_id FROM search_index WHERE search_index MATCH 'surviving'"
).all() == [(1, "Preserved note", 7)]
columns = connection.exec_driver_sql("PRAGMA table_info(search_index)")
assert ("script_ngrams" in {row[1] for row in columns}) == upgrade
assert (
connection.exec_driver_sql(
"SELECT name FROM sqlite_master WHERE name = 'search_index_rebuilt'"
).all()
== []
)
finally:
engine.dispose()
69 changes: 69 additions & 0 deletions tests/repository/test_sqlite_vector_search_repository.py
Original file line number Diff line number Diff line change
Expand Up @@ -1087,6 +1087,75 @@ async def test_sqlite_vec_delete_requires_pending_source_generation(search_repos


@pytest.mark.asyncio
async def test_sqlite_vec_delete_batches_above_expression_depth_limit(search_repository):
"""Deleting more than ~995 pending chunks must not hit SQLITE_MAX_EXPR_DEPTH (#1850)."""
if not isinstance(search_repository, SQLiteSearchRepository):
pytest.skip("sqlite-vec deletion behavior is local SQLite-only.")

_enable_semantic(search_repository)
await search_repository.init_search_index()
index = cast(SQLiteVecIndex, search_repository._semantic_vector_index)
project_id = search_repository.project_id
count = 1200
deletions: list[VectorDeletion] = []

async with db.scoped_session(search_repository.session_maker) as session:
await index._ensure_loaded(session)
for i in range(count):
chunk_key = f"entity:1850:{i}"
row_id = 1850000 + i
await session.execute(
text(
"INSERT INTO search_vector_chunks ("
"id, entity_id, project_id, chunk_key, source_type, source_row_id, "
"chunk_index, chunk_text, source_hash, entity_fingerprint, "
"embedding_model, vector_index, embedding_status"
") VALUES ("
":id, 1850, :project_id, :chunk_key, 'entity', 1850, :chunk_index, "
"'text', 'hash-1850', 'fingerprint', :embedding_model, "
"'sqlite-vec', 'pending')"
),
{
"id": row_id,
"project_id": project_id,
"chunk_key": chunk_key,
"chunk_index": i,
"embedding_model": search_repository._embedding_model_key(),
},
)
await session.execute(
text(
"INSERT INTO search_vector_embeddings (rowid, embedding, source_hash) "
"VALUES (:rowid, :embedding, 'hash-1850')"
),
{"rowid": row_id, "embedding": "[1.0, 0.0, 0.0, 0.0]"},
)
deletions.append(
VectorDeletion(
key=VectorKey(entity_id=1850, chunk_key=chunk_key),
source_hash="hash-1850",
)
)
await session.commit()

await index.delete(project_id, deletions)

async with db.scoped_session(search_repository.session_maker) as session:
remaining_chunks = await session.scalar(
text(
"SELECT COUNT(*) FROM search_vector_chunks "
"WHERE project_id = :project_id AND entity_id = 1850"
),
{"project_id": project_id},
)
remaining_embeddings = await session.scalar(
text("SELECT COUNT(*) FROM search_vector_embeddings WHERE rowid BETWEEN :lo AND :hi"),
{"lo": 1850000, "hi": 1850000 + count - 1},
)
assert remaining_chunks == 0
assert remaining_embeddings == 0


async def test_sqlite_chunk_upsert_and_delete_lifecycle(search_repository):
"""sync_entity_vectors updates changed chunks and clears vectors when source rows disappear."""
if not isinstance(search_repository, SQLiteSearchRepository):
Expand Down
40 changes: 40 additions & 0 deletions tests/test_script_ngram_search_migration.py
Original file line number Diff line number Diff line change
Expand Up @@ -101,3 +101,43 @@ def test_postgres_upgrade_and_downgrade_manage_both_script_indexes(monkeypatch)
assert "idx_search_index_fts_chunks_script_ngrams_fts" in sql
assert "to_tsvector('simple', script_ngrams)" in sql
assert "DROP COLUMN IF EXISTS script_ngrams" in sql


def test_sqlite_upgrade_survives_leftover_search_index_rebuilt(tmp_path, monkeypatch) -> None:
"""Issue #1853: interrupted upgrade left search_index_rebuilt; retry must succeed."""
database_path = tmp_path / "leftover-rebuilt.db"
engine = create_engine(f"sqlite:///{database_path}")
with engine.begin() as connection:
connection.exec_driver_sql("""
CREATE VIRTUAL TABLE search_index USING fts5(
id UNINDEXED, title, content_stems, content_snippet, permalink,
file_path UNINDEXED, type UNINDEXED, project_id UNINDEXED,
from_id UNINDEXED, to_id UNINDEXED, relation_type UNINDEXED,
entity_id UNINDEXED, category UNINDEXED, metadata UNINDEXED,
created_at UNINDEXED, updated_at UNINDEXED,
tokenize='unicode61 tokenchars 0x2F', prefix='1,2,3,4'
)
""")
connection.exec_driver_sql("""
INSERT INTO search_index (id, title, content_stems, project_id)
VALUES (1, 'Probe', 'words', 1)
""")
# Simulate the leftover empty FTS5 table from a killed CREATE.
connection.exec_driver_sql("CREATE VIRTUAL TABLE search_index_rebuilt USING fts5(x)")
monkeypatch.setattr(
migration,
"op",
Operations(MigrationContext.configure(connection)),
)
migration.upgrade()

with sqlite3.connect(database_path) as connection:
leftover = connection.execute(
"SELECT 1 FROM sqlite_master WHERE type = 'table' AND name = 'search_index_rebuilt'"
).fetchone()
assert leftover is None
columns = [row[1] for row in connection.execute("PRAGMA table_info(search_index)")]
assert "script_ngrams" in columns
assert connection.execute(
"SELECT id, title, script_ngrams FROM search_index"
).fetchall() == [(1, "Probe", "")]
Loading