fix(storage): 严格解析本地策略并原子执行数据库迁移
This commit is contained in:
@@ -32,8 +32,12 @@ def connect() -> sqlite3.Connection:
|
||||
# 关闭 Python sqlite3 的隐式事务,提交时机由 transaction() 或显式 commit 控制。
|
||||
conn.isolation_level = None
|
||||
conn.execute("PRAGMA foreign_keys = ON")
|
||||
_load_extension(conn)
|
||||
migrate(conn)
|
||||
try:
|
||||
_load_extension(conn)
|
||||
migrate(conn)
|
||||
except BaseException:
|
||||
conn.close()
|
||||
raise
|
||||
return conn
|
||||
|
||||
|
||||
|
||||
@@ -6,6 +6,7 @@
|
||||
"""
|
||||
|
||||
from datetime import datetime, timezone
|
||||
import sqlite3
|
||||
|
||||
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:
|
||||
"""把尚未应用的迁移脚本按序应用到给定连接。"""
|
||||
conn.execute(
|
||||
@@ -145,9 +158,28 @@ def migrate(conn) -> None:
|
||||
for idx, script in enumerate(MIGRATIONS, start=1):
|
||||
if idx in applied:
|
||||
continue
|
||||
conn.executescript(script)
|
||||
conn.execute(
|
||||
"INSERT INTO schema_migrations (version, applied_at) VALUES (?, ?)",
|
||||
(idx, datetime.now(timezone.utc).isoformat()),
|
||||
)
|
||||
conn.commit()
|
||||
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(
|
||||
"INSERT INTO schema_migrations (version, applied_at) VALUES (?, ?)",
|
||||
(idx, datetime.now(timezone.utc).isoformat()),
|
||||
)
|
||||
conn.execute("COMMIT")
|
||||
except BaseException:
|
||||
if conn.in_transaction:
|
||||
conn.execute("ROLLBACK")
|
||||
raise
|
||||
|
||||
@@ -13,7 +13,10 @@ from dataclasses import dataclass, field
|
||||
from datetime import datetime
|
||||
from pathlib import Path
|
||||
|
||||
import yaml
|
||||
|
||||
from app.contracts import NoteBlock
|
||||
from app.errors import ApiError
|
||||
from app.textutils import count_tokens
|
||||
|
||||
_HEADING_RE = re.compile(r"^(#{1,6})[ \t]+(.*?)\s*$")
|
||||
@@ -70,7 +73,7 @@ def parse_note(
|
||||
created_at=created_at,
|
||||
updated_at=updated_at,
|
||||
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
|
||||
|
||||
|
||||
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]:
|
||||
"""极简 frontmatter 解析,只提取 key: value 行。"""
|
||||
if not markdown.startswith("---"):
|
||||
|
||||
@@ -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)
|
||||
note = await create_transcript_note(job.job_id, TranscriptNoteRequest(title='Private'))
|
||||
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())
|
||||
assert len(calls) >= 2 and all(calls)
|
||||
assert len(calls) >= 3 and all(calls)
|
||||
asyncio.run(scenario())
|
||||
|
||||
@@ -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')
|
||||
Reference in New Issue
Block a user