diff --git a/backend/app/database/db.py b/backend/app/database/db.py index b2de78c..8183929 100644 --- a/backend/app/database/db.py +++ b/backend/app/database/db.py @@ -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 diff --git a/backend/app/database/migrations.py b/backend/app/database/migrations.py index 5fd1bd8..e6ff960 100644 --- a/backend/app/database/migrations.py +++ b/backend/app/database/migrations.py @@ -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 diff --git a/backend/app/knowledge/parser.py b/backend/app/knowledge/parser.py index d3217f5..fa2ed89 100644 --- a/backend/app/knowledge/parser.py +++ b/backend/app/knowledge/parser.py @@ -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("---"): diff --git a/backend/tests/test_media_jobs.py b/backend/tests/test_media_jobs.py index 20ecfa4..5b73980 100644 --- a/backend/tests/test_media_jobs.py +++ b/backend/tests/test_media_jobs.py @@ -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()) diff --git a/backend/tests/test_policy_and_migrations.py b/backend/tests/test_policy_and_migrations.py new file mode 100644 index 0000000..c65cff2 --- /dev/null +++ b/backend/tests/test_policy_and_migrations.py @@ -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')