fix(storage): 严格解析本地策略并原子执行数据库迁移

This commit is contained in:
2026-09-04 19:57:26 +08:00
parent 6661f4f8c6
commit b1c85b3ff7
5 changed files with 190 additions and 10 deletions
+4
View File
@@ -32,8 +32,12 @@ def connect() -> sqlite3.Connection:
# 关闭 Python sqlite3 的隐式事务,提交时机由 transaction() 或显式 commit 控制。 # 关闭 Python sqlite3 的隐式事务,提交时机由 transaction() 或显式 commit 控制。
conn.isolation_level = None conn.isolation_level = None
conn.execute("PRAGMA foreign_keys = ON") conn.execute("PRAGMA foreign_keys = ON")
try:
_load_extension(conn) _load_extension(conn)
migrate(conn) migrate(conn)
except BaseException:
conn.close()
raise
return conn return conn
+34 -2
View File
@@ -6,6 +6,7 @@
""" """
from datetime import datetime, timezone from datetime import datetime, timezone
import sqlite3
from app.constants import EMBEDDING_DIM from app.constants import EMBEDDING_DIM
@@ -134,6 +135,18 @@ MIGRATIONS: list[str] = [
] ]
def _statements(script: str):
"""Split complete SQLite statements without executescript's implicit COMMIT."""
pending = ""
for char in script:
pending += char
if char == ";" and sqlite3.complete_statement(pending):
yield pending
pending = ""
if pending.strip():
yield pending
def migrate(conn) -> None: def migrate(conn) -> None:
"""把尚未应用的迁移脚本按序应用到给定连接。""" """把尚未应用的迁移脚本按序应用到给定连接。"""
conn.execute( conn.execute(
@@ -145,9 +158,28 @@ def migrate(conn) -> None:
for idx, script in enumerate(MIGRATIONS, start=1): for idx, script in enumerate(MIGRATIONS, start=1):
if idx in applied: if idx in applied:
continue continue
conn.executescript(script) conn.execute("BEGIN IMMEDIATE")
try:
# Another connection may have migrated while this one waited.
if not conn.execute("SELECT 1 FROM schema_migrations WHERE version=?", (idx,)).fetchone():
recovered_v6 = False
if idx == 6:
column = next((row for row in conn.execute("PRAGMA table_info(blocks)")
if row["name"] == "embedding_local_only"), None)
if column is not None:
# Recover the precise partial state left by the old v6 runner.
if column["type"].upper() != "INTEGER" or column["notnull"] != 1 or column["dflt_value"] != "0":
raise sqlite3.DatabaseError("Unexpected embedding_local_only column schema")
recovered_v6 = True
if not recovered_v6:
for statement in _statements(script):
conn.execute(statement)
conn.execute( conn.execute(
"INSERT INTO schema_migrations (version, applied_at) VALUES (?, ?)", "INSERT INTO schema_migrations (version, applied_at) VALUES (?, ?)",
(idx, datetime.now(timezone.utc).isoformat()), (idx, datetime.now(timezone.utc).isoformat()),
) )
conn.commit() conn.execute("COMMIT")
except BaseException:
if conn.in_transaction:
conn.execute("ROLLBACK")
raise
+33 -1
View File
@@ -13,7 +13,10 @@ from dataclasses import dataclass, field
from datetime import datetime from datetime import datetime
from pathlib import Path from pathlib import Path
import yaml
from app.contracts import NoteBlock from app.contracts import NoteBlock
from app.errors import ApiError
from app.textutils import count_tokens from app.textutils import count_tokens
_HEADING_RE = re.compile(r"^(#{1,6})[ \t]+(.*?)\s*$") _HEADING_RE = re.compile(r"^(#{1,6})[ \t]+(.*?)\s*$")
@@ -70,7 +73,7 @@ def parse_note(
created_at=created_at, created_at=created_at,
updated_at=updated_at, updated_at=updated_at,
blocks=blocks, blocks=blocks,
embedding_local_only=str(frontmatter.get("embedding_local_only", "")).lower() == "true", embedding_local_only=_embedding_policy(markdown),
) )
@@ -184,6 +187,35 @@ def _utf16_len(text: str) -> int:
return len(text.encode("utf-16-le")) // 2 return len(text.encode("utf-16-le")) // 2
def _embedding_policy(markdown: str) -> bool:
end = markdown.find("\n---", 3) if markdown.startswith("---") else -1
if end == -1:
return False
try:
# Compose nodes without constructing objects. This accepts YAML comments,
# quoted keys and indentation while retaining duplicate-key information.
node = yaml.compose(markdown[3:end], Loader=yaml.SafeLoader)
except yaml.YAMLError as exc:
raise ApiError(422, "INVALID_EMBEDDING_POLICY", "Frontmatter YAML 无效,无法确认本地索引策略。") from exc
if node is None:
return False
if not isinstance(node, yaml.MappingNode):
raise ApiError(422, "INVALID_EMBEDDING_POLICY", "Frontmatter 必须是 YAML 键值映射。")
if any(key.tag == "tag:yaml.org,2002:merge" for key, _ in node.value):
raise ApiError(422, "INVALID_EMBEDDING_POLICY", "Frontmatter 不支持 YAML 合并键,请显式声明索引策略。")
values = [value for key, value in node.value
if isinstance(key, yaml.ScalarNode) and key.value.lower() == "embedding_local_only"]
if not values:
return False
if len(values) > 1:
raise ApiError(422, "INVALID_EMBEDDING_POLICY", "embedding_local_only 不能重复声明。")
value = values[0]
if (not isinstance(value, yaml.ScalarNode) or value.tag != "tag:yaml.org,2002:bool"
or value.value.lower() not in {"true", "false", "yes", "no", "on", "off"}):
raise ApiError(422, "INVALID_EMBEDDING_POLICY", "embedding_local_only 必须是 YAML 布尔值 true 或 false。")
return value.value.lower() in {"true", "yes", "on"}
def _extract_frontmatter(markdown: str) -> dict[str, str]: def _extract_frontmatter(markdown: str) -> dict[str, str]:
"""极简 frontmatter 解析,只提取 key: value 行。""" """极简 frontmatter 解析,只提取 key: value 行。"""
if not markdown.startswith("---"): if not markdown.startswith("---"):
+3 -1
View File
@@ -125,6 +125,8 @@ def test_local_only_export_and_rebuild_keep_local_embedding_policy(monkeypatch):
job = await jobs.create_transcription('lecture.txt', local_only=True) job = await jobs.create_transcription('lecture.txt', local_only=True)
note = await create_transcript_note(job.job_id, TranscriptNoteRequest(title='Private')) note = await create_transcript_note(job.job_id, TranscriptNoteRequest(title='Private'))
assert note.markdown.startswith('---\nembedding_local_only: true\n---') assert note.markdown.startswith('---\nembedding_local_only: true\n---')
await note_service.update_note(note.note_id, markdown=note.markdown.replace(
'embedding_local_only: true', 'embedding_local_only: true # keep local'))
await index_service.rebuild(IndexRebuildRequest()) await index_service.rebuild(IndexRebuildRequest())
assert len(calls) >= 2 and all(calls) assert len(calls) >= 3 and all(calls)
asyncio.run(scenario()) asyncio.run(scenario())
+110
View File
@@ -0,0 +1,110 @@
import sqlite3
from datetime import datetime, timezone
from concurrent.futures import ThreadPoolExecutor
import pytest
from app.database import migrations
from app.database.db import _load_extension
from app.errors import ApiError
from app.knowledge.parser import parse_note
def parsed(value):
return parse_note(markdown='---\nembedding_local_only: '+value+'\n---\nbody', file_path='note.md', folder='',
created_at=datetime.now(timezone.utc), updated_at=datetime.now(timezone.utc))
@pytest.mark.parametrize('value,expected', [('true', True), ('true # keep local', True), ('TRUE # comment', True), ('false # explicit', False)])
def test_policy_parses_yaml_boolean_with_comments(value, expected):
assert parsed(value).embedding_local_only is expected
@pytest.mark.parametrize('value', ['truth', '1', '', 'null', '"true"', '[true]', '{broken', 'true\nembedding_local_only: false'])
def test_invalid_policy_never_silently_enables_remote(value):
with pytest.raises(ApiError) as error:
parsed(value)
assert error.value.code == 'INVALID_EMBEDDING_POLICY'
def connection(path, factory=sqlite3.Connection):
conn = sqlite3.connect(path, isolation_level=None, factory=factory)
conn.row_factory = sqlite3.Row
_load_extension(conn)
return conn
def seed_v5(path, monkeypatch):
conn = connection(path)
with monkeypatch.context() as patch:
patch.setattr(migrations, 'MIGRATIONS', migrations.MIGRATIONS[:5])
migrations.migrate(conn)
conn.execute("INSERT INTO search_history(query) VALUES ('retained')")
conn.close()
@pytest.mark.parametrize('failure', [sqlite3.OperationalError, KeyboardInterrupt])
def test_migration_and_version_write_rollback_together(tmp_path, monkeypatch, failure):
path = tmp_path / 'migration.db'
seed_v5(path, monkeypatch)
class Interrupted(sqlite3.Connection):
def execute(self, sql, parameters=()):
if sql.startswith('INSERT INTO schema_migrations') and parameters[0] == 6:
raise failure('interrupted')
return super().execute(sql, parameters)
conn = connection(path, Interrupted)
try:
with pytest.raises(failure):
migrations.migrate(conn)
assert not conn.in_transaction
assert not any(r['name'] == 'embedding_local_only' for r in conn.execute('pragma table_info(blocks)'))
finally:
conn.close()
conn = connection(path)
try:
migrations.migrate(conn)
assert conn.execute('select count(*) from schema_migrations where version=6').fetchone()[0] == 1
assert conn.execute('select query from search_history').fetchone()[0] == 'retained'
finally:
conn.close()
def test_old_partial_v6_recovers_without_duplicate_column(tmp_path, monkeypatch):
path = tmp_path / 'partial.db'
seed_v5(path, monkeypatch)
conn = connection(path)
try:
conn.executescript(migrations.MIGRATIONS[5])
migrations.migrate(conn)
migrations.migrate(conn)
assert conn.execute('select count(*) from schema_migrations where version=6').fetchone()[0] == 1
assert conn.execute('select query from search_history').fetchone()[0] == 'retained'
finally:
conn.close()
def test_concurrent_connections_can_upgrade(tmp_path, monkeypatch):
path = tmp_path / 'concurrent.db'
seed_v5(path, monkeypatch)
def upgrade(_):
conn = connection(path)
try:
migrations.migrate(conn)
return conn.execute('select count(*) from schema_migrations where version=6').fetchone()[0]
finally:
conn.close()
with ThreadPoolExecutor(max_workers=2) as pool:
assert list(pool.map(upgrade, range(2))) == [1, 1]
@pytest.mark.parametrize('header', ['"embedding_local_only": true # comment', ' embedding_local_only: true', 'embedding_local_only:\n true', 'local: &local true\nembedding_local_only: *local'])
def test_policy_supports_yaml_key_and_scalar_forms(header):
note = parse_note(markdown='---\n'+header+'\n---\nbody',file_path='note.md',folder='',created_at=datetime.now(timezone.utc),updated_at=datetime.now(timezone.utc))
assert note.embedding_local_only
def test_merge_policy_is_rejected_instead_of_ignored():
with pytest.raises(ApiError):
parsed('true\n<<: {embedding_local_only: false}')
with pytest.raises(ApiError):
parsed('!!bool invalid')