python script to duplicate tenants

scripts/ps_duplicate_tenant.py

@@ -0,0 +1,73 @@
+# Usage from top of repo:
+#  poetry run python3 ./scripts/ps_duplicate_tenant.py c66e2e233057f7f05563caff664ecb14 .neon/remote_storage_local_fs
+import argparse
+import shutil
+import subprocess
+import time
+from pathlib import Path
+
+import sys
+
+sys.path.append("test_runner")
+
+from fixtures.pageserver.http import PageserverHttpClient
+from fixtures.types import TenantId
+
+parser = argparse.ArgumentParser(description="Duplicate tenant script.")
+parser.add_argument("initial_tenant", type=str, help="Initial tenant")
+parser.add_argument("remote_storage_local_fs_root", type=Path, help="Remote storage local fs root")
+parser.add_argument("--ncopies", type=int, help="Number of copies")
+parser.add_argument("--numthreads", type=int, default=1, help="Number of threads")
+parser.add_argument("--port", type=int, default=9898, help="Pageserver management api port")
+
+args = parser.parse_args()
+
+initial_tenant = args.initial_tenant
+remote_storage_local_fs_root: Path = args.remote_storage_local_fs_root
+ncopies = args.ncopies
+numthreads = args.numthreads
+
+new_tenant = TenantId.generate()
+print(f"New tenant: {new_tenant}")
+
+client = PageserverHttpClient(args.port, lambda: None)
+
+src_tenant_gen = int(client.tenant_status(initial_tenant)["generation"])
+
+assert remote_storage_local_fs_root.is_dir(), f"{remote_storage_local_fs_root} is not a directory"
+
+src_timelines_dir: Path = remote_storage_local_fs_root / "tenants" / initial_tenant / "timelines"
+assert src_timelines_dir.is_dir(), f"{src_timelines_dir} is not a directory"
+
+dst_timelines_dir: Path = remote_storage_local_fs_root / "tenants" / str(new_tenant) / "timelines"
+dst_timelines_dir.parent.mkdir(parents=False, exist_ok=False)
+dst_timelines_dir.mkdir(parents=False, exist_ok=False)
+
+for tl in src_timelines_dir.iterdir():
+    src_tl_dir = src_timelines_dir / tl.name
+    assert src_tl_dir.is_dir(), f"{src_tl_dir} is not a directory"
+    dst_tl_dir = dst_timelines_dir / tl.name
+    dst_tl_dir.mkdir(parents=False, exist_ok=False)
+    for file in tl.iterdir():
+        shutil.copy2(file, dst_tl_dir)
+        if "__" in file.name:
+            cmd = [
+                "./target/debug/pagectl",  # TODO: abstract this like the other binaries
+                "layer",
+                "rewrite-summary",
+                str(dst_tl_dir / file.name),
+                "--new-tenant-id",
+                str(new_tenant),
+            ]
+            subprocess.run(cmd, check=True)
+
+client.tenant_attach(new_tenant, generation=src_tenant_gen)
+
+while True:
+    status = client.tenant_status(new_tenant)
+    if status["state"]["slug"] == "Active":
+        break
+    print("Waiting for tenant to be active..., is: " + status["state"]["slug"])
+    time.sleep(1)
+
+print("Tenant is active: " + str(new_tenant))