229 lines
7.6 KiB
Python
229 lines
7.6 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
|
||
|
||
DB_PATH = "/tmp/contracts.db"
|
||
_db_session_key = None
|
||
_local = threading.local()
|
||
|
||
|
||
def init_db():
|
||
"""Создать/пересоздать БД + схему. Вызывается при старте и после cleanup."""
|
||
global _db_session_key
|
||
_db_session_key = time.time()
|
||
|
||
# Удалить старый файл + WAL-сателлиты
|
||
for f in (DB_PATH, DB_PATH + "-wal", DB_PATH + "-shm"):
|
||
try:
|
||
os.remove(f)
|
||
except FileNotFoundError:
|
||
pass
|
||
|
||
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 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 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);
|
||
""")
|
||
|
||
|
||
def cleanup_db():
|
||
"""Удалить БД полностью. Вызывает init_db() для создания новой."""
|
||
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 [])
|
||
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 [])
|
||
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", "")
|
||
# gen_random_uuid() → хекс-UUID через randomblob
|
||
if "gen_random_uuid()" in sql:
|
||
import uuid
|
||
sql = sql.replace("gen_random_uuid()", "?")
|
||
# 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
|