Files
scheduler/backend/db.py
T
2026-06-23 12:07:43 -04:00

556 lines
16 KiB
Python

import json
import sqlite3
import os
import re
from contextlib import contextmanager
from dataclasses import dataclass
from pathlib import Path
from typing import Any, Iterator
from graph import build_station_testpoint_map
APP_ROOT = Path(__file__).resolve().parent
DB_PATH = APP_ROOT / "scheduler.db"
# Hardware device names (from environment or defaults)
DUT = os.getenv("DUT", "CGW453").strip()
REF = os.getenv("REF", "CGW452").strip()
# Device names for the test database (use hardware device names)
DEVICE_DUT = DUT
DEVICE_REF = REF
@dataclass(frozen=True)
class TestRecord:
test_id: str
device: str
test_type: str
rotation: str | None
rx_tx: str | None
has_coe_pair: bool
coe_pairing: list[str]
priority: int
victim_band: str | None
config: dict[str, dict[str, str | None]]
throttled: bool
estimated_minutes: int
status: str = "pending"
excluded: bool = False
raw_payload: dict[str, Any] | None = None
station_testpoint_map: str | None = None
@dataclass(frozen=True)
class ScheduleRow:
test_id: str
device: str
scheduled_date: str
shift_index: int
sequence_in_shift: int
test_type: str
rotation: str | None
config: dict[str, dict[str, str | None]]
status: str
priority: int
estimated_minutes: int
@contextmanager
def get_connection(db_path: str | Path = DB_PATH) -> Iterator[sqlite3.Connection]:
conn = sqlite3.connect(str(db_path))
conn.row_factory = sqlite3.Row
try:
yield conn
conn.commit()
finally:
conn.close()
def init_db(db_path: str | Path = DB_PATH) -> None:
# Validate device names are set
if not DUT or not REF:
raise ValueError(f"Invalid device names: DUT={DUT!r}, REF={REF!r}")
with get_connection(db_path) as conn:
conn.executescript(
"""
CREATE TABLE IF NOT EXISTS tests (
id INTEGER PRIMARY KEY AUTOINCREMENT,
test_id TEXT NOT NULL,
device TEXT NOT NULL,
test_type TEXT NOT NULL CHECK (test_type IN ('P2P', 'COE', 'P3P')),
rotation TEXT,
rx_tx TEXT CHECK (rx_tx IN ('RX', 'TX') OR rx_tx IS NULL),
has_coe_pair INTEGER NOT NULL DEFAULT 0,
coe_pairing_json TEXT,
priority INTEGER NOT NULL CHECK (priority BETWEEN 1 AND 5),
victim_band TEXT,
config_json TEXT,
throttled INTEGER NOT NULL DEFAULT 0,
estimated_minutes INTEGER NOT NULL,
status TEXT NOT NULL CHECK (status IN ('pending', 'completed', 'failed', 'invalid')),
excluded INTEGER NOT NULL DEFAULT 0,
raw_payload TEXT,
created_at TEXT NOT NULL DEFAULT CURRENT_TIMESTAMP,
updated_at TEXT NOT NULL DEFAULT CURRENT_TIMESTAMP,
UNIQUE(test_id, device)
);
CREATE TABLE IF NOT EXISTS schedules (
id INTEGER PRIMARY KEY AUTOINCREMENT,
schedule_version INTEGER NOT NULL,
test_id TEXT NOT NULL,
device TEXT NOT NULL,
scheduled_date TEXT NOT NULL,
shift_index INTEGER NOT NULL CHECK (shift_index IN (1, 2, 3)),
sequence_in_shift INTEGER NOT NULL,
created_at TEXT NOT NULL DEFAULT CURRENT_TIMESTAMP,
FOREIGN KEY(test_id, device) REFERENCES tests(test_id, device)
);
CREATE TABLE IF NOT EXISTS settings (
key TEXT PRIMARY KEY,
value_json TEXT NOT NULL,
updated_at TEXT NOT NULL DEFAULT CURRENT_TIMESTAMP
);
CREATE TABLE IF NOT EXISTS holidays (
date TEXT PRIMARY KEY,
note TEXT
);
CREATE TABLE IF NOT EXISTS runtime_defaults (
test_type TEXT PRIMARY KEY,
minutes INTEGER NOT NULL
);
CREATE TABLE IF NOT EXISTS rerun_events (
id INTEGER PRIMARY KEY AUTOINCREMENT,
detected_date TEXT NOT NULL,
failed_test_ids_json TEXT NOT NULL,
estimated_rerun_minutes INTEGER NOT NULL,
rerun_during_day INTEGER
);
CREATE INDEX IF NOT EXISTS idx_tests_status ON tests(status);
CREATE INDEX IF NOT EXISTS idx_tests_priority ON tests(priority);
CREATE INDEX IF NOT EXISTS idx_schedules_date_shift ON schedules(scheduled_date, shift_index);
"""
)
_ensure_column(conn, "tests", "has_coe_pair", "INTEGER NOT NULL DEFAULT 0")
_ensure_column(conn, "tests", "coe_pairing_json", "TEXT")
_ensure_column(conn, "tests", "config_json", "TEXT")
_ensure_column(conn, "tests", "victim_band", "TEXT")
_ensure_column(conn, "tests", "excluded", "INTEGER NOT NULL DEFAULT 0")
_ensure_column(conn, "tests", "throttled", "INTEGER NOT NULL DEFAULT 0")
_ensure_column(conn, "tests", "station_testpoint_map", "TEXT")
# Seed runtime defaults for schedule estimation.
conn.execute(
"""
INSERT INTO runtime_defaults(test_type, minutes) VALUES
('P2P', 80),
('COE', 115),
('P3P', 105)
ON CONFLICT(test_type) DO NOTHING
"""
)
def _serialize_station_testpoint_map(config: dict[str, dict[str, str | None]]) -> str:
"""Compute and serialize the station-to-testpoint map from a test config."""
station_map = build_station_testpoint_map(config)
return json.dumps(station_map)
def upsert_tests(records: list[TestRecord], db_path: str | Path = DB_PATH) -> int:
if not records:
return 0
with get_connection(db_path) as conn:
conn.executemany(
"""
INSERT INTO tests(
test_id, device, test_type, rotation, rx_tx, has_coe_pair,
coe_pairing_json, config_json, priority, victim_band, throttled, estimated_minutes, status, excluded, raw_payload, station_testpoint_map
)
VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)
ON CONFLICT(test_id, device) DO UPDATE SET
test_type = excluded.test_type,
device = excluded.device,
rotation = excluded.rotation,
rx_tx = excluded.rx_tx,
has_coe_pair = excluded.has_coe_pair,
coe_pairing_json = excluded.coe_pairing_json,
config_json = excluded.config_json,
priority = excluded.priority,
victim_band = excluded.victim_band,
throttled = excluded.throttled,
estimated_minutes = excluded.estimated_minutes,
status = excluded.status,
excluded = excluded.excluded,
raw_payload = excluded.raw_payload,
station_testpoint_map = excluded.station_testpoint_map,
updated_at = CURRENT_TIMESTAMP
""",
[
(
r.test_id,
r.device,
r.test_type,
r.rotation,
r.rx_tx,
int(r.has_coe_pair),
json.dumps(r.coe_pairing or []),
json.dumps(r.config),
r.priority,
r.victim_band,
int(r.throttled),
r.estimated_minutes,
r.status,
int(r.excluded),
json.dumps(r.raw_payload or {}),
_serialize_station_testpoint_map(r.config),
)
for r in records
],
)
return len(records)
def _ensure_column(conn: sqlite3.Connection, table: str, column: str, column_type: str) -> None:
existing = conn.execute(f"PRAGMA table_info({table})").fetchall()
column_names = {row[1] for row in existing}
if column not in column_names:
conn.execute(f"ALTER TABLE {table} ADD COLUMN {column} {column_type}")
def save_settings(payload: dict[str, Any], db_path: str | Path = DB_PATH) -> None:
with get_connection(db_path) as conn:
for key, value in payload.items():
conn.execute(
"""
INSERT INTO settings(key, value_json, updated_at)
VALUES (?, ?, CURRENT_TIMESTAMP)
ON CONFLICT(key) DO UPDATE SET
value_json = excluded.value_json,
updated_at = CURRENT_TIMESTAMP
""",
(key, json.dumps(value)),
)
def read_settings(db_path: str | Path = DB_PATH) -> dict[str, Any]:
with get_connection(db_path) as conn:
rows = conn.execute("SELECT key, value_json FROM settings").fetchall()
return {row["key"]: json.loads(row["value_json"]) for row in rows}
def _parse_rule_tokens(rule: str | None) -> list[str]:
if not rule:
return []
return [token.strip().upper() for token in rule.split(",") if token.strip()]
def _normalize_text(value: Any) -> str:
if value is None:
return ""
return " ".join(str(value).strip().upper().split())
def _normalize_compact(value: Any) -> str:
return re.sub(r"[^A-Z0-9+\-]", "", _normalize_text(value))
def _extract_signed_int(value: Any) -> str:
match = re.search(r"[-+]?\d+", _normalize_text(value))
return match.group(0) if match else ""
def _band_entry(config: dict[str, dict[str, str | None]], band: str) -> dict[str, str | None]:
return config.get(band) or config.get(band.lower()) or {}
def _entry_value(entry: dict[str, str | None], key: str) -> str:
if not entry:
return ""
return _normalize_text(entry.get(key) or entry.get(key.capitalize()) or entry.get(key.upper()))
def _record_power_mode(record: TestRecord) -> str:
power_mode = _entry_value(_band_entry(record.config, "6G"), "power_mode")
if power_mode:
return power_mode
raw_payload = record.raw_payload or {}
return _normalize_text(raw_payload.get("6GHz Power Mode"))
def _match_any_band_value(record: TestRecord, key: str, token_suffix: str) -> bool:
if not token_suffix:
return False
target_num = _extract_signed_int(token_suffix)
target_compact = _normalize_compact(token_suffix)
keys = ("Station 1", "Station 2", "Station 3") if record.test_type == "P3P" else ("5G", "6G", "2G")
for band in keys:
entry = _band_entry(record.config, band)
value = _entry_value(entry, key)
if not value:
continue
if target_num:
if _extract_signed_int(value) == target_num:
return True
elif target_compact and _normalize_compact(value) == target_compact:
return True
return False
def _matches_atomic_exclusion_token(record: TestRecord, token: str) -> bool:
if not token:
return False
device = _normalize_text(record.device)
test_type = _normalize_text(record.test_type)
rotation = _normalize_text(record.rotation)
victim_band = _normalize_text(record.victim_band)
power_mode = _record_power_mode(record)
known_devices = {_normalize_text(DUT), _normalize_text(REF)}
if token in {"LPI", "SP"}:
return power_mode == token
if token in known_devices:
return device == token
if token in {"P2P", "P3P", "COE"}:
return test_type == token
if token in {"2G", "5G", "6G"}:
return victim_band == token
if token.startswith("BW"):
return _match_any_band_value(record, "bandwidth", token[2:])
if token.startswith("CH"):
return _match_any_band_value(record, "channel", token[2:])
if token.startswith("RSSI"):
return _match_any_band_value(record, "rssi", token[4:])
if token.startswith("R") and len(token) > 1:
return rotation == token
return False
def _matches_exclusion_token(record: TestRecord, token: str) -> bool:
if not token:
return False
parts = [part for part in token.split("_") if part]
if len(parts) > 1:
# Treat underscore as logical AND across sub-tokens.
return all(_matches_atomic_exclusion_token(record, part) for part in parts)
return _matches_atomic_exclusion_token(record, token)
def _should_exclude_by_rule(record: TestRecord, tokens: list[str]) -> bool:
return any(_matches_exclusion_token(record, token) for token in tokens)
def list_schedulable_tests(db_path: str | Path = DB_PATH, rule: str = "") -> list[TestRecord]:
tokens = _parse_rule_tokens(rule)
with get_connection(db_path) as conn:
rows = conn.execute(
"""
SELECT
test_id, device, test_type, rotation, rx_tx, has_coe_pair,
config_json, victim_band,
coe_pairing_json, priority, throttled, estimated_minutes,
excluded, status, raw_payload
FROM tests
WHERE excluded = 0
AND status != 'completed'
"""
).fetchall()
results: list[TestRecord] = []
for row in rows:
results.append(
TestRecord(
test_id=row["test_id"],
device=row["device"],
test_type=row["test_type"],
rotation=row["rotation"],
rx_tx=row["rx_tx"],
has_coe_pair=bool(row["has_coe_pair"]),
coe_pairing=json.loads(row["coe_pairing_json"] or "[]"),
priority=int(row["priority"]),
victim_band=row["victim_band"],
config=json.loads(row["config_json"] or "{}"),
throttled=bool(row["throttled"]),
estimated_minutes=int(row["estimated_minutes"]),
status=row["status"],
excluded=bool(row["excluded"]),
raw_payload=json.loads(row["raw_payload"] or "{}"),
)
)
if not tokens:
return results
return [record for record in results if not _should_exclude_by_rule(record, tokens)]
# Get all tests for a specific device, regardless of exclusion or completion status. Used for schedule display and management.
def list_tests_for_device(device: str, db_path: str | Path = DB_PATH) -> list[TestRecord]:
with get_connection(db_path) as conn:
rows = conn.execute(
"""
SELECT
test_id, device, test_type, rotation, rx_tx, has_coe_pair,
config_json, victim_band,
coe_pairing_json, priority, throttled, estimated_minutes,
excluded, status, raw_payload, station_testpoint_map
FROM tests
WHERE device = ?
""",
(device,),
).fetchall()
results: list[TestRecord] = []
for row in rows:
results.append(
TestRecord(
test_id=row["test_id"],
device=row["device"],
test_type=row["test_type"],
rotation=row["rotation"],
rx_tx=row["rx_tx"],
has_coe_pair=bool(row["has_coe_pair"]),
coe_pairing=json.loads(row["coe_pairing_json"] or "[]"),
priority=int(row["priority"]),
victim_band=row["victim_band"],
config=json.loads(row["config_json"] or "{}"),
throttled=bool(row["throttled"]),
estimated_minutes=int(row["estimated_minutes"]),
status=row["status"],
excluded=bool(row["excluded"]),
raw_payload=json.loads(row["raw_payload"] or "{}"),
station_testpoint_map=row["station_testpoint_map"],
)
)
return results
def list_holidays(db_path: str | Path = DB_PATH) -> set[str]:
with get_connection(db_path) as conn:
rows = conn.execute("SELECT date FROM holidays").fetchall()
return {row["date"] for row in rows}
def create_schedule_version(
entries: list[tuple[str, str, str, int, int]],
db_path: str | Path = DB_PATH,
) -> int:
with get_connection(db_path) as conn:
row = conn.execute("SELECT COALESCE(MAX(schedule_version), 0) AS current FROM schedules").fetchone()
next_version = int(row["current"]) + 1
if entries:
conn.executemany(
"""
INSERT INTO schedules(schedule_version, test_id, device, scheduled_date, shift_index, sequence_in_shift)
VALUES (?, ?, ?, ?, ?, ?)
""",
[
(next_version, test_id, device, scheduled_date, shift_index, sequence)
for test_id, device, scheduled_date, shift_index, sequence in entries
],
)
return next_version
def get_schedule_week(start_date: str, db_path: str | Path = DB_PATH) -> list[ScheduleRow]:
with get_connection(db_path) as conn:
version_row = conn.execute("SELECT MAX(schedule_version) AS latest FROM schedules").fetchone()
latest = version_row["latest"]
if latest is None:
return []
rows = conn.execute(
"""
SELECT
s.test_id,
s.device,
s.scheduled_date,
s.shift_index,
s.sequence_in_shift,
t.test_type,
t.rotation,
t.config_json,
t.status,
t.priority,
t.estimated_minutes
FROM schedules s
JOIN tests t ON t.test_id = s.test_id AND t.device = s.device
WHERE s.schedule_version = ?
AND s.scheduled_date >= ?
AND s.scheduled_date < date(?, '+7 day')
ORDER BY s.scheduled_date, s.shift_index, s.sequence_in_shift
""",
(latest, start_date, start_date),
).fetchall()
result: list[ScheduleRow] = []
for row in rows:
result.append(
ScheduleRow(
test_id=row["test_id"],
device=row["device"],
scheduled_date=row["scheduled_date"],
shift_index=int(row["shift_index"]),
sequence_in_shift=int(row["sequence_in_shift"]),
test_type=row["test_type"],
rotation=row["rotation"],
config=json.loads(row["config_json"] or "{}"),
status=row["status"],
priority=int(row["priority"]),
estimated_minutes=int(row["estimated_minutes"]),
)
)
return result
def mark_tests_completed(test_ids_with_device: list[tuple[str, str]], db_path: str | Path = DB_PATH) -> None:
if not test_ids_with_device:
return
with get_connection(db_path) as conn:
conn.executemany(
"""
UPDATE tests
SET status = 'completed', updated_at = CURRENT_TIMESTAMP
WHERE test_id = ?
AND device = ?
AND status != 'completed'
""",
test_ids_with_device,
)
def upsert_holidays(dates: list[str], db_path: str | Path = DB_PATH) -> None:
"""Replace all holidays with the provided list of YYYY-MM-DD date strings."""
with get_connection(db_path) as conn:
conn.execute("DELETE FROM holidays")
if dates:
conn.executemany(
"INSERT OR IGNORE INTO holidays(date) VALUES (?)",
[(d,) for d in dates],
)