init
This commit is contained in:
Executable
+476
@@ -0,0 +1,476 @@
|
||||
#!/usr/bin/env python3
|
||||
"""TSV registry helpers for Loop PG migration.
|
||||
|
||||
Columns:
|
||||
loop_id, domain, kube_context, yc_cluster, target_cluster,
|
||||
size_bytes, size_pretty, batch_id, status, notes
|
||||
"""
|
||||
from __future__ import annotations
|
||||
|
||||
import argparse
|
||||
import csv
|
||||
import fcntl
|
||||
import os
|
||||
import re
|
||||
import sys
|
||||
import tempfile
|
||||
import time
|
||||
from contextlib import contextmanager
|
||||
from pathlib import Path
|
||||
|
||||
COLUMNS = [
|
||||
"loop_id",
|
||||
"domain",
|
||||
"kube_context",
|
||||
"yc_cluster",
|
||||
"target_cluster",
|
||||
"size_bytes",
|
||||
"size_pretty",
|
||||
"batch_id",
|
||||
"status",
|
||||
"notes",
|
||||
]
|
||||
|
||||
# Fields registry-refresh may overwrite from live scan
|
||||
REFRESH_FIELDS = ("domain", "kube_context", "yc_cluster")
|
||||
|
||||
# Fields never clobbered by refresh if already set (unless --force-*)
|
||||
PRESERVE_IF_SET = ("target_cluster", "batch_id", "status", "size_bytes", "size_pretty", "notes")
|
||||
|
||||
STATUSES = {
|
||||
"pending",
|
||||
"sized",
|
||||
"migrating",
|
||||
"migrated",
|
||||
"verified",
|
||||
"cutover",
|
||||
"done",
|
||||
"failed",
|
||||
"size_error",
|
||||
"excluded",
|
||||
}
|
||||
|
||||
DONE_LIKE = {"migrated", "verified", "cutover", "done"}
|
||||
|
||||
|
||||
def registry_path(migrate_dir: Path) -> Path:
|
||||
return migrate_dir / "registry.tsv"
|
||||
|
||||
|
||||
def empty_row(loop_id: str = "") -> dict[str, str]:
|
||||
return {c: "" for c in COLUMNS} | {"loop_id": loop_id, "status": "pending"}
|
||||
|
||||
|
||||
@contextmanager
|
||||
def registry_lock(path: Path, timeout_sec: float = 120.0):
|
||||
"""Exclusive flock for registry.tsv (and its done/failed mirrors)."""
|
||||
path.parent.mkdir(parents=True, exist_ok=True)
|
||||
lock_path = path.parent / f"{path.name}.lock"
|
||||
fd = os.open(lock_path, os.O_CREAT | os.O_RDWR, 0o644)
|
||||
deadline = time.monotonic() + timeout_sec
|
||||
try:
|
||||
while True:
|
||||
try:
|
||||
fcntl.flock(fd, fcntl.LOCK_EX | fcntl.LOCK_NB)
|
||||
break
|
||||
except BlockingIOError:
|
||||
if time.monotonic() >= deadline:
|
||||
fcntl.flock(fd, fcntl.LOCK_EX) # block past timeout rather than fail cutover
|
||||
break
|
||||
time.sleep(0.05)
|
||||
yield
|
||||
finally:
|
||||
try:
|
||||
fcntl.flock(fd, fcntl.LOCK_UN)
|
||||
finally:
|
||||
os.close(fd)
|
||||
|
||||
|
||||
def _count_data_lines(path: Path) -> int:
|
||||
if not path.exists():
|
||||
return 0
|
||||
n = 0
|
||||
with path.open() as f:
|
||||
for line in f:
|
||||
if not line.strip():
|
||||
continue
|
||||
if line.startswith("loop_id\t") or line.strip() == "loop_id":
|
||||
continue
|
||||
n += 1
|
||||
return n
|
||||
|
||||
|
||||
def load(path: Path) -> dict[str, dict[str, str]]:
|
||||
"""Load registry.tsv. Tolerates header missing or moved (e.g. after `sort`)."""
|
||||
if not path.exists():
|
||||
return {}
|
||||
size = path.stat().st_size
|
||||
rows: dict[str, dict[str, str]] = {}
|
||||
with path.open(newline="") as f:
|
||||
sample = f.readline()
|
||||
if not sample:
|
||||
return {}
|
||||
# Normal file: first line is the header
|
||||
if sample.startswith("loop_id\t") or sample.strip() == "loop_id":
|
||||
f.seek(0)
|
||||
reader = csv.DictReader(f, delimiter="\t")
|
||||
else:
|
||||
# Header missing or sorted to the bottom — parse every line as data
|
||||
f.seek(0)
|
||||
reader = csv.DictReader(f, fieldnames=COLUMNS, delimiter="\t")
|
||||
for raw in reader:
|
||||
if not raw:
|
||||
continue
|
||||
loop_id = (raw.get("loop_id") or "").strip()
|
||||
if not loop_id or loop_id.startswith("#"):
|
||||
continue
|
||||
# stray header row in the body
|
||||
if loop_id == "loop_id":
|
||||
continue
|
||||
row = empty_row(loop_id)
|
||||
for c in COLUMNS:
|
||||
if c in raw and raw[c] is not None:
|
||||
row[c] = str(raw[c]).strip()
|
||||
rows[loop_id] = row
|
||||
# Non-empty file that parses to nothing = corruption / wrong format
|
||||
if size >= 500 and not rows:
|
||||
raise RuntimeError(f"registry {path} is {size} bytes but parsed 0 rows — refuse load")
|
||||
return rows
|
||||
|
||||
|
||||
def _save_unlocked(path: Path, rows: dict[str, dict[str, str]]) -> None:
|
||||
"""Atomic write via unique tempfile (safe under parallel cutover-list)."""
|
||||
path.parent.mkdir(parents=True, exist_ok=True)
|
||||
old_n = _count_data_lines(path)
|
||||
new_n = len(rows)
|
||||
# Hard guard: never collapse a full registry to a handful of rows (parallel race / empty load).
|
||||
# Override: REGISTRY_ALLOW_SHRINK=1
|
||||
allow_shrink = os.environ.get("REGISTRY_ALLOW_SHRINK", "").strip() in ("1", "true", "yes")
|
||||
if not allow_shrink and old_n >= 20 and new_n < max(10, old_n // 2):
|
||||
raise RuntimeError(
|
||||
f"refusing to shrink registry {path} from {old_n} to {new_n} rows "
|
||||
f"(set REGISTRY_ALLOW_SHRINK=1 to override)"
|
||||
)
|
||||
ordered = sorted(
|
||||
rows.values(),
|
||||
key=lambda r: (r.get("batch_id") or "zzz", r.get("kube_context") or "", r["loop_id"]),
|
||||
)
|
||||
fd, tmp_name = tempfile.mkstemp(
|
||||
prefix=f".{path.name}.",
|
||||
suffix=".tmp",
|
||||
dir=str(path.parent),
|
||||
)
|
||||
try:
|
||||
with os.fdopen(fd, "w", newline="") as f:
|
||||
w = csv.DictWriter(f, fieldnames=COLUMNS, delimiter="\t", lineterminator="\n")
|
||||
w.writeheader()
|
||||
for r in ordered:
|
||||
w.writerow({c: r.get(c, "") for c in COLUMNS})
|
||||
os.replace(tmp_name, path)
|
||||
except Exception:
|
||||
try:
|
||||
os.unlink(tmp_name)
|
||||
except OSError:
|
||||
pass
|
||||
raise
|
||||
|
||||
|
||||
def save(path: Path, rows: dict[str, dict[str, str]]) -> None:
|
||||
with registry_lock(path):
|
||||
_save_unlocked(path, rows)
|
||||
|
||||
|
||||
def upsert(rows: dict[str, dict[str, str]], patch: dict[str, str], *, force: bool = False) -> dict[str, str]:
|
||||
loop_id = patch["loop_id"]
|
||||
cur = rows.get(loop_id) or empty_row(loop_id)
|
||||
for k, v in patch.items():
|
||||
if k not in COLUMNS or v is None:
|
||||
continue
|
||||
v = str(v).strip()
|
||||
if k == "loop_id":
|
||||
cur[k] = v
|
||||
continue
|
||||
if not force and k in PRESERVE_IF_SET and cur.get(k):
|
||||
# allow status overwrite only when explicitly forced or empty
|
||||
if k == "status" and not force:
|
||||
continue
|
||||
if k != "status":
|
||||
continue
|
||||
if v != "" or force:
|
||||
cur[k] = v
|
||||
rows[loop_id] = cur
|
||||
return cur
|
||||
|
||||
|
||||
def _write_id_list(path: Path, ids: list[str]) -> None:
|
||||
text = "\n".join(ids) + ("\n" if ids else "")
|
||||
fd, tmp_name = tempfile.mkstemp(prefix=f".{path.name}.", suffix=".tmp", dir=str(path.parent))
|
||||
try:
|
||||
with os.fdopen(fd, "w") as f:
|
||||
f.write(text)
|
||||
os.replace(tmp_name, path)
|
||||
except Exception:
|
||||
try:
|
||||
os.unlink(tmp_name)
|
||||
except OSError:
|
||||
pass
|
||||
raise
|
||||
|
||||
|
||||
def _sync_legacy_mirrors(registry: Path, loop_id: str, status: str) -> None:
|
||||
"""done.txt / failed.txt are write-only mirrors of registry status. Caller holds registry_lock."""
|
||||
root = registry.parent
|
||||
done_p, fail_p = root / "done.txt", root / "failed.txt"
|
||||
|
||||
def _read_ids(path: Path) -> list[str]:
|
||||
if not path.exists():
|
||||
return []
|
||||
return [x for x in path.read_text().split() if x]
|
||||
|
||||
def _add(path: Path, lid: str) -> None:
|
||||
cur = _read_ids(path)
|
||||
if lid not in cur:
|
||||
cur.append(lid)
|
||||
_write_id_list(path, sorted(cur))
|
||||
|
||||
def _del(path: Path, lid: str) -> None:
|
||||
if not path.exists():
|
||||
return
|
||||
cur = [x for x in _read_ids(path) if x != lid]
|
||||
_write_id_list(path, cur)
|
||||
|
||||
if status in DONE_LIKE:
|
||||
_add(done_p, loop_id)
|
||||
_del(fail_p, loop_id)
|
||||
elif status == "failed":
|
||||
_add(fail_p, loop_id)
|
||||
_del(done_p, loop_id)
|
||||
elif status:
|
||||
_del(done_p, loop_id)
|
||||
_del(fail_p, loop_id)
|
||||
|
||||
|
||||
def set_fields(path: Path, loop_id: str, fields: dict[str, str], *, force: bool = True) -> None:
|
||||
# load→upsert→save under one lock so parallel cutover-list does not clobber rows
|
||||
with registry_lock(path):
|
||||
size = path.stat().st_size if path.exists() else 0
|
||||
rows = load(path)
|
||||
# Empty parse of a previously-large file would wipe the registry on save — abort
|
||||
if size >= 500 and not rows:
|
||||
raise RuntimeError(f"registry {path} unreadable ({size} bytes, 0 rows)")
|
||||
if size >= 5000 and len(rows) < 10:
|
||||
raise RuntimeError(
|
||||
f"registry {path} looks truncated ({size} bytes, {len(rows)} rows) — refuse update"
|
||||
)
|
||||
patch = {"loop_id": loop_id, **fields}
|
||||
# Never let force=True clear unrelated columns when caller only patches status/etc.
|
||||
# Only apply non-empty values unless key is status (explicit empty not used).
|
||||
clean = {"loop_id": loop_id}
|
||||
for k, v in fields.items():
|
||||
if k == "loop_id":
|
||||
continue
|
||||
if v is None:
|
||||
continue
|
||||
v = str(v)
|
||||
if v == "" and k != "notes":
|
||||
continue
|
||||
clean[k] = v
|
||||
upsert(rows, clean, force=force)
|
||||
_save_unlocked(path, rows)
|
||||
if "status" in fields:
|
||||
_sync_legacy_mirrors(path, loop_id, fields["status"])
|
||||
|
||||
|
||||
def is_done_status(status: str) -> bool:
|
||||
return status in DONE_LIKE
|
||||
|
||||
|
||||
# Known YC MDB cluster ids (fallback when `yc` CLI unavailable)
|
||||
KNOWN_CID_NAMES = {
|
||||
"c9q9hoccq1gs9jf2g3r9": "LP-PG-1",
|
||||
"c9q87tpl86oinfglqem4": "LP-PG-2",
|
||||
"c9q6tst6nutgn8ju2p5r": "LP-PG-3",
|
||||
"c9q9gqbgbnljg3apqkak": "LP-PG-4",
|
||||
"c9q0a21dumlqo9u3vmm0": "LP-PG-5",
|
||||
"c9qqfe1b3e43gkp4jju5": "LP-PG-6",
|
||||
"c9q4h551pgq3po55f7hu": "LP-PG-FF",
|
||||
"c9qfg20anmmknvmpqiuj": "LP-PG-LM",
|
||||
"c9qfkheuh46d1g731bqi": "postgresql120",
|
||||
"c9qnuaa4qqfrcj11h4rm": "4",
|
||||
"c9qbkifpjolmls057bij": "5",
|
||||
"c9qe5f6bo1c7isk8id9g": "6",
|
||||
}
|
||||
|
||||
# Original shared pool → restore cluster (host rewrite after source cluster delete/move)
|
||||
# old_name → (new_name, old_cid, new_cid)
|
||||
RESTORE_REWRITES = {
|
||||
"LP-PG-1": ("postgresql120", "c9q9hoccq1gs9jf2g3r9", "c9qfkheuh46d1g731bqi"),
|
||||
"LP-PG-4": ("4", "c9q9gqbgbnljg3apqkak", "c9qnuaa4qqfrcj11h4rm"),
|
||||
"LP-PG-5": ("5", "c9q0a21dumlqo9u3vmm0", "c9qbkifpjolmls057bij"),
|
||||
"LP-PG-6": ("6", "c9qqfe1b3e43gkp4jju5", "c9qe5f6bo1c7isk8id9g"),
|
||||
}
|
||||
|
||||
|
||||
def rw_host(cid: str) -> str:
|
||||
return f"c-{cid}.rw.mdb.yandexcloud.net"
|
||||
|
||||
|
||||
def restore_host_map() -> dict[str, tuple[str, str]]:
|
||||
"""old_fqdn → (new_fqdn, new_cluster_name)."""
|
||||
out: dict[str, tuple[str, str]] = {}
|
||||
for old_name, (new_name, old_cid, new_cid) in RESTORE_REWRITES.items():
|
||||
out[rw_host(old_cid)] = (rw_host(new_cid), new_name)
|
||||
return out
|
||||
|
||||
|
||||
def host_to_yc_cluster(host: str, cid_to_name: dict[str, str] | None = None) -> str:
|
||||
host = host or ""
|
||||
if "95.143.178.247" in host:
|
||||
return "CNPG-wilix"
|
||||
if "135.106.201.40" in host:
|
||||
return "CNPG-lm"
|
||||
if "135.106.201.38" in host:
|
||||
return "CNPG-ff"
|
||||
names = {**KNOWN_CID_NAMES, **(cid_to_name or {})}
|
||||
m = re.search(r"c-([a-z0-9]+)\.rw\.mdb", host)
|
||||
if m:
|
||||
cid = m.group(1)
|
||||
return names.get(cid, cid)
|
||||
if "mdb.yandexcloud.net" in host:
|
||||
return f"other-yc:{host.split(':')[0]}"
|
||||
if host:
|
||||
return host.split(":")[0]
|
||||
return ""
|
||||
|
||||
|
||||
def cmd_get(args: argparse.Namespace) -> int:
|
||||
rows = load(Path(args.registry))
|
||||
row = rows.get(args.loop_id)
|
||||
if not row:
|
||||
return 1
|
||||
if args.field:
|
||||
print(row.get(args.field, ""))
|
||||
else:
|
||||
print("\t".join(row.get(c, "") for c in COLUMNS))
|
||||
return 0
|
||||
|
||||
|
||||
def cmd_set(args: argparse.Namespace) -> int:
|
||||
fields = {}
|
||||
for item in args.field:
|
||||
if "=" not in item:
|
||||
raise SystemExit(f"bad field (want key=value): {item}")
|
||||
k, v = item.split("=", 1)
|
||||
if k not in COLUMNS or k == "loop_id":
|
||||
raise SystemExit(f"invalid field: {k}")
|
||||
fields[k] = v
|
||||
set_fields(Path(args.registry), args.loop_id, fields, force=True)
|
||||
return 0
|
||||
|
||||
|
||||
def cmd_export(args: argparse.Namespace) -> int:
|
||||
rows = load(Path(args.registry))
|
||||
out = sys.stdout
|
||||
close = False
|
||||
if args.output and args.output != "-":
|
||||
out = open(args.output, "w") # noqa: SIM115
|
||||
close = True
|
||||
try:
|
||||
n = 0
|
||||
for r in sorted(rows.values(), key=lambda x: x["loop_id"]):
|
||||
if args.batch and r.get("batch_id") != args.batch:
|
||||
continue
|
||||
if args.status:
|
||||
allowed = set(args.status.split(","))
|
||||
if r.get("status") not in allowed:
|
||||
continue
|
||||
if args.target and r.get("target_cluster") != args.target:
|
||||
continue
|
||||
if args.skip_done and is_done_status(r.get("status", "")):
|
||||
continue
|
||||
ctx = r.get("kube_context") or ""
|
||||
lid = r.get("loop_id") or ""
|
||||
if not ctx or not lid:
|
||||
continue
|
||||
# migrate-list line; optional target hint as 3rd col if cnpg/selectel
|
||||
line = f"{ctx} {lid}"
|
||||
if args.with_target_hint and r.get("target_cluster"):
|
||||
line += " cnpg"
|
||||
out.write(line + "\n")
|
||||
n += 1
|
||||
if args.output and args.output != "-":
|
||||
print(f"exported {n} rows", file=sys.stderr)
|
||||
finally:
|
||||
if close:
|
||||
out.close()
|
||||
return 0
|
||||
|
||||
|
||||
def cmd_summary(args: argparse.Namespace) -> int:
|
||||
rows = load(Path(args.registry))
|
||||
from collections import Counter
|
||||
|
||||
by_status = Counter(r.get("status") or "" for r in rows.values())
|
||||
by_batch = Counter(r.get("batch_id") or "(none)" for r in rows.values())
|
||||
by_yc = Counter(r.get("yc_cluster") or "(none)" for r in rows.values())
|
||||
total_bytes = 0
|
||||
sized = 0
|
||||
for r in rows.values():
|
||||
b = r.get("size_bytes") or ""
|
||||
if b.isdigit():
|
||||
total_bytes += int(b)
|
||||
sized += 1
|
||||
print(f"rows={len(rows)} sized={sized} total≈{total_bytes / (1024**3):.1f} GiB")
|
||||
print("status:", dict(by_status))
|
||||
print("batch:", dict(sorted(by_batch.items())))
|
||||
print("yc_cluster top:", by_yc.most_common(15))
|
||||
return 0
|
||||
|
||||
|
||||
def cmd_is_done(args: argparse.Namespace) -> int:
|
||||
rows = load(Path(args.registry))
|
||||
row = rows.get(args.loop_id)
|
||||
if row and is_done_status(row.get("status", "")):
|
||||
return 0
|
||||
return 1
|
||||
|
||||
|
||||
def main() -> int:
|
||||
p = argparse.ArgumentParser(description="Loop migrate registry.tsv helpers")
|
||||
p.add_argument("--registry", default="", help="path to registry.tsv")
|
||||
sub = p.add_subparsers(dest="cmd", required=True)
|
||||
|
||||
g = sub.add_parser("get", help="print row or one field")
|
||||
g.add_argument("loop_id")
|
||||
g.add_argument("field", nargs="?")
|
||||
g.set_defaults(func=cmd_get)
|
||||
|
||||
s = sub.add_parser("set", help="set fields: key=value …")
|
||||
s.add_argument("loop_id")
|
||||
s.add_argument("field", nargs="+")
|
||||
s.set_defaults(func=cmd_set)
|
||||
|
||||
e = sub.add_parser("export", help="export migrate-list lines")
|
||||
e.add_argument("--batch", default="")
|
||||
e.add_argument("--status", default="", help="comma-separated statuses")
|
||||
e.add_argument("--target", default="")
|
||||
e.add_argument("--skip-done", action="store_true")
|
||||
e.add_argument("--with-target-hint", action="store_true")
|
||||
e.add_argument("-o", "--output", default="-")
|
||||
e.set_defaults(func=cmd_export)
|
||||
|
||||
u = sub.add_parser("summary")
|
||||
u.set_defaults(func=cmd_summary)
|
||||
|
||||
d = sub.add_parser("is-done", help="exit 0 if status is migrated/verified/cutover/done")
|
||||
d.add_argument("loop_id")
|
||||
d.set_defaults(func=cmd_is_done)
|
||||
|
||||
args = p.parse_args()
|
||||
if not args.registry:
|
||||
args.registry = str(Path(__file__).resolve().parent / "registry.tsv")
|
||||
return args.func(args)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
sys.exit(main())
|
||||
Reference in New Issue
Block a user