Files
scheduler/backend/db.py
T
2026-07-23 14:02:14 -04:00

1019 lines
29 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 test_config import serialize_station_testpoint_map, resolve_test_config_keys
APP_ROOT = Path(__file__).resolve().parent
DB_PATH = Path(os.getenv("DB_PATH", str(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
MAX_SCHEDULE_VERSIONS = 50
@dataclass(frozen=True)
class TestRecord:
test_id: str
device: str
test_type: str
rotation: str | None
rx_tx: str | None
power_mode: 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),
power_mode TEXT,
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', 'rerun')),
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 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", "power_mode", "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")
_ensure_column(conn, "schedules", "status_snapshot", "TEXT")
conn.execute(
"""
UPDATE schedules
SET status_snapshot = COALESCE((
SELECT t.status
FROM tests t
WHERE t.test_id = schedules.test_id
AND t.device = schedules.device
), 'pending')
WHERE status_snapshot IS NULL
"""
)
prune_schedule_versions(conn=conn)
# 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."""
return serialize_station_testpoint_map(config)
def upsert_tests(records: list[TestRecord], db_path: str | Path = DB_PATH) -> int:
if not records:
return 0
values: list[tuple[Any, ...]] = []
for r in records:
station_testpoint_map = _serialize_station_testpoint_map(r.config)
values.append(
(
r.test_id,
r.device,
r.test_type,
r.rotation,
r.rx_tx,
r.power_mode,
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 {}),
station_testpoint_map,
)
)
with get_connection(db_path) as conn:
conn.executemany(
"""
INSERT INTO tests(
test_id, device, test_type, rotation, rx_tx, power_mode, 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,
power_mode = excluded.power_mode,
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 = CASE
WHEN tests.status IN ('completed', 'rerun') THEN tests.status
ELSE excluded.status
END,
excluded = excluded.excluded,
raw_payload = excluded.raw_payload,
station_testpoint_map = excluded.station_testpoint_map,
updated_at = CURRENT_TIMESTAMP
""",
values,
)
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:
return _normalize_text(record.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 get_completed_test_tcs(db_path: str | Path = DB_PATH) -> set[str]:
"""Return the set of TC names that have at least one completed test."""
with get_connection(db_path) as conn:
rows = conn.execute(
"SELECT config_json FROM tests WHERE status = 'completed' AND excluded = 0"
).fetchall()
tcs: set[str] = set()
for row in rows:
config = json.loads(row["config_json"] or "{}")
tcs.update(resolve_test_config_keys(config))
return tcs
def get_completed_dut_test_ids(db_path: str | Path = DB_PATH) -> set[str]:
"""Return test IDs that have been completed on DUT."""
with get_connection(db_path) as conn:
rows = conn.execute(
"SELECT test_id FROM tests WHERE status = 'completed' AND excluded = 0 AND device = ?",
(DUT,)
).fetchall()
return {row["test_id"] for row in rows}
# Get tests for device that are not excluded regardless of completion status.
def get_not_excluded_tests(device: str, db_path: str | Path = DB_PATH, rules: dict[str, Any] | None = None) -> list[TestRecord]:
tokens = _parse_rule_tokens(rules.get("rule") if rules else "") if rules else []
with get_connection(db_path) as conn:
rows = conn.execute(
"""
SELECT
test_id, device, test_type, rotation, rx_tx, has_coe_pair,
power_mode, config_json, victim_band,
coe_pairing_json, priority, throttled, estimated_minutes,
excluded, status, raw_payload
FROM tests
WHERE excluded = 0
AND 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"],
power_mode=row["power_mode"],
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)]
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,
power_mode, 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"],
power_mode=row["power_mode"],
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,
power_mode, 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"],
power_mode=row["power_mode"],
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, str]],
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,
status_snapshot
)
VALUES (?, ?, ?, ?, ?, ?, ?)
""",
[
(next_version, test_id, device, scheduled_date, shift_index, sequence, status_snapshot)
for test_id, device, scheduled_date, shift_index, sequence, status_snapshot in entries
],
)
prune_schedule_versions(conn=conn)
return next_version
def prune_schedule_versions(
max_versions: int = MAX_SCHEDULE_VERSIONS,
db_path: str | Path = DB_PATH,
conn: sqlite3.Connection | None = None,
) -> int:
if max_versions <= 0:
raise ValueError("max_versions must be positive")
def _prune(active_conn: sqlite3.Connection) -> int:
cursor = active_conn.execute(
"""
DELETE FROM schedules
WHERE schedule_version IN (
SELECT schedule_version
FROM (
SELECT schedule_version
FROM schedules
GROUP BY schedule_version
ORDER BY schedule_version DESC
LIMIT -1 OFFSET ?
)
)
""",
(max_versions,),
)
return cursor.rowcount
if conn is not None:
return _prune(conn)
with get_connection(db_path) as active_conn:
return _prune(active_conn)
def _resolve_requested_schedule_version(
conn: sqlite3.Connection,
version: int | None,
) -> tuple[int | None, int | None]:
latest_row = conn.execute("SELECT MAX(schedule_version) AS latest FROM schedules").fetchone()
latest = latest_row["latest"]
latest_version = int(latest) if latest is not None else None
if version is not None:
row = conn.execute(
"SELECT 1 AS exists_row FROM schedules WHERE schedule_version = ? LIMIT 1",
(version,),
).fetchone()
return (int(version) if row else None, latest_version)
return latest_version, latest_version
def _hydrate_schedule_rows(rows: list[sqlite3.Row]) -> list[ScheduleRow]:
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 get_schedule_versions(db_path: str | Path = DB_PATH) -> list[dict[str, Any]]:
with get_connection(db_path) as conn:
rows = conn.execute(
"""
SELECT
schedule_version,
MIN(created_at) AS created_at,
COUNT(*) AS entry_count
FROM schedules
GROUP BY schedule_version
ORDER BY schedule_version DESC
"""
).fetchall()
return [
{
"schedule_version": int(row["schedule_version"]),
"created_at": row["created_at"],
"entry_count": int(row["entry_count"]),
}
for row in rows
]
def resolve_schedule_version(
version: int | None = None,
db_path: str | Path = DB_PATH,
) -> int | None:
with get_connection(db_path) as conn:
selected_version, _latest_version = _resolve_requested_schedule_version(conn, version)
return selected_version
def get_schedule_week(
start_date: str,
version: int | None = None,
db_path: str | Path = DB_PATH,
) -> list[ScheduleRow]:
with get_connection(db_path) as conn:
selected_version, latest_version = _resolve_requested_schedule_version(conn, version)
if selected_version is None:
return []
use_live_status = latest_version is not None and selected_version == latest_version
if use_live_status:
# Only surface live 'rerun' for slots that are actually overdue (past).
# Future slots fall back to status_snapshot so they don't show red.
status_sql = """
CASE
WHEN t.status = 'completed' THEN 'completed'
WHEN t.status = 'rerun'
AND (
s.scheduled_date < date('now')
OR (s.scheduled_date = date('now') AND s.shift_index = 1)
)
THEN 'rerun'
WHEN t.status = 'rerun' THEN COALESCE(s.status_snapshot, 'pending')
ELSE COALESCE(t.status, s.status_snapshot, 'pending')
END
"""
else:
status_sql = "COALESCE(s.status_snapshot, 'pending')"
rows = conn.execute(
f"""
SELECT
s.test_id,
s.device,
s.scheduled_date,
s.shift_index,
s.sequence_in_shift,
t.test_type,
t.rotation,
t.config_json,
{status_sql} AS 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')
AND s.shift_index IN (2, 3)
)
OR (
s.scheduled_date > ?
AND s.scheduled_date < date(?, '+7 day')
AND s.shift_index = 1
)
OR (
s.scheduled_date = date(?, '+7 day')
AND s.shift_index = 1
)
)
ORDER BY s.scheduled_date, s.shift_index, s.sequence_in_shift
""",
(
selected_version,
start_date,
start_date,
start_date,
start_date,
start_date,
),
).fetchall()
return _hydrate_schedule_rows(rows)
def get_schedule_rows(
version: int | None = None,
db_path: str | Path = DB_PATH,
) -> list[ScheduleRow]:
with get_connection(db_path) as conn:
selected_version, latest_version = _resolve_requested_schedule_version(conn, version)
if selected_version is None:
return []
use_live_status = latest_version is not None and selected_version == latest_version
if use_live_status:
status_sql = """
CASE
WHEN t.status = 'completed' THEN 'completed'
WHEN t.status = 'rerun'
AND (
s.scheduled_date < date('now')
OR (s.scheduled_date = date('now') AND s.shift_index = 1)
)
THEN 'rerun'
WHEN t.status = 'rerun' THEN COALESCE(s.status_snapshot, 'pending')
ELSE COALESCE(t.status, s.status_snapshot, 'pending')
END
"""
else:
status_sql = "COALESCE(s.status_snapshot, 'pending')"
rows = conn.execute(
f"""
SELECT
s.test_id,
s.device,
s.scheduled_date,
s.shift_index,
s.sequence_in_shift,
t.test_type,
t.rotation,
t.config_json,
{status_sql} AS 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 = ?
ORDER BY s.scheduled_date, s.shift_index, s.sequence_in_shift
""",
(selected_version,),
).fetchall()
return _hydrate_schedule_rows(rows)
def reset_completed_to_pending(db_path: str | Path = DB_PATH) -> int:
"""Reset all 'completed' tests back to 'pending'. Used before a full directory rescan."""
with get_connection(db_path) as conn:
cursor = conn.execute(
"""
UPDATE tests
SET status = 'pending', updated_at = CURRENT_TIMESTAMP
WHERE status = 'completed'
"""
)
return cursor.rowcount
def mark_tests_completed(test_ids_with_device: list[tuple[str, str]], db_path: str | Path = DB_PATH) -> int:
if not test_ids_with_device:
return 0
with get_connection(db_path) as conn:
# executemany() doesn't properly return rowcount, so we track manually
total_updated = 0
for test_id, device in test_ids_with_device:
cursor = conn.execute(
"""
UPDATE tests
SET status = 'completed', updated_at = CURRENT_TIMESTAMP
WHERE test_id = ?
AND device = ?
AND status != 'completed'
""",
(test_id, device),
)
if cursor.rowcount > 0:
total_updated += cursor.rowcount
print(f"[db] marked {test_id} on {device} as completed")
else:
print(f"[db] {test_id} on {device}: no update (not found or already completed)")
return total_updated
def mark_overdue_as_rerun(db_path: str | Path = DB_PATH) -> int:
from datetime import date as _date, datetime as _datetime
now = _datetime.now()
today = now.date().isoformat()
shift1_ended = now.hour >= 10 # shift 1 ends ~10am
with get_connection(db_path) as conn:
# Use the most recent schedule version that has any overdue (past-due) entries.
# If a fresh schedule was just compiled, its entries are all in the future and
# MAX(schedule_version) would miss the overdue context from the prior version.
if shift1_ended:
overdue_date_clause = "s.scheduled_date < :today OR (s.scheduled_date = :today AND s.shift_index = 1)"
else:
overdue_date_clause = "s.scheduled_date < :today"
version_row = conn.execute(
f"""
SELECT MAX(s.schedule_version) AS v
FROM schedules s
JOIN tests t ON t.test_id = s.test_id AND t.device = s.device
WHERE t.status = 'pending'
AND ({overdue_date_clause})
""",
{"today": today},
).fetchone()
latest = version_row["v"] if version_row else None
if latest is None:
# No version has overdue pending tests — fall back to MAX to still run the
# update (it will simply mark zero rows).
fb = conn.execute("SELECT MAX(schedule_version) AS v FROM schedules").fetchone()
latest = fb["v"] if fb else None
if latest is None:
return 0
if shift1_ended:
rows = conn.execute(
"""
SELECT DISTINCT s.test_id, s.device
FROM schedules s
JOIN tests t ON t.test_id = s.test_id AND t.device = s.device
WHERE s.schedule_version = ?
AND t.status = 'pending'
AND (
s.scheduled_date < ?
OR (s.scheduled_date = ? AND s.shift_index = 1)
)
""",
(latest, today, today),
).fetchall()
else:
rows = conn.execute(
"""
SELECT DISTINCT s.test_id, s.device
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 t.status = 'pending'
""",
(latest, today),
).fetchall()
if not rows:
return 0
conn.executemany(
"""
UPDATE tests
SET status = 'rerun', updated_at = CURRENT_TIMESTAMP
WHERE test_id = ? AND device = ? AND status = 'pending'
""",
[(r["test_id"], r["device"]) for r in rows],
)
return len(rows)
def get_rerun_tests(db_path: str | Path = DB_PATH) -> list[dict]:
"""Return all overdue tests that have not been completed.
Checks the most recent schedule version that contains overdue pending entries
so that recompiling a fresh schedule does not erase the overdue context.
Overdue = scheduled before today (any shift) or today shift 1 if it has ended.
"""
from datetime import date as _date, datetime as _datetime
now = _datetime.now()
today = now.date().isoformat()
shift1_ended = now.hour >= 10
with get_connection(db_path) as conn:
if shift1_ended:
overdue_date_clause = "s.scheduled_date < :today OR (s.scheduled_date = :today AND s.shift_index = 1)"
else:
overdue_date_clause = "s.scheduled_date < :today"
version_row = conn.execute(
f"""
SELECT MAX(s.schedule_version) AS v
FROM schedules s
JOIN tests t ON t.test_id = s.test_id AND t.device = s.device
WHERE t.status != 'completed'
AND ({overdue_date_clause})
""",
{"today": today},
).fetchone()
latest = version_row["v"] if version_row else None
if latest is None:
return []
rows = conn.execute(
f"""
SELECT DISTINCT t.test_id, t.device, t.test_type, t.estimated_minutes
FROM tests t
JOIN schedules s ON s.test_id = t.test_id AND s.device = t.device
WHERE s.schedule_version = :v
AND t.status != 'completed'
AND ({overdue_date_clause})
ORDER BY t.test_id, t.device
""",
{"v": latest, "today": today},
).fetchall()
return [
{
"test_id": r["test_id"],
"device": r["device"],
"test_type": r["test_type"],
"estimated_minutes": r["estimated_minutes"],
}
for r in rows
]
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],
)
def reset_tests(db_path: str | Path = DB_PATH) -> int:
"""Clear all tests records keep empty tests table."""
with get_connection(db_path) as conn:
cursor = conn.execute("DELETE FROM tests")
return cursor.rowcount
def reset_all_data(db_path: str | Path = DB_PATH) -> dict[str, int]:
"""Clear all scheduler data from mutable tables.
Keeps schema and runtime default rows intact.
"""
with get_connection(db_path) as conn:
schedules_deleted = conn.execute("DELETE FROM schedules").rowcount
tests_deleted = conn.execute("DELETE FROM tests").rowcount
settings_deleted = conn.execute("DELETE FROM settings").rowcount
holidays_deleted = conn.execute("DELETE FROM holidays").rowcount
# Reset autoincrement counters for cleared tables.
conn.execute(
"DELETE FROM sqlite_sequence WHERE name IN ('tests', 'schedules')"
)
return {
"tests_deleted": tests_deleted,
"schedules_deleted": schedules_deleted,
"settings_deleted": settings_deleted,
"holidays_deleted": holidays_deleted,
}