#!/usr/bin/env python3"""Minimal BigQuery client: stdlib + `cryptography`. Dry-runs every query and refuses to run past a byte cap.Credentials: path in env GOOGLE_APPLICATION_CREDENTIALS (service-account JSON). Never logged.The NEAR public dataset is billed to the caller; the free sandbox gives 1 TB of scanned data a month."""import base64, json, os, pathlib, sys, time, urllib.error, urllib.requestsys.path.insert(0, str(pathlib.Path(__file__).resolve().parent))from common import MAX_BYTES, OPENER  # noqa: E402from cryptography.hazmat.primitives import hashes, serializationfrom cryptography.hazmat.primitives.asymmetric import paddingTOKEN_URL = "https://oauth2.googleapis.com/token"SCOPE = "https://www.googleapis.com/auth/bigquery"API = "https://bigquery.googleapis.com/bigquery/v2"def _b64(raw):    return base64.urlsafe_b64encode(raw).rstrip(b"=")class BQ:    def __init__(self, cred_path=None, max_bytes=120 * 10**9):        path = cred_path or os.environ.get("GOOGLE_APPLICATION_CREDENTIALS")        if not path:            raise SystemExit("GOOGLE_APPLICATION_CREDENTIALS is not set")        with open(path, encoding="utf8") as fh:            self._cred = json.load(fh)        self.project = self._cred["project_id"]        self.max_bytes = int(max_bytes)        self._token, self._exp = None, 0    def token(self):        if self._token and time.time() < self._exp - 60:            return self._token        now = int(time.time())        head = _b64(json.dumps({"alg": "RS256", "typ": "JWT"}).encode())        claims = _b64(json.dumps({"iss": self._cred["client_email"], "scope": SCOPE, "aud": TOKEN_URL,                                  "iat": now, "exp": now + 3600}).encode())        key = serialization.load_pem_private_key(self._cred["private_key"].encode(), password=None)        sig = key.sign(head + b"." + claims, padding.PKCS1v15(), hashes.SHA256())        body = ("grant_type=urn%3Aietf%3Aparams%3Aoauth%3Agrant-type%3Ajwt-bearer&assertion="                + (head + b"." + claims + b"." + _b64(sig)).decode()).encode()        req = urllib.request.Request(TOKEN_URL, data=body, headers={"content-type": "application/x-www-form-urlencoded"})        try:            with OPENER.open(req, timeout=30) as resp:                out = json.loads(resp.read(1000000))        except urllib.error.HTTPError as exc:            # Google's error code and description name the problem and carry no credentials            try:                info = json.loads(exc.read(4000))                detail = f"{info.get('error')}: {str(info.get('error_description'))[:120]}"            except (ValueError, AttributeError):                detail = "no detail"            raise SystemExit(f"Google token request failed, HTTP {exc.code}, {detail}") from None        self._token, self._exp = out["access_token"], now + int(out.get("expires_in", 3600))        return self._token    def _call(self, url, payload=None, timeout=300):        data = None if payload is None else json.dumps(payload).encode()        req = urllib.request.Request(url, data=data, method="POST" if data else "GET",                                     headers={"authorization": f"Bearer {self.token()}", "content-type": "application/json"})        try:            with OPENER.open(req, timeout=timeout) as resp:                return json.loads(resp.read(MAX_BYTES))        except urllib.error.HTTPError as exc:            # error bodies can name the project and account, and this log may be public: keep only the reason code            try:                reason = json.loads(exc.read(20000))["error"]["errors"][0]["reason"]            except (ValueError, KeyError, IndexError, TypeError):                reason = "unknown"            raise RuntimeError(f"BigQuery HTTP {exc.code}: {reason}") from None    def dry_run(self, sql):        out = self._call(f"{API}/projects/{self.project}/jobs", {            "configuration": {"dryRun": True, "query": {"query": sql, "useLegacySql": False}}})        return int(out["statistics"]["totalBytesProcessed"])    def query(self, sql, label="query", cap=None, poll_s=4, max_wait_s=900):        """Returns (rows as dicts, stats). Raises if the dry run exceeds the cap."""        cap = int(cap or self.max_bytes)        scanned = self.dry_run(sql)        if scanned > cap:            raise RuntimeError(f"{label}: dry run {scanned/1e9:.1f} GB exceeds cap {cap/1e9:.1f} GB")        job = self._call(f"{API}/projects/{self.project}/jobs", {            "configuration": {"query": {"query": sql, "useLegacySql": False, "maximumBytesBilled": str(cap)}}})        ref = job["jobReference"]        loc = ref.get("location", "US")        deadline = time.time() + max_wait_s        while True:            res = self._call(f"{API}/projects/{self.project}/queries/{ref['jobId']}?location={loc}&timeoutMs=20000&maxResults=20000")            if res.get("jobComplete"):                break            if time.time() > deadline:                raise RuntimeError(f"{label}: timed out")            time.sleep(poll_s)        fields = [f["name"] for f in res["schema"]["fields"]]        rows, page = [], res        while True:            rows.extend({n: c.get("v") for n, c in zip(fields, r["f"])} for r in page.get("rows", []))            tok = page.get("pageToken")            if not tok:                break            page = self._call(f"{API}/projects/{self.project}/queries/{ref['jobId']}?location={loc}&pageToken={tok}&maxResults=20000")        return rows, {"bytes_processed": int(res.get("totalBytesProcessed") or scanned), "dry_run_bytes": scanned,                      "rows": len(rows), "job_id": ref["jobId"]}