58 lines
2.0 KiB
Python
58 lines
2.0 KiB
Python
from contextlib import closing
|
|
import os
|
|
import sqlite3
|
|
from pathlib import Path
|
|
|
|
|
|
DEFAULT_DB_PATH = "/app/data/social_writer.sqlite3"
|
|
|
|
|
|
def connect(db_path: str | None = None) -> sqlite3.Connection:
|
|
path = Path(db_path or os.environ.get("BOT_DB_PATH", DEFAULT_DB_PATH))
|
|
path.parent.mkdir(parents=True, exist_ok=True)
|
|
connection = sqlite3.connect(path, timeout=10)
|
|
connection.row_factory = sqlite3.Row
|
|
connection.execute("PRAGMA busy_timeout = 10000")
|
|
connection.execute("PRAGMA journal_mode = WAL")
|
|
return connection
|
|
|
|
|
|
def init_db(db_path: str | None = None) -> None:
|
|
with closing(connect(db_path)) as connection, connection:
|
|
connection.execute(
|
|
"""
|
|
CREATE TABLE IF NOT EXISTS users (
|
|
telegram_user_id INTEGER PRIMARY KEY,
|
|
tone_of_voice TEXT NOT NULL DEFAULT '',
|
|
created_at TEXT NOT NULL DEFAULT CURRENT_TIMESTAMP,
|
|
updated_at TEXT NOT NULL DEFAULT CURRENT_TIMESTAMP
|
|
)
|
|
"""
|
|
)
|
|
|
|
|
|
def get_tone_of_voice(user_id: int, db_path: str | None = None) -> str:
|
|
with closing(connect(db_path)) as connection, connection:
|
|
connection.execute(
|
|
"INSERT OR IGNORE INTO users (telegram_user_id) VALUES (?)", (user_id,)
|
|
)
|
|
row = connection.execute(
|
|
"SELECT tone_of_voice FROM users WHERE telegram_user_id = ?", (user_id,)
|
|
).fetchone()
|
|
return str(row["tone_of_voice"] or "") if row else ""
|
|
|
|
|
|
def save_tone_of_voice(user_id: int, tone_of_voice: str, db_path: str | None = None) -> None:
|
|
tone = tone_of_voice.strip()[:1200]
|
|
with closing(connect(db_path)) as connection, connection:
|
|
connection.execute(
|
|
"""
|
|
INSERT INTO users (telegram_user_id, tone_of_voice, updated_at)
|
|
VALUES (?, ?, CURRENT_TIMESTAMP)
|
|
ON CONFLICT(telegram_user_id) DO UPDATE SET
|
|
tone_of_voice = excluded.tone_of_voice,
|
|
updated_at = CURRENT_TIMESTAMP
|
|
""",
|
|
(user_id, tone),
|
|
)
|