228 lines
7.3 KiB
Python
228 lines
7.3 KiB
Python
"""
|
||
Слой доступа к данным (SQLite).
|
||
|
||
Использует thread-local соединения — каждый поток получает своё соединение
|
||
и переиспользует его, вместо создания нового на каждый запрос.
|
||
"""
|
||
import json
|
||
import sqlite3
|
||
import threading
|
||
|
||
import config
|
||
|
||
_local = threading.local()
|
||
|
||
|
||
def _connect() -> sqlite3.Connection:
|
||
"""Возвращает соединение для текущего потока (создаёт при первом вызове)."""
|
||
conn = getattr(_local, "conn", None)
|
||
if conn is None:
|
||
conn = sqlite3.connect(config.DB_FILE, check_same_thread=False)
|
||
conn.row_factory = sqlite3.Row
|
||
_local.conn = conn
|
||
return conn
|
||
|
||
|
||
def init_db():
|
||
with _connect() as conn:
|
||
conn.execute("""
|
||
CREATE TABLE IF NOT EXISTS messages (
|
||
task_id INTEGER PRIMARY KEY,
|
||
message_id INTEGER NOT NULL,
|
||
task_number TEXT,
|
||
group_id INTEGER DEFAULT 0,
|
||
base_text TEXT
|
||
)
|
||
""")
|
||
for col, dfn in [
|
||
("group_id", "INTEGER DEFAULT 0"),
|
||
("base_text", "TEXT"),
|
||
("events_json", "TEXT"),
|
||
]:
|
||
try:
|
||
conn.execute(f"ALTER TABLE messages ADD COLUMN {col} {dfn}")
|
||
except Exception:
|
||
pass
|
||
|
||
conn.execute("""
|
||
CREATE TABLE IF NOT EXISTS ref_data (
|
||
entity TEXT NOT NULL,
|
||
id INTEGER NOT NULL,
|
||
name TEXT NOT NULL,
|
||
PRIMARY KEY (entity, id)
|
||
)
|
||
""")
|
||
|
||
conn.execute("""
|
||
CREATE TABLE IF NOT EXISTS responsible (
|
||
task_id INTEGER NOT NULL,
|
||
user_id INTEGER NOT NULL,
|
||
user_name TEXT NOT NULL,
|
||
PRIMARY KEY (task_id, user_id)
|
||
)
|
||
""")
|
||
|
||
conn.execute("""
|
||
CREATE TABLE IF NOT EXISTS gif_messages (
|
||
task_id INTEGER PRIMARY KEY,
|
||
message_id INTEGER NOT NULL
|
||
)
|
||
""")
|
||
try:
|
||
conn.execute("SELECT 1 FROM gif_messages LIMIT 1")
|
||
except Exception:
|
||
conn.execute("""
|
||
CREATE TABLE gif_messages (
|
||
task_id INTEGER PRIMARY KEY,
|
||
message_id INTEGER NOT NULL
|
||
)
|
||
""")
|
||
conn.commit()
|
||
|
||
|
||
# ── messages ──────────────────────────────────────────────────────────────────
|
||
|
||
def get_message_id(task_id: int):
|
||
row = _connect().execute(
|
||
"SELECT message_id FROM messages WHERE task_id = ?", (task_id,)
|
||
).fetchone()
|
||
return row["message_id"] if row else None
|
||
|
||
|
||
def get_message_data(task_id: int) -> dict | None:
|
||
row = _connect().execute(
|
||
"SELECT message_id, task_number, group_id, base_text FROM messages WHERE task_id = ?",
|
||
(task_id,),
|
||
).fetchone()
|
||
return dict(row) if row else None
|
||
|
||
|
||
def get_task_id_by_number(task_number: str) -> int | None:
|
||
row = _connect().execute(
|
||
"SELECT task_id FROM messages WHERE task_number = ?", (str(task_number),)
|
||
).fetchone()
|
||
return row["task_id"] if row else None
|
||
|
||
|
||
def save_message_id(
|
||
task_id: int,
|
||
message_id: int,
|
||
task_number: str = None,
|
||
group_id: int = 0,
|
||
base_text: str = None,
|
||
):
|
||
conn = _connect()
|
||
row = conn.execute(
|
||
"SELECT events_json FROM messages WHERE task_id = ?", (task_id,)
|
||
).fetchone()
|
||
existing_events = row["events_json"] if row else None
|
||
|
||
conn.execute(
|
||
"""INSERT OR REPLACE INTO messages
|
||
(task_id, message_id, task_number, group_id, base_text, events_json)
|
||
VALUES (?, ?, ?, ?, ?, ?)""",
|
||
(task_id, message_id, task_number, group_id or 0, base_text, existing_events),
|
||
)
|
||
conn.commit()
|
||
|
||
|
||
def get_stored_events(task_id: int) -> list[str]:
|
||
row = _connect().execute(
|
||
"SELECT events_json FROM messages WHERE task_id = ?", (task_id,)
|
||
).fetchone()
|
||
if not row or not row["events_json"]:
|
||
return []
|
||
try:
|
||
return json.loads(row["events_json"])
|
||
except Exception:
|
||
return []
|
||
|
||
|
||
def save_stored_events(task_id: int, events: list[str]):
|
||
conn = _connect()
|
||
conn.execute(
|
||
"UPDATE messages SET events_json = ? WHERE task_id = ?",
|
||
(json.dumps(events, ensure_ascii=False), task_id),
|
||
)
|
||
conn.commit()
|
||
|
||
|
||
# ── ref_data ──────────────────────────────────────────────────────────────────
|
||
|
||
def save_ref_data(entity: str, items: list[tuple[int, str]]):
|
||
conn = _connect()
|
||
conn.executemany(
|
||
"INSERT OR REPLACE INTO ref_data (entity, id, name) VALUES (?, ?, ?)",
|
||
[(entity, item_id, name) for item_id, name in items],
|
||
)
|
||
conn.commit()
|
||
|
||
|
||
def resolve(entity: str, item_id: int | None) -> str | None:
|
||
if item_id is None:
|
||
return None
|
||
row = _connect().execute(
|
||
"SELECT name FROM ref_data WHERE entity = ? AND id = ?",
|
||
(entity, item_id),
|
||
).fetchone()
|
||
return row["name"] if row else None
|
||
|
||
|
||
def get_all_ref(entity: str) -> list[dict]:
|
||
rows = _connect().execute(
|
||
"SELECT id, name FROM ref_data WHERE entity = ? ORDER BY name",
|
||
(entity,),
|
||
).fetchall()
|
||
return [{"id": r["id"], "name": r["name"]} for r in rows]
|
||
|
||
|
||
# ── responsible ───────────────────────────────────────────────────────────────
|
||
|
||
def add_responsible(task_id: int, user_id: int, user_name: str):
|
||
conn = _connect()
|
||
conn.execute(
|
||
"INSERT OR REPLACE INTO responsible (task_id, user_id, user_name) VALUES (?, ?, ?)",
|
||
(task_id, user_id, user_name),
|
||
)
|
||
conn.commit()
|
||
|
||
|
||
def remove_responsible(task_id: int, user_id: int):
|
||
conn = _connect()
|
||
conn.execute(
|
||
"DELETE FROM responsible WHERE task_id = ? AND user_id = ?",
|
||
(task_id, user_id),
|
||
)
|
||
conn.commit()
|
||
|
||
|
||
def get_responsible(task_id: int) -> list[dict]:
|
||
rows = _connect().execute(
|
||
"SELECT user_id, user_name FROM responsible WHERE task_id = ? ORDER BY rowid",
|
||
(task_id,),
|
||
).fetchall()
|
||
return [{"user_id": r["user_id"], "user_name": r["user_name"]} for r in rows]
|
||
|
||
|
||
# ── gif_messages ──────────────────────────────────────────────────────────────
|
||
|
||
def get_gif_message_id(task_id: int) -> int | None:
|
||
row = _connect().execute(
|
||
"SELECT message_id FROM gif_messages WHERE task_id = ?", (task_id,)
|
||
).fetchone()
|
||
return row["message_id"] if row else None
|
||
|
||
|
||
def save_gif_message_id(task_id: int, message_id: int):
|
||
conn = _connect()
|
||
conn.execute(
|
||
"INSERT OR REPLACE INTO gif_messages (task_id, message_id) VALUES (?, ?)",
|
||
(task_id, message_id),
|
||
)
|
||
conn.commit()
|
||
|
||
|
||
def delete_gif_message_id(task_id: int):
|
||
conn = _connect()
|
||
conn.execute("DELETE FROM gif_messages WHERE task_id = ?", (task_id,))
|
||
conn.commit()
|