269 lines
10 KiB
Python
269 lines
10 KiB
Python
"""원문 위치를 보존하는 SQLite 코퍼스 저장소.
|
|
|
|
검색용 전처리 텍스트와 증거 표시용 원문을 분리한다. 현재 수령한 XLSX처럼
|
|
페이지 정보가 없는 데이터도 ``coordinate_scope=episode`` 로 정직하게 기록하며,
|
|
향후 PDF/DOCX 재추출 시 동일 스키마에 page/paragraph/char offset을 채울 수 있다.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import hashlib
|
|
import json
|
|
import sqlite3
|
|
from dataclasses import dataclass, field
|
|
from pathlib import Path
|
|
from typing import Iterable, Iterator
|
|
|
|
|
|
SCHEMA_VERSION = 1
|
|
|
|
|
|
def text_sha256(text: str) -> str:
|
|
return hashlib.sha256(text.encode("utf-8")).hexdigest()
|
|
|
|
|
|
def stable_id(prefix: str, *parts: str) -> str:
|
|
payload = "\x1f".join(parts).encode("utf-8")
|
|
return f"{prefix}-{hashlib.sha256(payload).hexdigest()[:20]}"
|
|
|
|
|
|
@dataclass(frozen=True)
|
|
class DocumentRecord:
|
|
document_id: str
|
|
title: str
|
|
source_path: str | None = None
|
|
source_sha256: str | None = None
|
|
metadata: dict = field(default_factory=dict)
|
|
|
|
|
|
@dataclass(frozen=True)
|
|
class SegmentRecord:
|
|
segment_id: str
|
|
document_id: str
|
|
text: str
|
|
ordinal: str
|
|
coordinate_scope: str = "episode"
|
|
page_number: int | None = None
|
|
paragraph_number: int | None = None
|
|
char_start: int | None = None
|
|
char_end: int | None = None
|
|
source_locator: str | None = None
|
|
metadata: dict = field(default_factory=dict)
|
|
|
|
@property
|
|
def text_sha256(self) -> str:
|
|
return text_sha256(self.text)
|
|
|
|
|
|
class CorpusStore:
|
|
"""SQLite 기반 원문/세그먼트 저장소.
|
|
|
|
연결은 호출 단위로 열어 멀티프로세스 API에서도 안전하게 사용한다. WAL은
|
|
읽기 중 증분 적재를 허용한다.
|
|
"""
|
|
|
|
def __init__(self, path: str | Path):
|
|
self.path = Path(path)
|
|
|
|
def _connect(self) -> sqlite3.Connection:
|
|
self.path.parent.mkdir(parents=True, exist_ok=True)
|
|
con = sqlite3.connect(self.path)
|
|
con.row_factory = sqlite3.Row
|
|
con.execute("PRAGMA journal_mode=WAL")
|
|
con.execute("PRAGMA foreign_keys=ON")
|
|
return con
|
|
|
|
def initialize(self) -> None:
|
|
with self._connect() as con:
|
|
con.executescript(
|
|
"""
|
|
CREATE TABLE IF NOT EXISTS corpus_meta (
|
|
key TEXT PRIMARY KEY,
|
|
value TEXT NOT NULL
|
|
);
|
|
CREATE TABLE IF NOT EXISTS documents (
|
|
document_id TEXT PRIMARY KEY,
|
|
title TEXT NOT NULL,
|
|
source_path TEXT,
|
|
source_sha256 TEXT,
|
|
metadata_json TEXT NOT NULL DEFAULT '{}',
|
|
created_at TEXT NOT NULL DEFAULT CURRENT_TIMESTAMP,
|
|
updated_at TEXT NOT NULL DEFAULT CURRENT_TIMESTAMP
|
|
);
|
|
CREATE TABLE IF NOT EXISTS segments (
|
|
segment_id TEXT PRIMARY KEY,
|
|
document_id TEXT NOT NULL REFERENCES documents(document_id) ON DELETE CASCADE,
|
|
ordinal TEXT NOT NULL,
|
|
text TEXT NOT NULL,
|
|
text_sha256 TEXT NOT NULL,
|
|
coordinate_scope TEXT NOT NULL,
|
|
page_number INTEGER,
|
|
paragraph_number INTEGER,
|
|
char_start INTEGER,
|
|
char_end INTEGER,
|
|
source_locator TEXT,
|
|
metadata_json TEXT NOT NULL DEFAULT '{}',
|
|
created_at TEXT NOT NULL DEFAULT CURRENT_TIMESTAMP,
|
|
UNIQUE(document_id, text_sha256)
|
|
);
|
|
CREATE INDEX IF NOT EXISTS idx_segments_document ON segments(document_id);
|
|
CREATE INDEX IF NOT EXISTS idx_segments_hash ON segments(text_sha256);
|
|
"""
|
|
)
|
|
con.execute(
|
|
"INSERT OR REPLACE INTO corpus_meta(key, value) VALUES('schema_version', ?)",
|
|
(str(SCHEMA_VERSION),),
|
|
)
|
|
|
|
def upsert_document(self, record: DocumentRecord) -> None:
|
|
self.upsert_documents([record])
|
|
|
|
def upsert_documents(self, records: Iterable[DocumentRecord]) -> None:
|
|
self.initialize()
|
|
with self._connect() as con:
|
|
con.executemany(
|
|
"""
|
|
INSERT INTO documents(document_id,title,source_path,source_sha256,metadata_json)
|
|
VALUES(?,?,?,?,?)
|
|
ON CONFLICT(document_id) DO UPDATE SET
|
|
title=excluded.title,
|
|
source_path=COALESCE(excluded.source_path, documents.source_path),
|
|
source_sha256=COALESCE(excluded.source_sha256, documents.source_sha256),
|
|
metadata_json=excluded.metadata_json,
|
|
updated_at=CURRENT_TIMESTAMP
|
|
""",
|
|
[(
|
|
record.document_id, record.title, record.source_path,
|
|
record.source_sha256,
|
|
json.dumps(record.metadata, ensure_ascii=False, sort_keys=True),
|
|
) for record in records],
|
|
)
|
|
|
|
def add_segments(self, records: Iterable[SegmentRecord]) -> tuple[int, int]:
|
|
"""세그먼트를 추가하고 (inserted, duplicates) 반환."""
|
|
self.initialize()
|
|
inserted = duplicates = 0
|
|
with self._connect() as con:
|
|
for r in records:
|
|
before = con.total_changes
|
|
con.execute(
|
|
"""
|
|
INSERT OR IGNORE INTO segments(
|
|
segment_id,document_id,ordinal,text,text_sha256,coordinate_scope,
|
|
page_number,paragraph_number,char_start,char_end,source_locator,metadata_json
|
|
) VALUES(?,?,?,?,?,?,?,?,?,?,?,?)
|
|
""",
|
|
(
|
|
r.segment_id,
|
|
r.document_id,
|
|
r.ordinal,
|
|
r.text,
|
|
r.text_sha256,
|
|
r.coordinate_scope,
|
|
r.page_number,
|
|
r.paragraph_number,
|
|
r.char_start,
|
|
r.char_end,
|
|
r.source_locator,
|
|
json.dumps(r.metadata, ensure_ascii=False, sort_keys=True),
|
|
),
|
|
)
|
|
if con.total_changes > before:
|
|
inserted += 1
|
|
else:
|
|
duplicates += 1
|
|
return inserted, duplicates
|
|
|
|
def iter_segments(self) -> Iterator[SegmentRecord]:
|
|
if not self.path.exists():
|
|
return
|
|
with self._connect() as con:
|
|
rows = con.execute(
|
|
"""
|
|
SELECT s.*, d.title
|
|
FROM segments s JOIN documents d USING(document_id)
|
|
ORDER BY s.document_id, s.ordinal, s.segment_id
|
|
"""
|
|
)
|
|
for row in rows:
|
|
metadata = json.loads(row["metadata_json"] or "{}")
|
|
metadata.setdefault("document_title", row["title"])
|
|
yield SegmentRecord(
|
|
segment_id=row["segment_id"],
|
|
document_id=row["document_id"],
|
|
text=row["text"],
|
|
ordinal=row["ordinal"],
|
|
coordinate_scope=row["coordinate_scope"],
|
|
page_number=row["page_number"],
|
|
paragraph_number=row["paragraph_number"],
|
|
char_start=row["char_start"],
|
|
char_end=row["char_end"],
|
|
source_locator=row["source_locator"],
|
|
metadata=metadata,
|
|
)
|
|
|
|
def get_segments(self, segment_ids: Iterable[str]) -> dict[str, SegmentRecord]:
|
|
ids = list(dict.fromkeys(segment_ids))
|
|
if not ids or not self.path.exists():
|
|
return {}
|
|
result: dict[str, SegmentRecord] = {}
|
|
with self._connect() as con:
|
|
for start in range(0, len(ids), 500):
|
|
batch = ids[start : start + 500]
|
|
marks = ",".join("?" for _ in batch)
|
|
rows = con.execute(
|
|
f"""SELECT s.*, d.title FROM segments s
|
|
JOIN documents d USING(document_id)
|
|
WHERE segment_id IN ({marks})""",
|
|
batch,
|
|
)
|
|
for row in rows:
|
|
metadata = json.loads(row["metadata_json"] or "{}")
|
|
metadata.setdefault("document_title", row["title"])
|
|
result[row["segment_id"]] = SegmentRecord(
|
|
segment_id=row["segment_id"], document_id=row["document_id"],
|
|
text=row["text"], ordinal=row["ordinal"],
|
|
coordinate_scope=row["coordinate_scope"],
|
|
page_number=row["page_number"], paragraph_number=row["paragraph_number"],
|
|
char_start=row["char_start"], char_end=row["char_end"],
|
|
source_locator=row["source_locator"], metadata=metadata,
|
|
)
|
|
return result
|
|
|
|
def stats(self) -> dict[str, int]:
|
|
if not self.path.exists():
|
|
return {"documents": 0, "segments": 0, "characters": 0}
|
|
with self._connect() as con:
|
|
documents = con.execute("SELECT COUNT(*) FROM documents").fetchone()[0]
|
|
segments, characters = con.execute(
|
|
"SELECT COUNT(*), COALESCE(SUM(LENGTH(text)),0) FROM segments"
|
|
).fetchone()
|
|
return {"documents": documents, "segments": segments, "characters": characters}
|
|
|
|
def document_count(self) -> int:
|
|
if not self.path.exists():
|
|
return 0
|
|
with self._connect() as con:
|
|
return int(con.execute("SELECT COUNT(*) FROM documents").fetchone()[0])
|
|
|
|
def list_documents(self) -> list[dict]:
|
|
if not self.path.exists():
|
|
return []
|
|
with self._connect() as con:
|
|
rows = con.execute(
|
|
"""SELECT d.document_id, d.title, d.source_path,
|
|
COUNT(s.segment_id) AS segment_count,
|
|
COALESCE(SUM(LENGTH(s.text)),0) AS characters
|
|
FROM documents d LEFT JOIN segments s USING(document_id)
|
|
GROUP BY d.document_id ORDER BY d.title"""
|
|
).fetchall()
|
|
return [dict(row) for row in rows]
|
|
|
|
def delete_document(self, document_id: str) -> bool:
|
|
if not self.path.exists():
|
|
return False
|
|
with self._connect() as con:
|
|
before = con.total_changes
|
|
con.execute("DELETE FROM documents WHERE document_id=?", (document_id,))
|
|
return con.total_changes > before
|