From 9d903b6e8b6a91b79257de193d1c91bb247dbca5 Mon Sep 17 00:00:00 2001
From: Dmitry Ivanov <ivadmi5@gmail.com>
Date: Thu, 29 Sep 2022 17:40:36 +0300
Subject: [PATCH] Improvements

---
 hack/demo.py | 74 ++++++++++++++++++++++++++++++++++++++--------------
 1 file changed, 54 insertions(+), 20 deletions(-)
 mode change 100644 => 100755 hack/demo.py

diff --git a/hack/demo.py b/hack/demo.py
old mode 100644
new mode 100755
index 9c819dbbb2..d5a98b6923
--- a/hack/demo.py
+++ b/hack/demo.py
@@ -1,47 +1,81 @@
+#!/usr/bin/env python3
+
+import argparse
 import json
-from pathlib import Path
 import subprocess
-import testgres
 import sys
+import testgres
+import textwrap
+import uuid
+
+from pathlib import Path
 
 
-def from_backup_at(backup_dir: Path):
+def make_tarfile(output_filename, source_dir):
+    cmd = ["tar", "--transform=s/\\.//", "-C", str(source_dir), "-cvf", str(output_filename), "."]
+    print('Command: ', ' '.join(cmd))
+    r = subprocess.check_output(cmd).decode()
+    print(textwrap.indent(r, '> '))
+
+
+def create_tenant(tenant_id):
+    cmd = f"target/debug/neon_local tenant create --tenant-id {tenant_id}"
+    print("Run command:", cmd)
+    r = subprocess.check_output(cmd.split()).decode()
+    print(textwrap.indent(r, '> '))
+
+
+def from_backup_at(args, backup_dir: Path):
     manifest = json.loads((backup_dir / "data" / "backup_manifest").read_text())
     start_lsn = manifest["WAL-Ranges"][0]["Start-LSN"]
     end_lsn = manifest["WAL-Ranges"][0]["End-LSN"]
+    print("LSNs:", start_lsn, end_lsn)
+
+    print("Make tarball")
+    tar = Path('/tmp/base.tar')
+    make_tarfile(tar, backup_dir / 'data')
 
     cmd = (
         "target/debug/neon_local timeline import "
-        f"--tenant-id {tenant_id} "
+        f"--tenant-id {args.tenant_id} "
         f"--base-lsn {start_lsn} "
         f"--end-lsn {end_lsn} "
-        f"--base-tarfile {backup_dir / 'data' / 'base.tar'} "
-        f"--timeline-id {timeline_id} "
-        f"--node-name {node_name}"
+        f"--base-tarfile {tar} "
+        f"--timeline-id {args.timeline_id} "
+        f"--node-name {args.node}"
     )
 
-    r = subprocess.check_output(cmd.split())
-    print(r)
+    print("Run neon_local")
+    r = subprocess.check_output(cmd.split()).decode()
+    print(textwrap.indent(r, '> '))
 
 
-def main(tenant_id, timeline_id, node_name):
+def main(args):
+    print("Create a node")
     node = testgres.get_new_node()
-    node.init(allow_streaming=True).start()
-    node.pgbench_init(scale=2)
+    node.init(unix_sockets=False, allow_streaming=True).start()
+    node.execute("""
+        create table foo as select 1;
+    """)
+    # node.pgbench_init(scale=1)
 
-    backup = node.backup(backup_format="t")
+    print("Create a backup")
+    backup = node.backup()
     backup_dir = Path(backup.base_dir)
 
-    from_backup_at(backup_dir)
+    print("Import a backup")
+    create_tenant(args.tenant_id)
+    from_backup_at(args, backup_dir)
     print(backup_dir)
 
 
 if __name__ == "__main__":
-    tenant_id = "56fc742b0993a7adfd63fe37daa8a6ed"  # sys.argv[1]
-    timeline_id = "56fc742b0993a7adfd63fe37daa8a7ed"  # sys.argv[2]
-    node_name = sys.argv[3]
+    tenant_id = uuid.uuid4().hex
 
-    dir = sys.argv[4]
-    from_backup_at(Path(dir))
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--tenant-id", default=tenant_id)
+    parser.add_argument("--timeline-id", default=tenant_id)
+    parser.add_argument("node")
 
-    # main(tenant_id, timeline_id, node_name)
+    args = parser.parse_args(sys.argv[1:])
+    main(args)