388 lines
16 KiB
Python
388 lines
16 KiB
Python
"""Database connection — SQLite with thread-local connections.
|
||
|
||
Архитектура (Sonnet review 2026-07-15):
|
||
- threading.local() — каждому потоку своё соединение
|
||
- WAL-режим — readers не блокируют writer
|
||
- _db_session_key — защита от inode split-brain при cleanup+reuse потоков
|
||
- busy_timeout=5000 — ждать при конкурентной записи
|
||
"""
|
||
import sqlite3
|
||
import threading
|
||
import time
|
||
import os
|
||
from contextlib import contextmanager
|
||
|
||
DB_PATH = os.getenv("CONTRACTS_DB_PATH", "/tmp/contracts.db")
|
||
_db_session_key = None
|
||
_local = threading.local()
|
||
|
||
|
||
@contextmanager
|
||
def transaction():
|
||
"""Run DB writes atomically on the current thread connection."""
|
||
conn = get_conn()
|
||
conn.execute("BEGIN IMMEDIATE")
|
||
_local.in_transaction = True
|
||
try:
|
||
yield conn
|
||
conn.commit()
|
||
except Exception:
|
||
conn.rollback()
|
||
raise
|
||
finally:
|
||
_local.in_transaction = False
|
||
|
||
|
||
def init_db():
|
||
"""Create the database schema without deleting persisted data."""
|
||
global _db_session_key
|
||
_db_session_key = time.time_ns()
|
||
conn = get_conn()
|
||
_create_schema(conn)
|
||
return conn
|
||
|
||
|
||
def _create_schema(conn):
|
||
"""Создать все таблицы (idempotent)."""
|
||
conn.executescript("""
|
||
CREATE TABLE IF NOT EXISTS documents (
|
||
id TEXT PRIMARY KEY,
|
||
filename TEXT NOT NULL,
|
||
mime_type TEXT DEFAULT 'application/octet-stream',
|
||
original_bytes TEXT,
|
||
status TEXT DEFAULT 'uploaded',
|
||
error_message TEXT,
|
||
elements_json TEXT,
|
||
batch_id TEXT,
|
||
zip_source TEXT,
|
||
content_hash TEXT,
|
||
classify_status TEXT DEFAULT 'pending',
|
||
doc_type TEXT,
|
||
own_number TEXT,
|
||
parent_number TEXT,
|
||
doc_date TEXT,
|
||
counterparty TEXT,
|
||
classify_raw TEXT,
|
||
classify_input TEXT,
|
||
created_at TEXT DEFAULT (datetime('now'))
|
||
);
|
||
|
||
CREATE TABLE IF NOT EXISTS extraction_jobs (
|
||
batch_id TEXT PRIMARY KEY,
|
||
owner TEXT NOT NULL,
|
||
state TEXT NOT NULL,
|
||
total INTEGER NOT NULL,
|
||
done INTEGER NOT NULL DEFAULT 0,
|
||
failed INTEGER NOT NULL DEFAULT 0,
|
||
classes_json TEXT NOT NULL DEFAULT '{}',
|
||
result_json TEXT,
|
||
error TEXT,
|
||
updated_at REAL NOT NULL
|
||
);
|
||
|
||
CREATE TABLE IF NOT EXISTS contracts (
|
||
id TEXT PRIMARY KEY,
|
||
number TEXT NOT NULL,
|
||
client TEXT DEFAULT '',
|
||
created_at TEXT DEFAULT (datetime('now'))
|
||
);
|
||
|
||
CREATE TABLE IF NOT EXISTS supplements (
|
||
id TEXT PRIMARY KEY,
|
||
contract_id TEXT NOT NULL REFERENCES contracts(id),
|
||
document_id TEXT NOT NULL REFERENCES documents(id),
|
||
type TEXT DEFAULT 'additional',
|
||
created_at TEXT DEFAULT (datetime('now'))
|
||
);
|
||
|
||
CREATE TABLE IF NOT EXISTS spec_events (
|
||
id TEXT PRIMARY KEY,
|
||
contract_id TEXT NOT NULL,
|
||
supplement_id TEXT NOT NULL,
|
||
seq INTEGER NOT NULL DEFAULT 0,
|
||
action TEXT NOT NULL DEFAULT 'UNRESOLVED',
|
||
target_hash TEXT DEFAULT '',
|
||
new_values TEXT DEFAULT '{}',
|
||
comment TEXT DEFAULT '',
|
||
status TEXT DEFAULT 'unresolved',
|
||
prompt_version TEXT DEFAULT '',
|
||
source_document_id TEXT DEFAULT '',
|
||
raw_llm_response TEXT DEFAULT '{}',
|
||
created_at TEXT DEFAULT (datetime('now'))
|
||
);
|
||
|
||
CREATE TABLE IF NOT EXISTS spec_current (
|
||
id TEXT PRIMARY KEY,
|
||
contract_id TEXT NOT NULL,
|
||
name_hash TEXT NOT NULL,
|
||
name TEXT,
|
||
price REAL,
|
||
qty REAL,
|
||
sum REAL,
|
||
date_start TEXT,
|
||
last_event_id TEXT,
|
||
updated_at TEXT DEFAULT (datetime('now')),
|
||
created_at TEXT DEFAULT (datetime('now'))
|
||
);
|
||
|
||
CREATE TABLE IF NOT EXISTS prompts (
|
||
id TEXT PRIMARY KEY,
|
||
role TEXT NOT NULL,
|
||
name TEXT DEFAULT '',
|
||
body TEXT NOT NULL,
|
||
is_active INTEGER DEFAULT 0,
|
||
notes TEXT,
|
||
created_at TEXT DEFAULT (datetime('now'))
|
||
);
|
||
|
||
CREATE TABLE IF NOT EXISTS upload_chunks (
|
||
id TEXT PRIMARY KEY,
|
||
upload_id TEXT NOT NULL,
|
||
chunk_index INTEGER NOT NULL,
|
||
data TEXT,
|
||
created_at TEXT DEFAULT (datetime('now'))
|
||
);
|
||
|
||
-- Извлечённый паспорт хранится отдельно от исходного документа, чтобы
|
||
-- повторный запуск пайплайна не перезаписывал исходные байты и парсинг.
|
||
CREATE TABLE IF NOT EXISTS document_extractions (
|
||
id TEXT PRIMARY KEY,
|
||
document_id TEXT NOT NULL UNIQUE REFERENCES documents(id) ON DELETE CASCADE,
|
||
doc_class TEXT NOT NULL DEFAULT 'unknown',
|
||
intent_mode TEXT NOT NULL DEFAULT 'no_spec',
|
||
own_number TEXT,
|
||
own_date TEXT,
|
||
counterparty_name TEXT,
|
||
counterparty_inn TEXT,
|
||
referenced_numbers_json TEXT NOT NULL DEFAULT '[]',
|
||
sections_json TEXT NOT NULL DEFAULT '[]',
|
||
math_validation_json TEXT NOT NULL DEFAULT '{}',
|
||
status TEXT NOT NULL DEFAULT 'pending',
|
||
prompt_version TEXT NOT NULL DEFAULT '',
|
||
raw_llm_response TEXT NOT NULL DEFAULT '{}',
|
||
created_at TEXT DEFAULT (datetime('now')),
|
||
updated_at TEXT DEFAULT (datetime('now'))
|
||
);
|
||
|
||
-- Досье создаётся в рамках batch; отсутствие рамочного договора не
|
||
-- запрещает реконструкцию по спецификациям и дополнительным соглашениям.
|
||
CREATE TABLE IF NOT EXISTS contract_dossiers (
|
||
id TEXT PRIMARY KEY,
|
||
batch_id TEXT,
|
||
canonical_number TEXT NOT NULL,
|
||
alias_numbers_json TEXT NOT NULL DEFAULT '[]',
|
||
counterparty_name TEXT NOT NULL DEFAULT '',
|
||
counterparty_inn TEXT NOT NULL DEFAULT '',
|
||
has_frame_contract INTEGER NOT NULL DEFAULT 0,
|
||
status TEXT NOT NULL DEFAULT 'ACTIVE',
|
||
created_at TEXT DEFAULT (datetime('now'))
|
||
);
|
||
|
||
-- Связь задаёт принятую принадлежность документа и сохраняет её основание.
|
||
-- UNIQUE предотвращает дубли при повторной сборке групп одного batch.
|
||
CREATE TABLE IF NOT EXISTS dossier_documents (
|
||
id TEXT PRIMARY KEY,
|
||
dossier_id TEXT NOT NULL REFERENCES contract_dossiers(id) ON DELETE CASCADE,
|
||
document_id TEXT NOT NULL REFERENCES documents(id) ON DELETE CASCADE,
|
||
role TEXT NOT NULL DEFAULT 'supplement',
|
||
link_status TEXT NOT NULL DEFAULT 'CONFIRMED',
|
||
link_reason TEXT NOT NULL DEFAULT '',
|
||
seq_order INTEGER NOT NULL DEFAULT 0,
|
||
created_at TEXT DEFAULT (datetime('now')),
|
||
UNIQUE (dossier_id, document_id)
|
||
);
|
||
|
||
-- Идентичность услуги отделена от параметров и цены, которые меняются
|
||
-- по временным интервалам.
|
||
CREATE TABLE IF NOT EXISTS service_items (
|
||
id TEXT PRIMARY KEY,
|
||
dossier_id TEXT NOT NULL REFERENCES contract_dossiers(id) ON DELETE CASCADE,
|
||
service_family TEXT NOT NULL,
|
||
instance_anchor TEXT NOT NULL DEFAULT '',
|
||
canonical_name TEXT NOT NULL,
|
||
catalog_sku TEXT,
|
||
sku_status TEXT NOT NULL DEFAULT 'UNKNOWN',
|
||
created_at TEXT DEFAULT (datetime('now'))
|
||
);
|
||
|
||
-- Источник интервала каскадно удаляется вместе с документом. Документ,
|
||
-- закрывший интервал, очищается через SET NULL: сам интервал остаётся.
|
||
CREATE TABLE IF NOT EXISTS timeline_intervals (
|
||
id TEXT PRIMARY KEY,
|
||
dossier_id TEXT NOT NULL REFERENCES contract_dossiers(id) ON DELETE CASCADE,
|
||
service_item_id TEXT NOT NULL REFERENCES service_items(id) ON DELETE CASCADE,
|
||
charge_type TEXT NOT NULL DEFAULT 'recurring',
|
||
valid_from TEXT,
|
||
valid_to TEXT,
|
||
date_status TEXT NOT NULL DEFAULT 'explicit',
|
||
date_basis TEXT NOT NULL DEFAULT '',
|
||
raw_name TEXT NOT NULL,
|
||
config_params_json TEXT NOT NULL DEFAULT '{}',
|
||
price TEXT,
|
||
qty TEXT,
|
||
sum TEXT,
|
||
vat_rate TEXT,
|
||
vat_mode TEXT NOT NULL DEFAULT 'with_vat',
|
||
source_document_id TEXT REFERENCES documents(id) ON DELETE CASCADE,
|
||
closed_by_document_id TEXT REFERENCES documents(id) ON DELETE SET NULL,
|
||
table_ref TEXT NOT NULL DEFAULT '',
|
||
quote TEXT NOT NULL DEFAULT '',
|
||
status TEXT NOT NULL DEFAULT 'VERIFIED',
|
||
created_at TEXT DEFAULT (datetime('now'))
|
||
);
|
||
|
||
-- Исключения могут быть привязаны к разным этапам и разрешаться отдельно.
|
||
CREATE TABLE IF NOT EXISTS exceptions_queue (
|
||
id TEXT PRIMARY KEY,
|
||
batch_id TEXT,
|
||
dossier_id TEXT REFERENCES contract_dossiers(id) ON DELETE CASCADE,
|
||
document_id TEXT REFERENCES documents(id) ON DELETE CASCADE,
|
||
service_item_id TEXT REFERENCES service_items(id) ON DELETE SET NULL,
|
||
severity TEXT NOT NULL DEFAULT 'CONFLICT',
|
||
stage TEXT NOT NULL,
|
||
title TEXT NOT NULL,
|
||
source_quote TEXT NOT NULL DEFAULT '',
|
||
options_json TEXT NOT NULL DEFAULT '[]',
|
||
resolution_json TEXT,
|
||
status TEXT NOT NULL DEFAULT 'OPEN',
|
||
created_at TEXT DEFAULT (datetime('now')),
|
||
resolved_at TEXT
|
||
);
|
||
|
||
CREATE INDEX IF NOT EXISTS idx_documents_batch ON documents(batch_id);
|
||
CREATE INDEX IF NOT EXISTS idx_documents_hash ON documents(batch_id, content_hash);
|
||
CREATE INDEX IF NOT EXISTS idx_supplements_contract ON supplements(contract_id);
|
||
CREATE INDEX IF NOT EXISTS idx_spec_current_contract ON spec_current(contract_id);
|
||
CREATE INDEX IF NOT EXISTS idx_spec_events_supplement ON spec_events(supplement_id);
|
||
CREATE INDEX IF NOT EXISTS idx_extractions_document ON document_extractions(document_id);
|
||
CREATE INDEX IF NOT EXISTS idx_dossiers_batch ON contract_dossiers(batch_id);
|
||
CREATE INDEX IF NOT EXISTS idx_dossier_documents_dossier ON dossier_documents(dossier_id);
|
||
CREATE INDEX IF NOT EXISTS idx_dossier_documents_document ON dossier_documents(document_id);
|
||
CREATE INDEX IF NOT EXISTS idx_service_items_dossier ON service_items(dossier_id);
|
||
CREATE INDEX IF NOT EXISTS idx_timeline_dossier ON timeline_intervals(dossier_id);
|
||
CREATE INDEX IF NOT EXISTS idx_timeline_service ON timeline_intervals(service_item_id);
|
||
CREATE INDEX IF NOT EXISTS idx_exceptions_dossier ON exceptions_queue(dossier_id);
|
||
CREATE INDEX IF NOT EXISTS idx_exceptions_batch ON exceptions_queue(batch_id);
|
||
""")
|
||
# Existing /tmp databases survive application restarts. Additive migrations
|
||
# therefore inspect the live schema before ALTER TABLE, while fresh databases
|
||
# receive the same column from the CREATE TABLE definition above.
|
||
timeline_columns = {
|
||
row["name"] for row in conn.execute("PRAGMA table_info(timeline_intervals)").fetchall()
|
||
}
|
||
if "vat_rate" not in timeline_columns:
|
||
conn.execute("ALTER TABLE timeline_intervals ADD COLUMN vat_rate TEXT")
|
||
if not getattr(_local, "in_transaction", False):
|
||
conn.commit()
|
||
|
||
|
||
def cleanup_db():
|
||
"""Explicitly delete the database and recreate an empty schema."""
|
||
global _db_session_key
|
||
conn = getattr(_local, "conn", None)
|
||
if conn is not None:
|
||
conn.close()
|
||
_local.conn = None
|
||
_local.session_key = None
|
||
_db_session_key = time.time_ns()
|
||
for path in (DB_PATH, DB_PATH + "-wal", DB_PATH + "-shm"):
|
||
try:
|
||
os.remove(path)
|
||
except FileNotFoundError:
|
||
pass
|
||
return init_db()
|
||
|
||
|
||
def get_conn():
|
||
"""Thread-local соединение. Пересоздаётся при смене сессии."""
|
||
global _db_session_key
|
||
|
||
conn = getattr(_local, "conn", None)
|
||
local_key = getattr(_local, "session_key", None)
|
||
|
||
if conn is None or local_key != _db_session_key:
|
||
if conn is not None:
|
||
try:
|
||
conn.close()
|
||
except Exception:
|
||
pass
|
||
conn = sqlite3.connect(DB_PATH, timeout=5, check_same_thread=False)
|
||
conn.row_factory = sqlite3.Row
|
||
conn.execute("PRAGMA journal_mode=WAL")
|
||
conn.execute("PRAGMA busy_timeout=5000")
|
||
conn.execute("PRAGMA foreign_keys=ON")
|
||
_local.conn = conn
|
||
_local.session_key = _db_session_key
|
||
|
||
return conn
|
||
|
||
|
||
def query(sql, params=None):
|
||
"""SELECT → list[dict]."""
|
||
conn = get_conn()
|
||
sql = _pg_to_sqlite(sql)
|
||
cur = conn.execute(sql, params or [])
|
||
return [dict(r) for r in cur.fetchall()]
|
||
|
||
|
||
def execute(sql, params=None):
|
||
"""INSERT/UPDATE/DELETE → rowcount."""
|
||
conn = get_conn()
|
||
sql = _pg_to_sqlite(sql)
|
||
cur = conn.execute(sql, params or [])
|
||
if not getattr(_local, "in_transaction", False):
|
||
conn.commit()
|
||
return cur.rowcount
|
||
|
||
|
||
def execute_returning(sql, params=None):
|
||
"""INSERT с RETURNING → dict (эмулируется через lastrowid)."""
|
||
conn = get_conn()
|
||
table = _extract_table(sql)
|
||
sql = _pg_to_sqlite(sql)
|
||
|
||
if " RETURNING " in sql.upper():
|
||
sql = sql[:sql.upper().rfind(" RETURNING ")]
|
||
|
||
cur = conn.execute(sql, params or [])
|
||
# The outer transaction owns commit and rollback; committing here would make
|
||
# its rollback ineffective and could expose only part of a multi-write update.
|
||
if not getattr(_local, "in_transaction", False):
|
||
conn.commit()
|
||
rowid = cur.lastrowid
|
||
|
||
if table and rowid:
|
||
row = conn.execute(f"SELECT * FROM {table} WHERE rowid = ?", (rowid,)).fetchone()
|
||
if row:
|
||
return dict(row)
|
||
|
||
# Fallback: если нет таблицы или rowid — просто вернуть последнюю строку
|
||
if table:
|
||
row = conn.execute(f"SELECT * FROM {table} ORDER BY rowid DESC LIMIT 1").fetchone()
|
||
return dict(row) if row else None
|
||
return None
|
||
|
||
|
||
def _pg_to_sqlite(sql):
|
||
"""Конвертировать PostgreSQL-специфичный SQL в SQLite."""
|
||
# %s → ?
|
||
sql = sql.replace("%s", "?")
|
||
# ::jsonb → убрать (SQLite не типизирует)
|
||
sql = sql.replace("::jsonb", "")
|
||
# BOOLEAN → INTEGER
|
||
sql = sql.replace(" BOOLEAN ", " INTEGER ")
|
||
sql = sql.replace(" bool ", " INTEGER ")
|
||
# ILIKE → LIKE (SQLite LIKE case-insensitive для ASCII)
|
||
sql = sql.replace(" ILIKE ", " LIKE ")
|
||
# FALSE/TRUE → 0/1
|
||
sql = sql.replace(" FALSE", " 0").replace(" TRUE", " 1")
|
||
sql = sql.replace(" false", " 0").replace(" true", " 1")
|
||
return sql
|
||
|
||
|
||
def _extract_table(sql):
|
||
"""Извлечь имя таблицы из INSERT INTO <table>."""
|
||
import re
|
||
m = re.search(r"INSERT\s+INTO\s+(\w+)", sql, re.IGNORECASE)
|
||
return m.group(1) if m else None
|