Files
scheduler/backend/db.py
T
2026-06-16 15:07:59 -04:00

512 lines
14 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
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]]
estimated_minutes: int
status: str = "pending"
excluded: bool = False
raw_payload: dict[str, Any] | 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
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,
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")
# 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 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, estimated_minutes, status, excluded, raw_payload
)
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,
estimated_minutes = excluded.estimated_minutes,
status = excluded.status,
excluded = excluded.excluded,
raw_payload = excluded.raw_payload,
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,
r.estimated_minutes,
r.status,
int(r.excluded),
json.dumps(r.raw_payload or {}),
)
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 _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)
for band in ("5G", "6G", "2G"):
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 = _entry_value(_band_entry(record.config, "6G"), "power_mode")
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, 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 "{}"),
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, estimated_minutes,
excluded, status, raw_payload
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 "{}"),
estimated_minutes=int(row["estimated_minutes"]),
status=row["status"],
excluded=bool(row["excluded"]),
raw_payload=json.loads(row["raw_payload"] or "{}"),
)
)
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.status,
t.priority,
t.estimated_minutes
FROM schedules s
JOIN tests t ON t.test_id = s.test_id
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"],
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,
)