misc.python.materialize.biglake

  1# Copyright Materialize, Inc. and contributors. All rights reserved.
  2#
  3# Use of this software is governed by the Business Source License
  4# included in the LICENSE file at the root of this repository.
  5#
  6# As of the Change Date specified in that file, in accordance with
  7# the Business Source License, use of this software will be governed
  8# by the Apache License, Version 2.0.
  9
 10"""Helpers for bootstrapping Google Cloud BigLake Iceberg REST catalogs.
 11
 12BigLake (the API is still called BigLake, though it now lives under the
 13"Lakehouse" umbrella) exposes an Iceberg REST catalog at
 14``https://biglake.googleapis.com/iceberg/v1/restcatalog``.
 15
 16Unlike some Iceberg catalogs, BigLake does **not** auto-create catalogs or
 17namespaces on first write. A Materialize Iceberg sink only creates the *table*
 18(see ``load_or_create_table`` in ``src/storage/src/sink/iceberg.rs``), not the
 19namespace, so a sink targeting a missing namespace fails at runtime with::
 20
 21    Failed to create Iceberg table '<t>' in namespace '<ns>':
 22    Tried to create a table under a namespace that does not exist
 23
 24These helpers create the catalog and namespace out of band so the sink can then
 25create its tables. Shared by ``test/gcp`` (per-run, throwaway ``e2e_*``
 26namespaces) and ``test/canary-environment`` (a fixed, long-lived namespace).
 27"""
 28
 29import json
 30import time
 31import urllib.error
 32import urllib.parse
 33import urllib.request
 34
 35import jwt
 36
 37# Blanket scope used to mint tokens for both GCS and BigLake. Matches GCP_SCOPE
 38# in src/storage-types/src/connections/gcp.rs.
 39GCP_SCOPE = "https://www.googleapis.com/auth/cloud-platform"
 40
 41BIGLAKE_REST_BASE = "https://biglake.googleapis.com/iceberg/v1/restcatalog"
 42
 43
 44def mint_gcp_access_token(service_account: dict) -> str:
 45    """Mint an OAuth2 access token from a GCP service-account key (JWT bearer flow)."""
 46    now = int(time.time())
 47    assertion = jwt.encode(
 48        {
 49            "iss": service_account["client_email"],
 50            "scope": GCP_SCOPE,
 51            "aud": "https://oauth2.googleapis.com/token",
 52            "iat": now,
 53            "exp": now + 3600,
 54        },
 55        service_account["private_key"],
 56        algorithm="RS256",
 57    )
 58    body = urllib.parse.urlencode(
 59        {
 60            "grant_type": "urn:ietf:params:oauth:grant-type:jwt-bearer",
 61            "assertion": assertion,
 62        }
 63    ).encode()
 64    req = urllib.request.Request(
 65        "https://oauth2.googleapis.com/token",
 66        data=body,
 67        headers={"Content-Type": "application/x-www-form-urlencoded"},
 68    )
 69    with urllib.request.urlopen(req) as resp:
 70        return json.loads(resp.read())["access_token"]
 71
 72
 73def biglake_request(
 74    method: str, url: str, token: str, project: str
 75) -> urllib.request.Request:
 76    return urllib.request.Request(
 77        url,
 78        method=method,
 79        headers={
 80            "Authorization": f"Bearer {token}",
 81            "x-goog-user-project": project,
 82        },
 83    )
 84
 85
 86def ensure_catalog(token: str, project: str, bucket: str) -> None:
 87    """Create the BigLake Iceberg REST catalog for this bucket if it's missing.
 88
 89    The Iceberg REST `/v1/config?warehouse=gs://<bucket>` lookup resolves the catalog
 90    whose id equals the bucket name; it does not provision one on demand. Without it,
 91    every REST call returns a sanitized 403. Creating it here lets callers bootstrap
 92    their own catalog instead of depending on out-of-band (Pulumi/console) setup.
 93
 94    Credential mode is END_USER: the Materialize sink writes to GCS with the service
 95    account's own credentials (see `connect_rest` for the GCP branch in
 96    mz_storage_types::connections), not credentials vended by the catalog.
 97    """
 98    base = f"{BIGLAKE_REST_BASE}/extensions/projects/{project}/catalogs"
 99
100    # GET returns the catalog if it exists, 404 if not.
101    try:
102        urllib.request.urlopen(
103            biglake_request("GET", f"{base}/{bucket}", token, project)
104        )
105        return
106    except urllib.error.HTTPError as e:
107        if e.code != 404:
108            raise
109
110    create_url = f"{base}?iceberg-catalog-id={urllib.parse.quote(bucket, safe='')}"
111    req = biglake_request("POST", create_url, token, project)
112    req.add_header("Content-Type", "application/json")
113    req.data = json.dumps(
114        {
115            "catalog-type": "CATALOG_TYPE_GCS_BUCKET",
116            "credential-mode": "CREDENTIAL_MODE_END_USER",
117        }
118    ).encode()
119    try:
120        urllib.request.urlopen(req)
121        print(f"created BigLake catalog for gs://{bucket}")
122    except urllib.error.HTTPError as e:
123        # A concurrent run can create the catalog between our GET and POST.
124        if e.code == 409:
125            return
126        body = e.read().decode("utf-8", errors="replace")
127        print(f"BigLake catalog create failed: HTTP {e.code}\n{body}")
128        raise
129
130
131def resolve_warehouse_prefix(token: str, project: str, bucket: str) -> str:
132    """Return the catalog prefix BigLake assigns to this warehouse.
133
134    Iceberg REST clients call GET /v1/config?warehouse=... before any other
135    operation. The catalog's response includes an `overrides.prefix` that the
136    client splices in between `/v1/` and resource paths for every later call:
137
138        {uri}/v1/{prefix}/namespaces/{ns}/tables/{tbl}
139
140    See `RestCatalogConfig::url_prefixed` in iceberg-catalog-rest.
141    """
142    warehouse = f"gs://{bucket}"
143    url = (
144        f"{BIGLAKE_REST_BASE}/v1/config"
145        f"?warehouse={urllib.parse.quote(warehouse, safe='')}"
146    )
147    with urllib.request.urlopen(biglake_request("GET", url, token, project)) as resp:
148        config = json.loads(resp.read())
149    print(f"BigLake /v1/config response: {json.dumps(config)}")
150    return config.get("overrides", {}).get("prefix", "")
151
152
153def catalog_url(prefix: str, suffix: str) -> str:
154    middle = f"{prefix}/" if prefix else ""
155    return f"{BIGLAKE_REST_BASE}/v1/{middle}{suffix}"
156
157
158def table_url(prefix: str, namespace: str, table: str) -> str:
159    ns = urllib.parse.quote(namespace, safe="")
160    tbl = urllib.parse.quote(table, safe="")
161    return catalog_url(prefix, f"namespaces/{ns}/tables/{tbl}")
162
163
164def namespace_url(prefix: str, namespace: str) -> str:
165    ns = urllib.parse.quote(namespace, safe="")
166    return catalog_url(prefix, f"namespaces/{ns}")
167
168
169def create_namespace(
170    token: str, project: str, prefix: str, namespace: str, *, exist_ok: bool = True
171) -> None:
172    """Create the namespace. BigLake doesn't auto-create on first commit.
173
174    With ``exist_ok`` (the default), a 409 from a namespace that already exists
175    is treated as success, so this is safe to call repeatedly against a fixed
176    long-lived namespace.
177    """
178    req = biglake_request("POST", catalog_url(prefix, "namespaces"), token, project)
179    req.add_header("Content-Type", "application/json")
180    req.data = json.dumps({"namespace": [namespace]}).encode()
181    try:
182        urllib.request.urlopen(req)
183    except urllib.error.HTTPError as e:
184        if exist_ok and e.code == 409:
185            return
186        raise
187
188
189def bootstrap_namespace(service_account: dict, bucket: str, namespace: str) -> str:
190    """Ensure the catalog and namespace for ``gs://<bucket>`` exist.
191
192    Mints a token from ``service_account``, ensures the BigLake catalog for the
193    bucket exists, resolves the warehouse prefix, and creates ``namespace`` if it
194    is missing. Idempotent: a no-op once the catalog and namespace exist. Returns
195    the resolved catalog prefix.
196    """
197    project = service_account["project_id"]
198    token = mint_gcp_access_token(service_account)
199    ensure_catalog(token, project, bucket)
200    prefix = resolve_warehouse_prefix(token, project, bucket)
201    create_namespace(token, project, prefix, namespace)
202    print(f"BigLake namespace ready: {namespace} (gs://{bucket})")
203    return prefix