105 lines
4.2 KiB
Python
105 lines
4.2 KiB
Python
#!/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())
|