Sample docs loader (#365)

trustgraph-cli/scripts/tg-init-trustgraph

@@ -0,0 +1,221 @@
+#!/usr/bin/env python3
+
+"""
+Initialises Pulsar with Trustgraph tenant / namespaces & policy.
+"""
+
+import requests
+import time
+import argparse
+import json
+
+from trustgraph.clients.config_client import ConfigClient
+
+default_pulsar_admin_url = "http://pulsar:8080"
+default_pulsar_host = "pulsar://pulsar:6650"
+subscriber = "tg-init-pulsar"
+
+def get_clusters(url):
+
+    print("Get clusters...", flush=True)
+
+    resp = requests.get(f"{url}/admin/v2/clusters")
+
+    if resp.status_code != 200: raise RuntimeError("Could not fetch clusters")
+
+    return resp.json()
+
+def ensure_tenant(url, tenant, clusters):
+
+    resp = requests.get(f"{url}/admin/v2/tenants/{tenant}")
+
+    if resp.status_code == 200:
+        print(f"Tenant {tenant} already exists.", flush=True)
+        return
+
+    resp = requests.put(
+        f"{url}/admin/v2/tenants/{tenant}",
+        json={
+            "adminRoles": [],
+            "allowedClusters": clusters,
+        }
+    )
+
+    if resp.status_code != 204:
+        print(resp.text, flush=True)
+        raise RuntimeError("Tenant creation failed.")
+
+    print(f"Tenant {tenant} created.", flush=True)
+
+def ensure_namespace(url, tenant, namespace, config):
+
+    resp = requests.get(f"{url}/admin/v2/namespaces/{tenant}/{namespace}")
+
+    if resp.status_code == 200:
+        print(f"Namespace {tenant}/{namespace} already exists.", flush=True)
+        return
+
+    resp = requests.put(
+        f"{url}/admin/v2/namespaces/{tenant}/{namespace}",
+        json=config,
+    )
+
+    if resp.status_code != 204:
+        print(resp.status_code, flush=True)
+        print(resp.text, flush=True)
+        raise RuntimeError(f"Namespace {tenant}/{namespace} creation failed.")
+
+    print(f"Namespace {tenant}/{namespace} created.", flush=True)
+
+def ensure_config(config, pulsar_host, pulsar_api_key):
+
+    cli = ConfigClient(
+        subscriber=subscriber,
+        pulsar_host=pulsar_host,
+        pulsar_api_key=pulsar_api_key,
+    )
+
+    while True:
+
+        try:
+
+            print("Get current config...", flush=True)
+            current, version = cli.config(timeout=5)
+
+        except Exception as e:
+
+            print("Exception:", e, flush=True)
+            time.sleep(2)
+            print("Retrying...", flush=True)
+            continue
+
+        print("Current config version is", version, flush=True)
+
+        if version != 0:
+            print("Already updated, not updating config.  Done.", flush=True)
+            return
+
+        print("Config is version 0, updating...", flush=True)
+
+        batch = []
+
+        for type in config:
+            for key in config[type]:
+                print(f"Adding {type}/{key} to update.", flush=True)
+                batch.append({
+                    "type": type,
+                    "key": key,
+                    "value": json.dumps(config[type][key]),
+                })
+
+        try:
+            cli.put(batch, timeout=10)
+            print("Update succeeded.", flush=True)
+            break
+        except Exception as e:
+            print("Exception:", e, flush=True)
+            time.sleep(2)
+            print("Retrying...", flush=True)
+            continue
+    
+def init(pulsar_admin_url, pulsar_host, pulsar_api_key, config, tenant):
+
+    clusters = get_clusters(pulsar_admin_url)
+
+    ensure_tenant(pulsar_admin_url, tenant, clusters)
+
+    ensure_namespace(pulsar_admin_url, tenant, "flow", {})
+
+    ensure_namespace(pulsar_admin_url, tenant, "request", {})
+
+    ensure_namespace(pulsar_admin_url, tenant, "response", {
+        "retention_policies": {
+            "retentionSizeInMB": -1,
+            "retentionTimeInMinutes": 3,
+            "subscriptionExpirationTimeMinutes": 30,
+        }
+    })
+
+    ensure_namespace(pulsar_admin_url, tenant, "config", {
+        "retention_policies": {
+            "retentionSizeInMB": 10,
+            "retentionTimeInMinutes": -1,
+            "subscriptionExpirationTimeMinutes": 5,
+        }
+    })
+
+    if config is not None:
+
+        try:
+            print("Decoding config...", flush=True)
+            dec = json.loads(config)
+            print("Decoded.", flush=True)
+        except Exception as e:
+            print("Exception:", e, flush=True)
+            raise e
+
+        ensure_config(dec, pulsar_host, pulsar_api_key)
+
+    else:
+        print("No config to update.", flush=True)
+
+def main():
+
+    parser = argparse.ArgumentParser(
+        prog='tg-init-trustgraph',
+        description=__doc__,
+    )
+
+    parser.add_argument(
+        '-p', '--pulsar-admin-url',
+        default=default_pulsar_admin_url,
+        help=f'Pulsar admin URL (default: {default_pulsar_admin_url})',
+    )
+
+    parser.add_argument(
+        '--pulsar-host',
+        default=default_pulsar_host,
+        help=f'Pulsar host (default: {default_pulsar_host})',
+    )
+
+    parser.add_argument(
+        '--pulsar-api-key',
+        help=f'Pulsar API key',
+    )
+
+    parser.add_argument(
+        '-c', '--config',
+        help=f'Initial configuration to load',
+    )
+
+    parser.add_argument(
+        '-t', '--tenant',
+        default="tg",
+        help=f'Tenant (default: tg)',
+    )
+
+    args = parser.parse_args()
+
+    while True:
+
+        try:
+
+            print(flush=True)
+            print(
+                f"Initialising with Pulsar {args.pulsar_admin_url}...",
+                flush=True
+            )
+            init(**vars(args))
+            print("Initialisation complete.", flush=True)
+            break
+
+        except Exception as e:
+
+            print("Exception:", e, flush=True)
+
+            print("Sleeping...", flush=True)
+            time.sleep(2)
+            print("Will retry...", flush=True)
+
+main()
+