tradingview-alert-dispatcher/app/db.py

from __future__ import annotations

import json
import os
import sqlite3
from contextlib import contextmanager
from datetime import datetime, timedelta, timezone
from typing import Any, Iterator

from app.auth import hash_password
from app.config import Settings


UTC = timezone.utc


def now_iso() -> str:
    return datetime.now(UTC).replace(microsecond=0).isoformat()


def to_json(value: Any) -> str:
    return json.dumps(value, ensure_ascii=False, separators=(",", ":"))


def from_json(value: str | None, default: Any = None) -> Any:
    if not value:
        return default
    return json.loads(value)


class Database:
    def __init__(self, settings: Settings):
        self.path = settings.database_path
        os.makedirs(os.path.dirname(self.path) or ".", exist_ok=True)

    @contextmanager
    def connect(self) -> Iterator[sqlite3.Connection]:
        conn = sqlite3.connect(self.path)
        conn.row_factory = sqlite3.Row
        conn.execute("PRAGMA foreign_keys = ON")
        try:
            yield conn
            conn.commit()
        except Exception:
            conn.rollback()
            raise
        finally:
            conn.close()

    def migrate(self, settings: Settings) -> None:
        with self.connect() as conn:
            conn.executescript(
                """
                CREATE TABLE IF NOT EXISTS admin_settings (
                    id INTEGER PRIMARY KEY CHECK (id = 1),
                    password_hash TEXT NOT NULL,
                    created_at TEXT NOT NULL,
                    updated_at TEXT NOT NULL
                );

                CREATE TABLE IF NOT EXISTS webhook_targets (
                    id INTEGER PRIMARY KEY AUTOINCREMENT,
                    name TEXT NOT NULL,
                    webhook_url TEXT NOT NULL,
                    enabled INTEGER NOT NULL DEFAULT 1,
                    created_at TEXT NOT NULL,
                    updated_at TEXT NOT NULL
                );

                CREATE TABLE IF NOT EXISTS routing_rules (
                    id INTEGER PRIMARY KEY AUTOINCREMENT,
                    name TEXT NOT NULL,
                    timeframe TEXT NOT NULL,
                    symbol TEXT NOT NULL,
                    strategy TEXT NOT NULL,
                    priority INTEGER NOT NULL DEFAULT 100,
                    enabled INTEGER NOT NULL DEFAULT 1,
                    target_ids TEXT NOT NULL,
                    created_at TEXT NOT NULL,
                    updated_at TEXT NOT NULL
                );

                CREATE INDEX IF NOT EXISTS idx_rules_match
                    ON routing_rules(enabled, timeframe, symbol, strategy, priority);

                CREATE TABLE IF NOT EXISTS alerts (
                    id INTEGER PRIMARY KEY AUTOINCREMENT,
                    timeframe TEXT NOT NULL,
                    symbol TEXT NOT NULL,
                    strategy TEXT NOT NULL,
                    action TEXT,
                    price REAL,
                    payload TEXT NOT NULL,
                    matched_rule_id INTEGER,
                    status TEXT NOT NULL,
                    error TEXT,
                    created_at TEXT NOT NULL,
                    FOREIGN KEY(matched_rule_id) REFERENCES routing_rules(id) ON DELETE SET NULL
                );

                CREATE TABLE IF NOT EXISTS deliveries (
                    id INTEGER PRIMARY KEY AUTOINCREMENT,
                    alert_id INTEGER NOT NULL,
                    rule_id INTEGER,
                    target_id INTEGER,
                    target_name TEXT NOT NULL,
                    webhook_url TEXT NOT NULL,
                    status TEXT NOT NULL,
                    attempts INTEGER NOT NULL DEFAULT 0,
                    next_attempt_at TEXT,
                    last_attempt_at TEXT,
                    response_code INTEGER,
                    response_body TEXT,
                    error TEXT,
                    created_at TEXT NOT NULL,
                    updated_at TEXT NOT NULL,
                    FOREIGN KEY(alert_id) REFERENCES alerts(id) ON DELETE CASCADE,
                    FOREIGN KEY(rule_id) REFERENCES routing_rules(id) ON DELETE SET NULL,
                    FOREIGN KEY(target_id) REFERENCES webhook_targets(id) ON DELETE SET NULL
                );

                CREATE INDEX IF NOT EXISTS idx_deliveries_retry
                    ON deliveries(status, next_attempt_at);
                """
            )
            existing_columns = {
                row["name"] for row in conn.execute("PRAGMA table_info(routing_rules)").fetchall()
            }
            if "message_type" not in existing_columns:
                conn.execute("ALTER TABLE routing_rules ADD COLUMN message_type TEXT NOT NULL DEFAULT 'card'")
            if "card_title_template" not in existing_columns:
                conn.execute(
                    "ALTER TABLE routing_rules ADD COLUMN card_title_template TEXT NOT NULL DEFAULT 'TradingView {{symbol}} {{action}}'"
                )
            if "card_body_template" not in existing_columns:
                conn.execute(
                    "ALTER TABLE routing_rules ADD COLUMN card_body_template TEXT NOT NULL DEFAULT '{{symbol}} {{timeframe}} {{strategy}} {{action}} @ {{price}}'"
                )
            admin = conn.execute("SELECT id FROM admin_settings WHERE id = 1").fetchone()
            if not admin:
                now = now_iso()
                conn.execute(
                    "INSERT INTO admin_settings (id, password_hash, created_at, updated_at) VALUES (1, ?, ?, ?)",
                    (hash_password(settings.admin_password), now, now),
                )

    def cleanup_old_logs(self, retention_days: int) -> int:
        cutoff = (datetime.now(UTC) - timedelta(days=retention_days)).replace(microsecond=0).isoformat()
        with self.connect() as conn:
            cur = conn.execute("DELETE FROM alerts WHERE created_at < ?", (cutoff,))
            return cur.rowcount