ripple-sql 0.1.0__py3-none-any.whl

ripple/usage/collect.py

@@ -0,0 +1,315 @@
+"""Run the user's already-authenticated warehouse CLI and ingest the result.
+
+Ripple never opens a warehouse connection of its own: the child process is
+`snow`, `bq`, or `databricks`, using whatever auth the user set up for it.
+The rows land in a temp file that is aggregated and deleted inside one call,
+so query text never reaches the caller (human or model) and is never stored.
+
+Validation is strict on purpose: a clean exit code AND stdout that parses as
+complete JSON of the expected shape, or the whole collect is refused. Partial
+output silently ingested is how usage numbers end up confidently wrong.
+"""
+
+from __future__ import annotations
+
+import json
+import re
+import subprocess
+import tempfile
+import time
+from pathlib import Path
+
+from ripple.usage import discover
+from ripple.usage.ingest import ingest_file
+
+DEFAULT_TIMEOUT = 60.0
+SNOWFLAKE_INFOSCHEMA_MAX_DAYS = 7
+ACCOUNT_ROW_LIMIT = 100_000
+DATABRICKS_PAGE_SIZE = 1000
+
+SNOWFLAKE_GRANT_HINT = (
+    "Reading everyone's queries needs one admin grant. Have an admin run:\n"
+    "  GRANT IMPORTED PRIVILEGES ON DATABASE SNOWFLAKE TO ROLE <your role>;\n"
+    "or rerun with scope='mine' (your own queries, no permission needed)."
+)
+BIGQUERY_GRANT_HINT = (
+    "Reading everyone's jobs needs bigquery.jobs.listAll on the project "
+    "(roles/bigquery.resourceViewer). Or rerun with scope='mine'."
+)
+
+_COLUMNS = "query_id, query_text, database_name, schema_name, user_name"
+_SNOW_TIME = "TO_VARCHAR(start_time, 'YYYY-MM-DD\"T\"HH24:MI:SSTZH:TZM') AS start_time"
+
+
+class CollectError(Exception):
+    """A refusal with the reason and, when known, the exact fix."""
+
+
+def _snowflake_sql(days: int, scope: str) -> str:
+    if scope == "account":
+        return (
+            f"SELECT {_COLUMNS}, {_SNOW_TIME} FROM SNOWFLAKE.ACCOUNT_USAGE.QUERY_HISTORY "
+            f"WHERE start_time >= DATEADD('day', -{days}, CURRENT_TIMESTAMP()) "
+            f"LIMIT {ACCOUNT_ROW_LIMIT}"
+        )
+    days = min(days, SNOWFLAKE_INFOSCHEMA_MAX_DAYS)
+    return (
+        f"SELECT {_COLUMNS}, {_SNOW_TIME} FROM TABLE(INFORMATION_SCHEMA.QUERY_HISTORY("
+        f"END_TIME_RANGE_START => DATEADD('day', -{days}, CURRENT_TIMESTAMP()), "
+        f"RESULT_LIMIT => 10000))"
+    )
+
+
+def _bigquery_sql(days: int, scope: str, region: str) -> str:
+    view = "JOBS_BY_PROJECT" if scope == "account" else "JOBS_BY_USER"
+    return (
+        "SELECT job_id AS query_id, query AS query_text, user_email AS user_name, "
+        "FORMAT_TIMESTAMP('%Y-%m-%dT%H:%M:%S+00:00', creation_time) AS start_time "
+        f"FROM `region-{region}`.INFORMATION_SCHEMA.{view} "
+        f"WHERE creation_time > TIMESTAMP_SUB(CURRENT_TIMESTAMP(), INTERVAL {days} DAY) "
+        "AND job_type = 'QUERY' AND query IS NOT NULL "
+        "ORDER BY creation_time DESC LIMIT 10000"
+    )
+
+
+def _parse_json_stdout(stdout: str, cli: str):
+    text = stdout.strip()
+    if not text:
+        raise CollectError(f"{cli} exited cleanly but printed nothing; nothing was ingested.")
+    try:
+        return json.loads(text)
+    except json.JSONDecodeError as e:
+        raise CollectError(
+            f"{cli} printed something that is not complete JSON ({e}); refusing to "
+            "ingest a partial or polluted export. First bytes: " + text[:120]
+        ) from None
+
+
+def _rows_from_array(payload, cli: str) -> list[dict]:
+    if isinstance(payload, list) and payload and all(isinstance(r, list) for r in payload):
+        payload = [row for block in payload for row in block]  # multi-statement output
+    if not isinstance(payload, list) or not all(isinstance(r, dict) for r in payload):
+        raise CollectError(f"{cli} returned JSON of an unexpected shape (wanted an array of rows).")
+    return payload
+
+
+def _normalize_snowflake(stdout: str) -> tuple[list[dict], bool]:
+    return _rows_from_array(_parse_json_stdout(stdout, "snow"), "snow"), False
+
+
+def _normalize_bigquery(stdout: str) -> tuple[list[dict], bool]:
+    return _rows_from_array(_parse_json_stdout(stdout, "bq"), "bq"), False
+
+
+def _normalize_databricks(stdout: str) -> tuple[list[dict], bool]:
+    payload = _parse_json_stdout(stdout, "databricks")
+    if not isinstance(payload, dict) or not isinstance(payload.get("res"), list):
+        raise CollectError(
+            "databricks returned JSON without a 'res' list (wanted the "
+            "/api/2.0/sql/history/queries response)."
+        )
+    rows = []
+    for q in payload["res"]:
+        if not isinstance(q, dict):
+            continue
+        rows.append(
+            {
+                "query_id": q.get("query_id"),
+                "query_text": q.get("query_text"),
+                "user_name": q.get("user_name"),
+                "start_time": q.get("query_start_time_ms"),
+            }
+        )
+    return rows, bool(payload.get("has_next_page"))
+
+
+def resolve_connection(platform: str, requested: str | None) -> tuple[str | None, str]:
+    """(connection to pass, how it was chosen). Refuses to pick among several.
+
+    An explicit request always wins. Otherwise only the user's own declared
+    default (default_connection_name, the active gcloud config, a DEFAULT
+    profile) or a single unambiguous entry is used; several candidates with
+    no default is an error that lists them, never a silent pick.
+    """
+    if requested:
+        return requested, "requested"
+    if platform == "snowflake":
+        found = discover.snowflake_connections()
+        names = sorted(found["connections"])
+        if found["default"]:
+            return found["default"], "default_connection_name"
+        if len(names) == 1:
+            return names[0], "only connection configured"
+        if not names:
+            return None, "no connections.toml entries; snow's own default applies"
+        raise CollectError(
+            f"Several Snowflake connections are configured ({', '.join(names)}) and none "
+            "is marked default. Pass connection=<name>; picking one silently could run "
+            "against the wrong account or role."
+        )
+    if platform == "bigquery":
+        found = discover.bigquery_configurations()
+        active = found["active"]
+        project = found["configurations"].get(active) if active else None
+        if project:
+            return None, f"active gcloud configuration '{active}' (project {project})"
+        return None, "bq's own default project applies"
+    if platform == "databricks":
+        profiles = discover.databricks_profiles()
+        names = sorted(profiles)
+        if "DEFAULT" in profiles:
+            return None, "DEFAULT profile"
+        if len(names) == 1:
+            return names[0], "only profile configured"
+        if not names:
+            return None, "no ~/.databrickscfg profiles; the CLI's own auth applies"
+        raise CollectError(
+            f"Several Databricks profiles are configured ({', '.join(names)}) and none is "
+            "DEFAULT. Pass connection=<profile>; picking one silently could run against "
+            "the wrong workspace."
+        )
+    raise CollectError(f"unknown platform '{platform}' (snowflake, bigquery, or databricks)")
+
+
+def _build_command(
+    platform: str, connection: str | None, days: int, scope: str, region: str
+) -> list[str]:
+    if platform == "snowflake":
+        argv = ["snow", "sql", "-q", _snowflake_sql(days, scope), "--format", "json"]
+        if connection:
+            argv += ["-c", connection]
+        return argv
+    if platform == "bigquery":
+        argv = ["bq", "query", "--nouse_legacy_sql", "--format=json", "--max_rows=10000"]
+        if connection:
+            argv += ["--project_id", connection]
+        return argv + [_bigquery_sql(days, scope, region)]
+    if platform == "databricks":
+        body = {
+            "max_results": DATABRICKS_PAGE_SIZE,
+            "filter_by": {
+                "query_start_time_range": {
+                    "start_time_ms": int((time.time() - days * 86400) * 1000)
+                }
+            },
+        }
+        argv = ["databricks", "api", "post", "/api/2.0/sql/history/queries"]
+        if connection:
+            argv += ["--profile", connection]
+        return argv + ["--json", json.dumps(body)]
+    raise CollectError(f"unknown platform '{platform}' (snowflake, bigquery, or databricks)")
+
+
+NORMALIZERS = {
+    "snowflake": _normalize_snowflake,
+    "bigquery": _normalize_bigquery,
+    "databricks": _normalize_databricks,
+}
+
+
+def _permission_hint(platform: str, scope: str, stderr: str) -> str:
+    lowered = stderr.lower()
+    if scope != "account":
+        return ""
+    if platform == "snowflake" and ("not authorized" in lowered or "does not exist" in lowered):
+        return "\n" + SNOWFLAKE_GRANT_HINT
+    if platform == "bigquery" and ("access denied" in lowered or "permission" in lowered):
+        return "\n" + BIGQUERY_GRANT_HINT
+    return ""
+
+
+def _run_cli(argv: list[str], timeout: float, platform: str, scope: str) -> str:
+    try:
+        proc = subprocess.run(argv, capture_output=True, text=True, timeout=timeout)
+    except FileNotFoundError:
+        raise CollectError(
+            f"{argv[0]} is not installed (or not on PATH). Install and authenticate it "
+            "once, then collect works with no further setup."
+        ) from None
+    except subprocess.TimeoutExpired:
+        raise CollectError(
+            f"{argv[0]} did not finish within {timeout:.0f}s. If it opened a browser "
+            "window to sign in, finish signing in there and run collect again."
+        ) from None
+    if proc.returncode != 0:
+        tail = (proc.stderr or proc.stdout or "").strip()[-600:]
+        raise CollectError(
+            f"{argv[0]} exited with code {proc.returncode}; nothing was ingested.\n{tail}"
+            + _permission_hint(platform, scope, tail)
+        )
+    return proc.stdout
+
+
+def collect(
+    platform: str,
+    model_names: list[str],
+    connection: str | None = None,
+    days: int = 7,
+    scope: str = "mine",
+    region: str = "us",
+    timeout: float = DEFAULT_TIMEOUT,
+    model_aliases: dict[str, set[str]] | None = None,
+) -> dict:
+    """One collect: resolve connection, run the CLI, validate, ingest, clean up.
+
+    Returns the usage store (aggregates only). Raises CollectError with the
+    reason and the fix; a failed collect never half-writes anything.
+    """
+    if platform not in NORMALIZERS:
+        raise CollectError(f"unknown platform '{platform}' (snowflake, bigquery, or databricks)")
+    if scope not in ("mine", "account"):
+        raise CollectError(f"unknown scope '{scope}' (mine or account)")
+    if not re.fullmatch(r"[A-Za-z0-9-]+", region or ""):
+        # the region lands inside a backtick-quoted identifier; anything
+        # beyond a plain region name could smuggle SQL into the user's
+        # authenticated bq session
+        raise CollectError(f"'{region}' is not a BigQuery region name (letters, digits, hyphens)")
+    requested_days = max(1, int(days))
+    days = requested_days
+    if platform == "snowflake" and scope == "mine":
+        # INFORMATION_SCHEMA retains 7 days; asking for more must not
+        # silently claim a window the query never covered
+        days = min(days, SNOWFLAKE_INFOSCHEMA_MAX_DAYS)
+    connection, chosen_by = resolve_connection(platform, connection)
+    argv = _build_command(platform, connection, days, scope, region)
+    stdout = _run_cli(argv, timeout, platform, scope)
+    records, cli_truncated = NORMALIZERS[platform](stdout)
+    if not records:
+        window = f"{days} day(s)" + (
+            f" (clamped from {requested_days}; INFORMATION_SCHEMA keeps 7 days)"
+            if days != requested_days
+            else ""
+        )
+        raise CollectError(
+            f"The export came back empty: no queries in the last {window} for this "
+            "connection and scope. Try a longer window (days=...) or scope='account'."
+        )
+    with tempfile.TemporaryDirectory(prefix="ripple-collect-") as tmp:
+        spool = Path(tmp) / f"{platform}-history.jsonl"
+        spool.write_text(
+            "".join(json.dumps(r, default=str) + "\n" for r in records), encoding="utf-8"
+        )
+        store = ingest_file(spool, model_names, dialect=platform, model_aliases=model_aliases)
+    store["manifest"]["collected"] = {
+        "platform": platform,
+        "connection": connection or "(CLI default)",
+        "chosen_by": chosen_by,
+        # the Databricks history API has no mine/account split: it returns
+        # whatever this login can see, which for an admin is everyone
+        "scope": "visible-to-this-login" if platform == "databricks" else scope,
+        "days": days,
+        "command": " ".join(argv),
+    }
+    if days != requested_days:
+        store["manifest"]["collected"]["requested_days"] = requested_days
+    if cli_truncated:
+        store["manifest"]["truncated"] = (
+            f"yes: the CLI returned its first page of {DATABRICKS_PAGE_SIZE} results and "
+            "reported more; the window was cut short"
+        )
+    elif platform == "snowflake" and scope == "account" and len(records) == ACCOUNT_ROW_LIMIT:
+        store["manifest"]["truncated"] = (
+            f"likely: exactly {ACCOUNT_ROW_LIMIT:,} rows is the collect limit for "
+            "account scope, so the window was probably cut short"
+        )
+    return store

ripple/usage/discover.py

@@ -0,0 +1,190 @@
+"""Find the warehouse CLIs already on this machine and their configured
+connections, by reading config files only. No subprocess, no network.
+
+Only connection NAMES and non-secret fields (authenticator, host, project)
+ever leave this module. Tokens and passwords in the same files are dropped
+at parse time so no caller can surface one by accident.
+"""
+
+from __future__ import annotations
+
+import configparser
+import os
+import re
+import shutil
+from pathlib import Path
+
+CLIS = {"snowflake": "snow", "bigquery": "bq", "databricks": "databricks"}
+
+SAFE_CONNECTION_KEYS = ("authenticator", "account", "host")
+
+
+def _parse_toml_tables(text: str) -> tuple[dict[str, dict], dict]:
+    """(tables, top_level) for the flat TOML that snow writes.
+
+    tomllib on 3.11+; a line parser on 3.10. Either way, only string-ish
+    values survive and only SAFE_CONNECTION_KEYS are kept per table.
+    """
+    try:
+        import tomllib
+
+        try:
+            data = tomllib.loads(text)
+        except tomllib.TOMLDecodeError:
+            return {}, {}
+        top = {k: v for k, v in data.items() if isinstance(v, str)}
+        tables = {}
+        for name, body in data.items():
+            if isinstance(body, dict):
+                tables[name] = body
+        return tables, top
+    except ImportError:
+        pass
+    tables: dict[str, dict] = {}
+    top: dict = {}
+    section = None
+    for line in text.splitlines():
+        line = line.strip()
+        if not line or line.startswith("#"):
+            continue
+        header = re.match(r"\[\s*([A-Za-z0-9_.-]+)\s*\]$", line)
+        if header:
+            section = header.group(1)
+            tables.setdefault(section, {})
+            continue
+        kv = re.match(r"([A-Za-z0-9_-]+)\s*=\s*(.+)$", line)
+        if not kv:
+            continue
+        value = kv.group(2).strip().strip("\"'")
+        if section is None:
+            top[kv.group(1)] = value
+        else:
+            tables[section][kv.group(1)] = value
+    return tables, top
+
+
+def _connection_entry(body: dict) -> dict:
+    return {k: str(body[k]) for k in SAFE_CONNECTION_KEYS if isinstance(body.get(k), str)}
+
+
+def snowflake_connections() -> dict:
+    """{"connections": {name: safe fields}, "default": name | None}.
+
+    connections.toml uses [name] sections; config.toml nests them as
+    [connections.name] and may set default_connection_name. Both files and
+    both shapes are read; SNOWFLAKE_HOME relocates the directory.
+    """
+    home = Path(os.environ.get("SNOWFLAKE_HOME") or Path.home() / ".snowflake")
+    connections: dict[str, dict] = {}
+    default = os.environ.get("SNOWFLAKE_DEFAULT_CONNECTION_NAME")
+    for filename in ("connections.toml", "config.toml"):
+        try:
+            text = (home / filename).read_text(encoding="utf-8")
+        except OSError:
+            continue
+        tables, top = _parse_toml_tables(text)
+        if not default and isinstance(top.get("default_connection_name"), str):
+            default = top["default_connection_name"]
+        for name, body in tables.items():
+            if name == "connections":
+                # [connections] with nested tables (tomllib path)
+                for sub, sub_body in body.items():
+                    if isinstance(sub_body, dict):
+                        connections.setdefault(sub, _connection_entry(sub_body))
+                continue
+            if name.startswith("connections."):
+                name = name.split(".", 1)[1]
+            elif filename == "config.toml":
+                continue  # config.toml top tables (cli, logs, ...) are not connections
+            connections.setdefault(name, _connection_entry(body))
+    return {"connections": connections, "default": default}
+
+
+def bigquery_configurations() -> dict:
+    """{"configurations": {name: project | None}, "active": name | None}.
+
+    bq authenticates through gcloud, so a "connection" here is a gcloud
+    configuration and the project it points at.
+    """
+    home = Path(os.environ.get("CLOUDSDK_CONFIG") or Path.home() / ".config" / "gcloud")
+    active = os.environ.get("CLOUDSDK_ACTIVE_CONFIG_NAME")
+    if not active:
+        try:
+            active = (home / "active_config").read_text(encoding="utf-8").strip() or None
+        except OSError:
+            active = None
+    configurations: dict[str, str | None] = {}
+    try:
+        config_files = sorted((home / "configurations").glob("config_*"))
+    except OSError:
+        config_files = []
+    for path in config_files:
+        name = path.name.removeprefix("config_")
+        parser = configparser.ConfigParser()
+        try:
+            parser.read_string(path.read_text(encoding="utf-8"))
+        except (OSError, configparser.Error):
+            continue
+        configurations[name] = parser.get("core", "project", fallback=None)
+    return {"configurations": configurations, "active": active}
+
+
+def databricks_profiles() -> dict:
+    """{name: {"host": ...}}. The token field in the same file is never read out."""
+    path = Path(os.environ.get("DATABRICKS_CONFIG_FILE") or Path.home() / ".databrickscfg")
+    # default_section="@" so [DEFAULT] shows up as an ordinary profile
+    parser = configparser.ConfigParser(default_section="@", interpolation=None)
+    try:
+        parser.read_string(path.read_text(encoding="utf-8"))
+    except (OSError, configparser.Error):
+        return {}
+    profiles: dict[str, dict] = {}
+    for name in parser.sections():
+        entry = {}
+        host = parser.get(name, "host", fallback=None)
+        if host:
+            entry["host"] = host
+        profiles[name] = entry
+    return profiles
+
+
+def capabilities() -> dict:
+    """What is installed and configured, per platform. Names only, no secrets.
+
+    This is the first-call answer for collect: the caller shows it to the
+    user and asks which connection to use. Nothing here picks one.
+    """
+    snowflake = snowflake_connections()
+    bigquery = bigquery_configurations()
+    databricks = databricks_profiles()
+    return {
+        "snowflake": {
+            "cli": "snow",
+            "cli_found": shutil.which("snow") is not None,
+            "connections": sorted(snowflake["connections"]),
+            "default": snowflake["default"],
+            "browser_sso": sorted(
+                name
+                for name, body in snowflake["connections"].items()
+                if body.get("authenticator", "").lower() == "externalbrowser"
+            ),
+        },
+        "bigquery": {
+            "cli": "bq",
+            "cli_found": shutil.which("bq") is not None,
+            "configurations": [
+                {"name": name, "project": project}
+                for name, project in sorted(bigquery["configurations"].items())
+            ],
+            "active": bigquery["active"],
+        },
+        "databricks": {
+            "cli": "databricks",
+            "cli_found": shutil.which("databricks") is not None,
+            "profiles": [{"name": name, **body} for name, body in sorted(databricks.items())],
+        },
+        "note": (
+            "Pick the platform and connection with the user; when several "
+            "connections exist, never choose one for them."
+        ),
+    }