tests/ts_table_create_deepdive.py

#!/usr/bin/env python3
"""
ts_table_create_deepdive.py — isolate WHY logical-table TML import takes ~250s.

Hypotheses tested (2026-08-10 deep dive):
  H1 "instance"  : sebe (staging, measured ~250s in July) is unhealthy; prod
                   secloud is fine.  -> run the same probe on both envs.
  H2 "scoping"   : the app's connection TML has no `database` property, so the
                   connection sees ~507 Snowflake DBs; table create scales with
                   visible external metadata.  -> compare an unscoped vs a
                   database-pinned connection on the SAME instance.

Each probe: create a throwaway connection (variant-specific), TML-import ONE
trivial 2-column DONT_INDEX table against a pre-existing empty physical table,
time it, then delete the logical table + the throwaway connection.

Usage:
  python scratch/ts_table_create_deepdive.py --env "secloud - primary" --variant unscoped
  python scratch/ts_table_create_deepdive.py --env "sebe - se" --variant scoped
  python scratch/ts_table_create_deepdive.py --env "sebe - se" --variant both

Physical table: creates DEMOBUILD.ZPERFTEST.DDPROBE1 in Snowflake if missing
(empty, 2 columns). Snowflake + TS credentials come from .env / Supabase admin
settings exactly like the app.
"""
import argparse
import os
import sys
import time

import requests
import yaml
from dotenv import load_dotenv

ROOT = os.path.abspath(os.path.join(os.path.dirname(__file__), ".."))
sys.path.insert(0, ROOT)
load_dotenv(os.path.join(ROOT, ".env"))

from supabase_client import get_admin_setting  # noqa: E402

DB = os.getenv("DD_PROBE_DB", "DEMOBUILD")
SCHEMA = os.getenv("DD_PROBE_SCHEMA", "ZPERFTEST")
PHYS_TABLE = "DDPROBE1"

TABLE_TML = """\
guid: null
table:
  name: {name}
  db: {db}
  schema: {schema}
  db_table: {phys}
  connection:
    name: {connection}
  columns:
  - name: ID
    db_column_name: ID
    properties:
      column_type: MEASURE
      aggregation: SUM
      index_type: DONT_INDEX
    db_column_properties:
      data_type: INT64
  - name: LABEL
    db_column_name: LABEL
    properties:
      column_type: ATTRIBUTE
      index_type: DONT_INDEX
    db_column_properties:
      data_type: VARCHAR
"""


def resolve_env(label):
    i = 1
    while True:
        lbl = os.getenv(f"TS_ENV_{i}_LABEL", "").strip()
        if not lbl:
            return None, None
        if lbl == label:
            return (os.getenv(f"TS_ENV_{i}_URL", "").strip().rstrip("/"),
                    os.getenv(f"TS_ENV_{i}_KEY_VAR", "").strip())
        i += 1


def get_private_key_pem():
    raw = get_admin_setting("SNOWFLAKE_KP_PK")
    if not raw.startswith("-----BEGIN"):
        import base64
        try:
            raw = base64.b64decode(raw).decode("utf-8")
        except Exception:
            pass
    return raw


def connection_tml(name, scoped_db=None):
    props = [
        {"key": "accountName", "value": get_admin_setting("SNOWFLAKE_ACCOUNT")},
        {"key": "user", "value": get_admin_setting("SNOWFLAKE_KP_USER")},
        {"key": "private_key", "value": get_private_key_pem()},
        {"key": "passphrase", "value": get_admin_setting("SNOWFLAKE_KP_PASSPHRASE", required=False)},
        {"key": "role", "value": get_admin_setting("SNOWFLAKE_ROLE")},
        {"key": "warehouse", "value": get_admin_setting("SNOWFLAKE_WAREHOUSE")},
    ]
    if scoped_db:
        props.append({"key": "database", "value": scoped_db})
    return yaml.dump({
        "guid": None,
        "connection": {
            "name": name,
            "type": "RDBMS_SNOWFLAKE",
            "authentication_type": "KEY_PAIR",
            "properties": props,
            "description": "deep-dive probe — safe to delete",
        },
    }, default_flow_style=False, sort_keys=False)


def ensure_physical_table():
    from snowflake_auth import get_snowflake_connection
    conn = get_snowflake_connection()
    cur = conn.cursor()
    cur.execute(f"CREATE DATABASE IF NOT EXISTS {DB}")
    cur.execute(f"CREATE SCHEMA IF NOT EXISTS {DB}.{SCHEMA}")
    cur.execute(f"CREATE TABLE IF NOT EXISTS {DB}.{SCHEMA}.{PHYS_TABLE} (ID NUMBER, LABEL VARCHAR)")
    cur.close()
    conn.close()
    print(f"[snowflake] {DB}.{SCHEMA}.{PHYS_TABLE} ready")


def auth_session(url, user, secret):
    s = requests.Session()
    s.headers.update({"Content-Type": "application/json", "X-Requested-By": "ThoughtSpot"})
    t0 = time.time()
    r = s.post(f"{url}/api/rest/2.0/auth/token/full",
               json={"username": user, "secret_key": secret, "validity_time_in_sec": 3600}, timeout=60)
    print(f"[auth] HTTP {r.status_code} in {time.time()-t0:.2f}s")
    r.raise_for_status()
    s.headers["Authorization"] = f"Bearer {r.json()['token']}"
    return s


def tml_import(s, url, tmls, timeout=360):
    t0 = time.time()
    r = s.post(f"{url}/api/rest/2.0/metadata/tml/import",
               json={"metadata_tmls": tmls, "import_policy": "PARTIAL", "create_new": True},
               timeout=timeout)
    elapsed = time.time() - t0
    status, guid, err = None, None, None
    if r.status_code == 200:
        obj = (r.json() or [{}])[0]
        resp = obj.get("response", obj)
        status = (resp.get("status") or {}).get("status_code")
        err = (resp.get("status") or {}).get("error_message")
        guid = (resp.get("header") or {}).get("id_guid")
    return elapsed, r.status_code, status, guid, err


def delete_logical_table(s, url, guid):
    r = s.post(f"{url}/api/rest/2.0/metadata/delete",
               json={"metadata": [{"type": "LOGICAL_TABLE", "identifier": guid}]}, timeout=120)
    print(f"[cleanup] delete table {guid}: HTTP {r.status_code}")


def delete_connection(s, url, guid):
    # NB: connection/delete, NOT metadata/delete (CONNECTION not in that enum)
    r = s.post(f"{url}/api/rest/2.0/connection/delete",
               json={"connection_identifier": guid}, timeout=120)
    print(f"[cleanup] delete connection {guid}: HTTP {r.status_code}")


def run_variant(s, url, variant, suffix):
    scoped = DB if variant == "scoped" else None
    conn_name = f"ZPERF_DD_{variant.upper()}_{suffix}"
    print(f"\n=== VARIANT {variant} — connection {conn_name}"
          + (f" (database pinned to {DB})" if scoped else " (no database property, app-style)") + " ===")

    t0 = time.time()
    elapsed, http, status, conn_guid, err = None, None, None, None, None
    r = s.post(f"{url}/api/rest/2.0/metadata/tml/import",
               json={"metadata_tmls": [connection_tml(conn_name, scoped)], "import_policy": "PARTIAL"},
               timeout=300)
    conn_create_s = time.time() - t0
    obj = (r.json() or [{}])[0] if r.status_code == 200 else {}
    resp = obj.get("response", obj)
    conn_guid = (resp.get("header") or {}).get("id_guid")
    conn_status = (resp.get("status") or {}).get("status_code")
    print(f"[conn-create] HTTP {r.status_code} status={conn_status} in {conn_create_s:.2f}s guid={conn_guid}")
    if not conn_guid:
        print(f"[conn-create] FAILED: {str(resp)[:500]}")
        return None

    table_name = f"DDT_{variant.upper()}_{suffix}"
    tml = TABLE_TML.format(name=table_name, db=DB, schema=SCHEMA, phys=PHYS_TABLE, connection=conn_name)
    print(f"[import] creating logical table {table_name} ({len(tml)} bytes) ...")
    elapsed, http, status, table_guid, err = tml_import(s, url, [tml])
    flag = "  <-- GATEWAY TIMEOUT" if http in (502, 503, 504) else ""
    print(f"[import] HTTP {http} status={status} in {elapsed:.1f}s{flag}" + (f" err={err}" if err else ""))

    if table_guid:
        delete_logical_table(s, url, table_guid)
    delete_connection(s, url, conn_guid)
    return {"variant": variant, "conn_create_s": round(conn_create_s, 2),
            "import_s": round(elapsed, 1), "http": http, "status": status}


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--env", required=True, help='TS env label from .env (e.g. "sebe - se")')
    ap.add_argument("--variant", choices=["unscoped", "scoped", "both"], default="both")
    args = ap.parse_args()

    url, secret = resolve_env(args.env)
    user = os.getenv("TEST_USER") or os.getenv("THOUGHTSPOT_USERNAME")
    if not (url and user and secret):
        raise SystemExit(f"Could not resolve env '{args.env}' from .env")

    print(f"Host: {url}\nUser: {user}\nSource: {DB}.{SCHEMA}.{PHYS_TABLE}")
    ensure_physical_table()
    s = auth_session(url, user, secret)

    suffix = time.strftime("%H%M%S")
    variants = ["unscoped", "scoped"] if args.variant == "both" else [args.variant]
    results = [run_variant(s, url, v, suffix) for v in variants]

    print("\n" + "=" * 64)
    print(f"{'variant':12}{'conn create':>12}{'table import':>14}  http/status")
    for res in results:
        if res:
            print(f"{res['variant']:12}{res['conn_create_s']:>11}s{res['import_s']:>13}s  {res['http']}/{res['status']}")
    print("=" * 64)


if __name__ == "__main__":
    main()