trustgraph/trustgraph-cli/trustgraph/cli/init_trustgraph.py

"""
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, tenant,
        config, config_file,
):

    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)

    elif config_file is not None:

        try:
            print("Decoding config...", flush=True)
            dec = json.load(open(config_file))
            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(
        '-C', '--config-file',
        help=f'Initial configuration to load from file',
    )

    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)

if __name__ == "__main__":
    main()