added export tests feature
This commit is contained in:
+61
-2
@@ -2,15 +2,17 @@ from contextlib import asynccontextmanager
|
||||
from pathlib import Path
|
||||
from typing import Any
|
||||
from datetime import date, datetime, timedelta
|
||||
import io
|
||||
import os
|
||||
import zipfile
|
||||
|
||||
from fastapi import FastAPI, HTTPException
|
||||
from fastapi import FastAPI, HTTPException, Response
|
||||
from fastapi.middleware.cors import CORSMiddleware
|
||||
from pydantic import BaseModel, Field
|
||||
|
||||
import db as db
|
||||
from parser import CsvValidationError
|
||||
from file_manager import resolve_requested_csv_path
|
||||
from file_manager import resolve_requested_csv_path, read_file_bytes
|
||||
|
||||
# Ensure new_scheduler resolves the same DUT/REF labels as db records.
|
||||
os.environ.setdefault("DUT", db.DUT)
|
||||
@@ -526,6 +528,63 @@ def get_schedule_week(start: str | None = None, version: int | None = None) -> d
|
||||
"windows": _build_schedule_windows(week_start_date, all_rows, holiday_dates),
|
||||
}
|
||||
|
||||
@app.get("/api/schedule/export")
|
||||
def export_window(window_id: str, version: int | None = None) -> Response:
|
||||
resolved_version = db.resolve_schedule_version(version, DB_PATH)
|
||||
if version is not None and resolved_version is None:
|
||||
raise HTTPException(status_code=404, detail=f"Schedule version {version} was not found")
|
||||
|
||||
all_rows = db.get_schedule_rows(resolved_version, DB_PATH)
|
||||
if not all_rows:
|
||||
raise HTTPException(status_code=404, detail="No schedule found")
|
||||
|
||||
holiday_dates = db.list_holidays(DB_PATH)
|
||||
week_start = min(datetime.strptime(r.scheduled_date, "%Y-%m-%d").date() for r in all_rows)
|
||||
windows = _build_schedule_windows(week_start, all_rows, holiday_dates)
|
||||
|
||||
target_window = next((w for w in windows if w["window_id"] == window_id), None)
|
||||
if target_window is None:
|
||||
raise HTTPException(status_code=404, detail=f"Window '{window_id}' not found")
|
||||
|
||||
wanted = {(t["test_id"], t["device"]) for t in target_window["tests"]}
|
||||
if not wanted:
|
||||
raise HTTPException(status_code=400, detail="This window has no scheduled tests")
|
||||
|
||||
# Look up stored file paths from DB — avoids re-scanning and handles R<n> prefixes
|
||||
with db.get_connection(DB_PATH) as conn:
|
||||
rows = conn.execute(
|
||||
"SELECT test_id, device, file_path FROM tests WHERE file_path IS NOT NULL"
|
||||
).fetchall()
|
||||
path_lookup: dict[tuple[str, str], str] = {
|
||||
(row["test_id"], row["device"]): row["file_path"] for row in rows
|
||||
}
|
||||
file_paths = [path_lookup[k] for k in sorted(wanted) if k in path_lookup]
|
||||
|
||||
settings = db.read_settings(DB_PATH)
|
||||
smb_credentials = _smb_credentials_from_settings(settings)
|
||||
|
||||
buf = io.BytesIO()
|
||||
with zipfile.ZipFile(buf, mode="w", compression=zipfile.ZIP_DEFLATED) as zf:
|
||||
for fp in file_paths:
|
||||
try:
|
||||
data = read_file_bytes(fp, smb_credentials)
|
||||
except OSError as exc:
|
||||
print(f"[export] Could not read {fp}: {exc}")
|
||||
continue
|
||||
zf.writestr(os.path.basename(fp), data)
|
||||
|
||||
start_date_str = datetime.strptime(target_window["start_date"], "%Y-%m-%d").strftime("%m-%d-%Y")
|
||||
devices = "_".join(sorted({t["device"] for t in target_window["tests"]}))
|
||||
suffix = "daytime_tests" if target_window["window_type"] == "daytime" else "tests"
|
||||
zip_name = f"{start_date_str}_{devices}_{suffix}.zip"
|
||||
|
||||
return Response(
|
||||
content=buf.getvalue(),
|
||||
media_type="application/zip",
|
||||
headers={"Content-Disposition": f'attachment; filename="{zip_name}"'},
|
||||
)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
import uvicorn
|
||||
uvicorn.run("app:app", host="0.0.0.0", port=8000, reload=True)
|
||||
|
||||
+6
-2
@@ -37,6 +37,7 @@ class TestRecord:
|
||||
excluded: bool = False
|
||||
raw_payload: dict[str, Any] | None = None
|
||||
station_testpoint_map: str | None = None
|
||||
file_path: str | None = None
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
@@ -134,6 +135,7 @@ def init_db(db_path: str | Path = DB_PATH) -> None:
|
||||
_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, "tests", "file_path", "TEXT")
|
||||
_ensure_column(conn, "schedules", "status_snapshot", "TEXT")
|
||||
conn.execute(
|
||||
"""
|
||||
@@ -186,6 +188,7 @@ def upsert_tests(records: list[TestRecord], db_path: str | Path = DB_PATH) -> in
|
||||
int(r.excluded),
|
||||
json.dumps(r.raw_payload or {}),
|
||||
station_testpoint_map,
|
||||
r.file_path,
|
||||
)
|
||||
)
|
||||
|
||||
@@ -194,9 +197,9 @@ def upsert_tests(records: list[TestRecord], db_path: str | Path = DB_PATH) -> in
|
||||
"""
|
||||
INSERT INTO tests(
|
||||
test_id, device, test_type, rotation, rx_tx, power_mode, has_coe_pair,
|
||||
coe_pairing_json, config_json, victim_band, throttled, estimated_minutes, status, excluded, raw_payload, station_testpoint_map
|
||||
coe_pairing_json, config_json, victim_band, throttled, estimated_minutes, status, excluded, raw_payload, station_testpoint_map, file_path
|
||||
)
|
||||
VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)
|
||||
VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)
|
||||
ON CONFLICT(test_id, device) DO UPDATE SET
|
||||
test_type = excluded.test_type,
|
||||
device = excluded.device,
|
||||
@@ -216,6 +219,7 @@ def upsert_tests(records: list[TestRecord], db_path: str | Path = DB_PATH) -> in
|
||||
excluded = excluded.excluded,
|
||||
raw_payload = excluded.raw_payload,
|
||||
station_testpoint_map = excluded.station_testpoint_map,
|
||||
file_path = excluded.file_path,
|
||||
updated_at = CURRENT_TIMESTAMP
|
||||
""",
|
||||
values,
|
||||
|
||||
@@ -205,6 +205,17 @@ def open_csv_handle(path: str, smb_credentials: dict[str, Any] | None):
|
||||
return smbclient.open_file(path, mode="r", encoding="utf-8-sig", newline="")
|
||||
return Path(path).open("r", encoding="utf-8-sig", newline="")
|
||||
|
||||
|
||||
def read_file_bytes(path: str, smb_credentials: dict[str, Any] | None = None) -> bytes:
|
||||
if is_unc_path(path):
|
||||
if smbclient is None:
|
||||
raise ModuleNotFoundError("smbclient is required to read UNC paths")
|
||||
_register_smb_session_if_needed(path, smb_credentials=smb_credentials)
|
||||
with smbclient.open_file(path, mode="rb") as fh:
|
||||
return fh.read()
|
||||
with open(path, "rb") as fh:
|
||||
return fh.read()
|
||||
|
||||
def join_path(path, name):
|
||||
if is_unc_path(path):
|
||||
base = path.rstrip("\\")
|
||||
|
||||
@@ -286,6 +286,7 @@ def process_targets(target_dir, csv_paths, smb_credentials=None, runtime_overrid
|
||||
status=parsed["status"],
|
||||
excluded=False,
|
||||
raw_payload=csv_entry.get("raw_payload", None),
|
||||
file_path=join_path(parent_path, filename),
|
||||
))
|
||||
|
||||
|
||||
|
||||
Reference in New Issue
Block a user