fix

scripts/prepare_logs_e2e_data.py

@@ -0,0 +1,104 @@
+#!/usr/bin/env python3
+import argparse
+import json
+from pathlib import Path
+from typing import Any, Dict
+
+import psycopg2
+import yaml
+
+
+def load_yaml(path: Path) -> Dict[str, Any]:
+    return yaml.safe_load(path.read_text(encoding="utf-8"))
+
+
+def parse_pg_dsn(dsn: str) -> str:
+    parts = dsn.split()
+    kept = []
+    timezone = None
+    for part in parts:
+        if "=" not in part:
+            kept.append(part)
+            continue
+        k, v = part.split("=", 1)
+        if k.lower() == "timezone":
+            timezone = v
+            continue
+        kept.append(part)
+    if timezone:
+        kept.append(f"options='-c timezone={timezone}'")
+    return " ".join(kept)
+
+
+def main() -> int:
+    parser = argparse.ArgumentParser(description="日志管理 E2E 测试数据准备脚本")
+    parser.add_argument(
+        "--config",
+        default="d:/work/ops/logs/etc/logs_dev.yaml",
+        help="logs 配置文件路径",
+    )
+    parser.add_argument("--run-id", required=True, help="本次测试 run id")
+    parser.add_argument("--cleanup-only", action="store_true", help="仅清理历史测试数据")
+    args = parser.parse_args()
+
+    cfg = load_yaml(Path(args.config))
+    dsn = parse_pg_dsn(cfg["Databases"]["Source"][0])
+    run_id = args.run_id
+    marker = f"%[E2E:{run_id}]%"
+
+    summary: Dict[str, Any] = {"run_id": run_id, "cleanup": {}, "seed": {}}
+    with psycopg2.connect(dsn) as conn:
+        with conn.cursor() as cur:
+            cur.execute("DELETE FROM logs_alert_outbox WHERE id IN (SELECT id FROM logs_alert_outbox ORDER BY id DESC LIMIT 0)")
+            cur.execute("DELETE FROM logs_syslog_rules WHERE name LIKE %s", (marker,))
+            summary["cleanup"]["logs_syslog_rules"] = cur.rowcount
+            cur.execute("DELETE FROM logs_trap_rules WHERE name LIKE %s", (marker,))
+            summary["cleanup"]["logs_trap_rules"] = cur.rowcount
+            cur.execute("DELETE FROM logs_trap_dictionary WHERE title LIKE %s", (marker,))
+            summary["cleanup"]["logs_trap_dictionary"] = cur.rowcount
+            cur.execute("DELETE FROM logs_trap_shields WHERE name LIKE %s", (marker,))
+            summary["cleanup"]["logs_trap_shields"] = cur.rowcount
+            cur.execute(
+                "DELETE FROM logs_resource_event_dedup WHERE event_id LIKE %s",
+                (f"e2e-{run_id}-%",),
+            )
+            summary["cleanup"]["logs_resource_event_dedup"] = cur.rowcount
+            if not args.cleanup_only:
+                cur.execute(
+                    """
+                    INSERT INTO logs_resource_mappings(
+                        resource_type, resource_id, resource_name, ips_json, hostnames_json, labels_json,
+                        version, is_deleted, last_event_id, event_time, created_at, updated_at
+                    ) VALUES (%s,%s,%s,%s,%s,%s,%s,false,%s,now(),now(),now())
+                    ON CONFLICT (resource_type, resource_id)
+                    DO UPDATE SET
+                        resource_name=EXCLUDED.resource_name,
+                        ips_json=EXCLUDED.ips_json,
+                        hostnames_json=EXCLUDED.hostnames_json,
+                        labels_json=EXCLUDED.labels_json,
+                        version=EXCLUDED.version,
+                        is_deleted=false,
+                        last_event_id=EXCLUDED.last_event_id,
+                        event_time=now(),
+                        updated_at=now()
+                    """,
+                    (
+                        "server",
+                        f"seed-{run_id}",
+                        f"E2E-Seed-{run_id}",
+                        json.dumps(["127.0.0.1"]),
+                        json.dumps([f"e2e-host-{run_id}"]),
+                        json.dumps({"source": "prepare_script"}),
+                        1,
+                        f"e2e-{run_id}-seed",
+                    ),
+                )
+                summary["seed"]["resource_mapping"] = {"resource_type": "server", "resource_id": f"seed-{run_id}"}
+        conn.commit()
+
+    print(json.dumps(summary, ensure_ascii=False, indent=2))
+    return 0
+
+
+if __name__ == "__main__":
+    raise SystemExit(main())