#!/usr/bin/env python3 """Discover allowed BigQuery metadata without reading table rows.""" from __future__ import annotations import argparse import hashlib import json import os import time from concurrent.futures import ThreadPoolExecutor from pathlib import Path from typing import Any, Callable, Optional from runtime_config import load_runtime_config, runtime_cache_root from semantic_model import ( allowed_projects, is_dataset_allowed, is_table_allowed, load_semantic_model, resolve_entity, ) def write_json(value: Any) -> None: print(json.dumps(value, indent=2, ensure_ascii=False)) def cache_path(operation: str, identifier: str) -> Path: digest = hashlib.sha256(identifier.encode("utf-8")).hexdigest()[:20] return runtime_cache_root() / "catalog" / f"{operation}-{digest}.json" def cached( operation: str, identifier: str, ttl_seconds: int, refresh: bool, loader: Callable[[], dict[str, Any]], ) -> dict[str, Any]: path = cache_path(operation, identifier) if not refresh and path.is_file() and time.time() - path.stat().st_mtime <= ttl_seconds: value = json.loads(path.read_text(encoding="utf-8")) value["cache"] = {"hit": True, "path": str(path)} return value value = loader() value["cache"] = {"hit": False, "path": str(path)} path.parent.mkdir(parents=True, exist_ok=True) path.write_text(json.dumps(value, indent=2, ensure_ascii=False) + "\n", encoding="utf-8") return value def split_dataset(value: str) -> tuple[str, str]: parts = value.strip().strip("`").split(".") if len(parts) != 2 or not all(parts): raise ValueError("Dataset must be fully qualified as project.dataset.") return parts[0], parts[1] def split_table(value: str) -> tuple[str, str, str]: parts = value.strip().strip("`").split(".") if len(parts) != 3 or not all(parts): raise ValueError("Table must be fully qualified as project.dataset.table.") return parts[0], parts[1], parts[2] def field_json(field: Any) -> dict[str, Any]: value = { "name": field.name, "type": field.field_type, "mode": field.mode, "description": field.description, } if field.fields: value["fields"] = [field_json(child) for child in field.fields] return value def table_description(client: Any, common: dict[str, Any], table_id: str) -> dict[str, Any]: table = client.get_table(table_id) return { **common, "table": table_id, "type": table.table_type, "description": table.description, "location": table.location, "num_rows": table.num_rows, "num_bytes": table.num_bytes, "partition_field": getattr(table.time_partitioning, "field", None), "clustering_fields": table.clustering_fields, "schema": [field_json(field) for field in table.schema], } def parse_args() -> argparse.Namespace: parser = argparse.ArgumentParser(description="Read allowed BigQuery catalog metadata.") parser.add_argument("--project", default=os.environ.get("BQ_JOB_PROJECT") or None) parser.add_argument("--semantic-model", type=Path) parser.add_argument("--cache-ttl", type=int, default=3600) parser.add_argument("--refresh", action="store_true") subparsers = parser.add_subparsers(dest="operation", required=True) list_datasets = subparsers.add_parser("list-datasets") list_datasets.add_argument("--source-project") list_tables = subparsers.add_parser("list-tables") list_tables.add_argument("--dataset", required=True) describe = subparsers.add_parser("describe-table") describe.add_argument("table") search = subparsers.add_parser("search-columns") search.add_argument("--dataset", required=True) search.add_argument("--contains", required=True) entity = subparsers.add_parser("resolve-entity") entity.add_argument("name") return parser.parse_args() def main() -> int: config_path = load_runtime_config() args = parse_args() model = load_semantic_model(args.semantic_model) if args.operation == "resolve-entity": entity = resolve_entity(model, args.name) write_json( { "entity": entity, "found": entity is not None, "semantic_model_path": model["_path"], } ) return 0 if entity else 4 try: from google.cloud import bigquery except ImportError: raise RuntimeError("google-cloud-bigquery is required; run bootstrap.py.") client = bigquery.Client(project=args.project) common = { "config_path": str(config_path) if config_path else None, "job_project": client.project, "semantic_model_path": model["_path"], } if args.operation == "list-datasets": projects = [args.source_project] if args.source_project else allowed_projects(model) disallowed = [project for project in projects if project not in allowed_projects(model)] if disallowed: raise ValueError(f"Project is outside the configured allowlist: {disallowed[0]}") def load_datasets() -> dict[str, Any]: datasets = [] for project in projects: for item in client.list_datasets(project=project): dataset_id = f"{item.project}.{item.dataset_id}" if is_dataset_allowed(model, dataset_id): datasets.append({"dataset": dataset_id}) return {**common, "datasets": sorted(datasets, key=lambda item: item["dataset"])} write_json(cached("datasets", ",".join(projects), args.cache_ttl, args.refresh, load_datasets)) return 0 if args.operation in {"list-tables", "search-columns"}: project, dataset = split_dataset(args.dataset) dataset_id = f"{project}.{dataset}" if not is_dataset_allowed(model, dataset_id): raise ValueError(f"Dataset {dataset_id} is outside the configured allowlist.") def load_tables() -> dict[str, Any]: tables = [] for item in client.list_tables(dataset_id): table_id = f"{item.project}.{item.dataset_id}.{item.table_id}" if is_table_allowed(model, table_id): tables.append({"table": table_id, "type": item.table_type}) return {**common, "dataset": dataset_id, "tables": tables} table_listing = cached("tables", dataset_id, args.cache_ttl, args.refresh, load_tables) if args.operation == "list-tables": write_json(table_listing) return 0 needle = args.contains.lower() def describe(item: dict[str, Any]) -> dict[str, Any]: table_id = item["table"] return cached( "table", table_id, args.cache_ttl, args.refresh, lambda: table_description(client, common, table_id), ) matches = [] with ThreadPoolExecutor(max_workers=12) as executor: descriptions = executor.map(describe, table_listing["tables"]) for description in descriptions: for field in description["schema"]: if needle in field["name"].lower() or needle in (field["description"] or "").lower(): matches.append( { "table": description["table"], "column": field["name"], "type": field["type"], "mode": field["mode"], "description": field["description"], } ) write_json( { **common, "dataset": dataset_id, "contains": args.contains, "matches": matches, "table_cache": table_listing["cache"], } ) return 0 project, dataset, table_name = split_table(args.table) table_id = f"{project}.{dataset}.{table_name}" if not is_table_allowed(model, table_id): raise ValueError(f"Table {table_id} is outside the configured allowlist.") def load_description() -> dict[str, Any]: return table_description(client, common, table_id) write_json(cached("table", table_id, args.cache_ttl, args.refresh, load_description)) return 0 if __name__ == "__main__": raise SystemExit(main())