next-tjwater-server/scripts/seed_peripheral_sample.py

"""Idempotently seed synthetic peripheral demo resources for one explicit project."""
from __future__ import annotations

import argparse
import json
import sys
from pathlib import Path
from uuid import UUID, uuid5

from psycopg.types.json import Jsonb

ROOT = Path(__file__).resolve().parents[1]
sys.path.insert(0, str(ROOT))

from app.services.content_digest import digest
from samples.peripheral.auth import Identity
from samples.peripheral.config import SampleSettings
from samples.peripheral.demo_data import build_demo_content
from samples.peripheral.schemas import CreateTestDataRecord
from samples.peripheral.seed_identity import SEED_ACTOR_ID, SEED_NAMESPACE
from samples.peripheral.store import PostgresStore


ACTOR_ID = SEED_ACTOR_ID


def request_id(kind: str) -> UUID:
    return uuid5(SEED_NAMESPACE, kind)


TEST_DATA = (
    ("DEMO-P001", "合成压力 1 · 00:00", "2026-09-17T00:00:00Z", 24.37, "m", "valid"),
    ("DEMO-P001", "合成压力 1 · 00:15", "2026-09-17T00:15:00Z", 24.18, "m", "valid"),
    ("DEMO-P001", "合成压力 1 · 00:30 缺测", "2026-09-17T00:30:00Z", None, "m", "missing"),
    ("DEMO-P001", "合成压力 1 · 00:45", "2026-09-17T00:45:00Z", 24.42, "m", "valid"),
    ("DEMO-P002", "合成压力 2 · 00:00", "2026-09-17T00:00:00Z", 31.26, "m", "valid"),
    ("DEMO-P002", "合成压力 2 · 00:15", "2026-09-17T00:15:00Z", 31.31, "m", "valid"),
    ("DEMO-P002", "合成压力 2 · 00:30", "2026-09-17T00:30:00Z", 31.18, "m", "valid"),
    ("DEMO-P002", "合成压力 2 · 00:45", "2026-09-17T00:45:00Z", 31.36, "m", "valid"),
    ("DEMO-F001", "合成流量 00:00", "2026-09-17T00:00:00Z", 186.43, "m3/h", "valid"),
    ("DEMO-F001", "合成流量 00:15", "2026-09-17T00:15:00Z", 185.90, "m3/h", "valid"),
    ("DEMO-F001", "合成零流量 00:30", "2026-09-17T00:30:00Z", 0.0, "m3/h", "valid"),
    ("DEMO-F001", "合成流量 00:45", "2026-09-17T00:45:00Z", 187.12, "m3/h", "valid"),
)


def previous_seed_result(
    store: PostgresStore,
    project_id: str,
    *,
    seed_request_id: UUID,
    command_key: str,
) -> dict | None:
    """Read this script's earlier receipt without weakening API idempotency checks."""
    with store.connection(readonly=True) as connection:
        row = connection.execute(
            """SELECT command_key,result
               FROM peripheral_sample.requests
               WHERE project_id=%s AND actor_id=%s AND request_id=%s""",
            (project_id, str(ACTOR_ID), seed_request_id),
        ).fetchone()
    if row is None:
        return None
    if row["command_key"] != command_key:
        raise RuntimeError("seed request id is already used by another command")
    return row["result"]


def restore_seeded_test_data(
    store: PostgresStore,
    project_id: str,
    record_ids: list[str],
) -> None:
    """Restore deterministic seed rows that were soft-deleted after an earlier run."""
    if not record_ids:
        return

    with store.connection() as connection:
        connection.execute(
            """UPDATE peripheral_sample.test_data_records
                SET document=document - 'deleted_at' - 'deleted_by',
                    deleted_at=NULL, deleted_by=NULL
                WHERE project_id=%s AND record_id=ANY(%s::uuid[])""",
            (project_id, record_ids),
        )


def ensure_demo_dataset(store: PostgresStore, project_id: str) -> str:
    content = build_demo_content()
    content_digest = digest(content)
    with store.connection() as connection:
        row = connection.execute(
            """SELECT digest FROM peripheral_sample.demo_datasets
               WHERE project_id=%s AND dataset_id='water-demo-v1'""",
            (project_id,),
        ).fetchone()
        if row is None:
            connection.execute(
                """INSERT INTO peripheral_sample.demo_datasets
                   (project_id,dataset_id,digest,document)
                   VALUES (%s,'water-demo-v1',%s,%s)""",
                (project_id, content_digest, Jsonb(content)),
            )
        elif row["digest"] != content_digest:
            connection.execute(
                """UPDATE peripheral_sample.demo_datasets
                   SET digest=%s, document=%s
                   WHERE project_id=%s AND dataset_id='water-demo-v1'""",
                (content_digest, Jsonb(content), project_id),
            )
    return content_digest


def main() -> int:
    parser = argparse.ArgumentParser()
    parser.add_argument("--project-id", required=True, type=UUID)
    args = parser.parse_args()
    project_id = str(args.project_id)
    store = PostgresStore(SampleSettings(db_name="tjwater_v2", enable_faults=False))
    who = Identity(project_id, str(ACTOR_ID), True)

    dataset_digest = ensure_demo_dataset(store, project_id)
    record_ids = []
    for point_id, label, observed_at, value, unit, quality in TEST_DATA:
        seed_request_id = request_id(f"test-data/{point_id}/{observed_at}")
        command = CreateTestDataRecord(
            request_id=seed_request_id,
            label=label,
            point_id=point_id,
            observed_at=observed_at,
            value=value,
            unit=unit,
            quality=quality,
            notes="外围样板固定合成数据，不是现场监测记录。",
        )
        previous = previous_seed_result(
            store,
            project_id,
            seed_request_id=seed_request_id,
            command_key="create_test_data",
        )
        result = previous or store.create_test_data(who, command)
        record_ids.append(result["record_id"])
    restore_seeded_test_data(store, project_id, record_ids)

    result = {
        "project_id": project_id,
        "actor_id": str(ACTOR_ID),
        "origin": "synthetic",
        "demo_dataset_digest": dataset_digest,
        "test_data_records": {"count": len(record_ids), "ids": record_ids},
    }
    print(json.dumps(result, ensure_ascii=False, indent=2, sort_keys=True))
    return 0


if __name__ == "__main__":
    raise SystemExit(main())