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 TABLE IF NOT EXISTS users (
                id            INTEGER PRIMARY KEY AUTOINCREMENT,
                username      TEXT NOT NULL UNIQUE,
                password_hash TEXT NOT NULL,
                role          TEXT NOT NULL,
                is_active     INTEGER NOT NULL DEFAULT 1,
                created_at    TEXT NOT NULL DEFAULT (datetime('now'))
            );

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

        for col_def in [
            "tput_mbps REAL",
            "dl_rssi_dbm REAL",
            "ul_rssi_dbm REAL",
            "tput_results TEXT",
            "throttled TEXT",
            "coe_pair TEXT",
            "p3p_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 set_p3p_pair(test_row_id, p3p_pair):
    json_value = json.dumps(p3p_pair or [])
    with _tx():
        _conn.execute(
            """
            UPDATE tests
            SET p3p_pair = ?
            WHERE id = ?
            """,
            (json_value, test_row_id),
        )


def update_all_p3p_pairs_sql():
    with _tx():
        cursor = _conn.execute(
            """
            UPDATE tests AS t
            SET p3p_pair = (
                SELECT CASE
                    WHEN COUNT(*) = 0 THEN '[]'
                    ELSE '[' || GROUP_CONCAT('"' || m.device || '_' || m.test_id || '"') || ']'
                END
                FROM tests AS m
                WHERE m.interference = 'P3P'
                  AND UPPER(m.device) = UPPER(t.device)
                  AND UPPER(m.test_id) = CASE
                      WHEN INSTR(UPPER(t.test_id), 'TH') > 0 THEN REPLACE(UPPER(t.test_id), 'TH', 'UT')
                      WHEN INSTR(UPPER(t.test_id), 'UT') > 0 THEN REPLACE(UPPER(t.test_id), 'UT', 'TH')
                      ELSE '__NO_MATCH__'
                  END
            )
            WHERE t.interference = 'P3P'
            """
        )
    return cursor.rowcount


def update_all_p2p_coe_pairs_sql():
    with _tx():
        cursor = _conn.execute(
            """
            UPDATE tests AS t
            SET coe_pair = (
                SELECT CASE
                    WHEN COUNT(*) = 0 THEN '[]'
                    ELSE '[' || GROUP_CONCAT('"' || pair_value || '"') || ']'
                END
                FROM (
                    SELECT DISTINCT m.device || '_' || m.test_id AS pair_value
                    FROM tests AS m
                    WHERE m.interference = 'COE'
                      AND IFNULL(m.device, '') = IFNULL(t.device, '')
                      AND IFNULL(m.rotation, '') = IFNULL(t.rotation, '')
                      AND IFNULL(m.test_point, '') = IFNULL(t.test_point, '')
                      AND IFNULL(m.rssi, '') = IFNULL(t.rssi, '')
                      AND IFNULL(m.station, '') = IFNULL(t.station, '')
                      AND IFNULL(m.band, '') = IFNULL(t.band, '')
                      AND IFNULL(m.channel, '') = IFNULL(t.channel, '')
                      AND IFNULL(m.bandwidth, '') = IFNULL(t.bandwidth, '')
                      AND IFNULL(m.direction, '') = IFNULL(t.direction, '')
                    ORDER BY pair_value
                )
            )
            WHERE t.interference = 'P2P'
            """
        )
    return cursor.rowcount


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 reset_all_results_state():
    with _tx():
        _conn.execute(
            """
            UPDATE tests
            SET completed = 0,
                completed_at = NULL,
                duration_seconds = NULL,
                tput_results = NULL
            """
        )


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"]


def get_user_by_username(username):
    with _lock:
        row = _conn.execute(
            """
            SELECT id, username, password_hash, role, is_active, created_at
            FROM users
            WHERE username = ?
            """,
            (username,),
        ).fetchone()
    return dict(row) if row else None


def create_user(username, password_hash, role="viewer", is_active=1):
    with _tx():
        cursor = _conn.execute(
            """
            INSERT INTO users (username, password_hash, role, is_active)
            VALUES (?, ?, ?, ?)
            """,
            (username, password_hash, role, is_active),
        )
    return cursor.lastrowid


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


def get_user_by_id(user_id):
    with _lock:
        row = _conn.execute(
            """
            SELECT id, username, password_hash, role, is_active, created_at
            FROM users
            WHERE id = ?
            """,
            (user_id,),
        ).fetchone()
    return dict(row) if row else None


def get_all_users():
    with _lock:
        rows = _conn.execute(
            """
            SELECT id, username, role, is_active, created_at
            FROM users
            ORDER BY username ASC
            """
        ).fetchall()
    return [dict(row) for row in rows]

def clear_users():
    with _tx():
        _conn.execute("DELETE FROM users")


_init_db()