test_dashboard/server/db_py.py

import json
import os
import sqlite3
import threading
from contextlib import contextmanager

BASE_DIR = os.path.dirname(os.path.abspath(__file__))
DB_PATH = os.path.join(BASE_DIR, "dashboard.db")
CONFIG_JSON = os.path.join(BASE_DIR, "config.json")

_conn = sqlite3.connect(DB_PATH, check_same_thread=False)
_conn.row_factory = sqlite3.Row
_lock = threading.RLock()


@contextmanager
def _tx():
    with _lock:
        try:
            yield
            _conn.commit()
        except Exception:
            _conn.rollback()
            raise


def _init_db():
    with _tx():
        _conn.execute("PRAGMA journal_mode=WAL;")
        _conn.executescript(
            """
            CREATE TABLE IF NOT EXISTS config (
                key   TEXT PRIMARY KEY,
                value TEXT NOT NULL
            );

            CREATE TABLE IF NOT EXISTS tests (
                id               TEXT PRIMARY KEY,
                test_id          TEXT,
                parent_dir       TEXT,
                filename         TEXT,
                interference     TEXT,
                device           TEXT,
                rotation         TEXT,
                test_point       TEXT,
                station          TEXT,
                band             TEXT,
                channel          TEXT,
                bandwidth        TEXT,
                rssi             TEXT,
                direction        TEXT,
                completed        INTEGER NOT NULL DEFAULT 0,
                completed_at     TEXT,
                duration_seconds REAL
            );

            CREATE INDEX IF NOT EXISTS idx_tests_test_id ON tests (test_id);
            CREATE INDEX IF NOT EXISTS idx_tests_device ON tests (device);
            """
        )

        for col_def in [
            "tput_mbps REAL",
            "dl_rssi_dbm REAL",
            "ul_rssi_dbm REAL",
            "tput_results TEXT",
            "throttled TEXT",
            "coe_pair TEXT",
        ]:
            try:
                _conn.execute(f"ALTER TABLE tests ADD COLUMN {col_def}")
            except sqlite3.OperationalError:
                pass


def get_config(key):
    with _lock:
        row = _conn.execute("SELECT value FROM config WHERE key = ?", (key,)).fetchone()
    return row["value"] if row else None


def set_config(key, value):
    with _tx():
        _conn.execute(
            "INSERT OR REPLACE INTO config (key, value) VALUES (?, ?)",
            (key, value),
        )


def del_config(key):
    with _tx():
        _conn.execute("DELETE FROM config WHERE key = ?", (key,))


def upsert_test(test):
    with _tx():
        _conn.execute(
            """
            INSERT INTO tests
                (id, test_id, parent_dir, filename, interference, device, rotation,
                 test_point, station, band, channel, bandwidth, rssi, direction, throttled)
            VALUES
                (:id, :test_id, :parent_dir, :filename, :interference, :device, :rotation,
                 :test_point, :station, :band, :channel, :bandwidth, :rssi, :direction, :throttled)
            ON CONFLICT(id) DO UPDATE SET
                test_id      = excluded.test_id,
                parent_dir   = excluded.parent_dir,
                filename     = excluded.filename,
                interference = excluded.interference,
                device       = excluded.device,
                rotation     = excluded.rotation,
                test_point   = excluded.test_point,
                station      = excluded.station,
                band         = excluded.band,
                channel      = excluded.channel,
                bandwidth    = excluded.bandwidth,
                rssi         = excluded.rssi,
                direction    = excluded.direction,
                throttled    = excluded.throttled
            """,
            test,
        )


def mark_completed(test_id, device, completed_at, duration_seconds, tput_results=None):
    json_value = json.dumps(tput_results) if tput_results else None
    with _tx():
        _conn.execute(
            """
            UPDATE tests
            SET completed = 1,
                completed_at = ?,
                duration_seconds = ?,
                tput_results = ?
            WHERE test_id = ?
              AND (? IS NULL OR device = ?)
            """,
            (completed_at, duration_seconds, json_value, test_id, device, device),
        )


def set_coe_pair(test_row_id, coe_pair):
    json_value = json.dumps(coe_pair or [])
    with _tx():
        _conn.execute(
            """
            UPDATE tests
            SET coe_pair = ?
            WHERE id = ?
            """,
            (json_value, test_row_id),
        )


def get_station_for_test(test_id, device):
    with _lock:
        row = _conn.execute(
            "SELECT station FROM tests WHERE test_id = ? AND device = ? LIMIT 1",
            (test_id, device),
        ).fetchone()
    return row["station"] if row else None


def reset_by_file_id_and_device(test_id, device):
    with _tx():
        _conn.execute(
            """
            UPDATE tests
            SET completed = 0,
                completed_at = NULL,
                duration_seconds = NULL,
                tput_results = NULL
            WHERE test_id = ?
              AND (? IS NULL OR device = ?)
            """,
            (test_id, device, device),
        )


def clear_tests():
    with _tx():
        _conn.execute("DELETE FROM tests")


def get_all_tests():
    with _lock:
        rows = _conn.execute("SELECT * FROM tests").fetchall()
    return [dict(row) for row in rows]


def count_tests():
    with _lock:
        row = _conn.execute("SELECT COUNT(*) AS n FROM tests").fetchone()
    return row["n"]


_init_db()