"""BigQuery data-access driver. DAX-021 creates the dedicated ``data_access/bigquery`` registry entry so dataset effects do not resolve through CMDB cloud integrations. DAX-022 defines the locator-scoped IAM interface. DAX-023 adds the IAM read path for one requested dataset locator. The 2026-08-19 schema-browser extension adds only table metadata reads within one requestable dataset; writes remain reconciler-owned. """ from __future__ import annotations import asyncio import json import re from collections import defaultdict from typing import Any from app.domains.data_access.apply_gate import ( APPLY_MODE_CONFIG_KEY, PERSONAL_APPLY_ENABLED_CONFIG_KEY, data_access_apply_mode_enabled, validate_dataset_locator, ) from app.domains.data_access.constants import ( BIGQUERY_DATA_VIEWER_ROLE, BIGQUERY_GROUP_MEMBER_PREFIX, BIGQUERY_SERVICE_ACCOUNT_MEMBER_PREFIX, BIGQUERY_USER_MEMBER_PREFIX, GCP_SERVICE_ACCOUNT_EMAIL_PATTERN, ) from app.integrations.driver_base import ( DataAccessDatasetFieldSchema, DataAccessDatasetIamBinding, DataAccessDatasetIamPolicy, DataAccessDatasetIamPolicyPatch, DataAccessDatasetIamUpdateResult, DataAccessDatasetTablePage, DataAccessDatasetTableSchema, DataAccessDatasetTableSummary, DataAccessDriver, ) from app.integrations.registry import register_driver from app.integrations.schemas import ( IntegrationFieldInfo, IntegrationTestResult, ) BIGQUERY_DATA_ACCESS_OAUTH_SCOPE = "https://www.googleapis.com/auth/bigquery" BIGQUERY_DATA_ACCESS_REQUIRED_PERMISSIONS = ( "bigquery.datasets.get", "bigquery.datasets.getIamPolicy", "bigquery.datasets.update", "bigquery.datasets.setIamPolicy", "bigquery.tables.get", "bigquery.tables.getIamPolicy", "bigquery.tables.list", "bigquery.tables.setIamPolicy", ) _DATASET_UPDATE_MODE_ACL = "UPDATE_ACL" _ACCESS_POLICY_VERSION_WITH_CONDITIONS = 3 # google-cloud-bigquery table IAM currently accepts only policy version 1. # This is separate from the dataset ACL REST path above, which supports v3. _TABLE_IAM_POLICY_VERSION = 1 _REQUIRED_SERVICE_ACCOUNT_FIELDS = ( "type", "client_email", "private_key", "project_id", ) _BQ_SCOPE = BIGQUERY_DATA_ACCESS_OAUTH_SCOPE _PRIVATE_KEY_RE = re.compile( r"-----BEGIN PRIVATE KEY-----.*?-----END PRIVATE KEY-----", re.DOTALL, ) _DATASET_ACCESS_ROLE_TO_PROVIDER_ROLE = { "READER": BIGQUERY_DATA_VIEWER_ROLE, "WRITER": "roles/bigquery.dataEditor", "OWNER": "roles/bigquery.dataOwner", } def _validate_dataset_locator(project_id: str, dataset_id: str) -> None: validate_dataset_locator(project_id, dataset_id) def _raw_dataset_access_entries(policy: DataAccessDatasetIamPolicy) -> list[Any]: raw_policy = policy.raw_policy or {} raw_access = raw_policy.get("access") if raw_access is None: raw_access = raw_policy.get("access_entries", []) if raw_access is None: return [] if not isinstance(raw_access, list): raise RuntimeError("BigQuery dataset IAM patch requires access entries") return raw_access def _access_entry_member(entry: dict[str, Any]) -> str | None: iam_member = entry.get("iamMember") if isinstance(iam_member, str): member = iam_member.strip() if member.startswith(BIGQUERY_GROUP_MEMBER_PREFIX): return member if member.startswith(BIGQUERY_SERVICE_ACCOUNT_MEMBER_PREFIX): email = member.removeprefix(BIGQUERY_SERVICE_ACCOUNT_MEMBER_PREFIX) if _is_service_account_email(email): return f"{BIGQUERY_SERVICE_ACCOUNT_MEMBER_PREFIX}{email}" if member.startswith(BIGQUERY_USER_MEMBER_PREFIX): email = member.removeprefix(BIGQUERY_USER_MEMBER_PREFIX) if _is_person_email(email): return f"{BIGQUERY_USER_MEMBER_PREFIX}{email}" group_email = entry.get("groupByEmail") if isinstance(group_email, str) and group_email.strip(): return f"{BIGQUERY_GROUP_MEMBER_PREFIX}{group_email.strip()}" service_account_email = entry.get("userByEmail") if ( isinstance(service_account_email, str) and _is_service_account_email(service_account_email) ): return ( f"{BIGQUERY_SERVICE_ACCOUNT_MEMBER_PREFIX}" f"{service_account_email.strip().lower()}" ) if isinstance(service_account_email, str) and _is_person_email(service_account_email): return f"{BIGQUERY_USER_MEMBER_PREFIX}{service_account_email.strip().lower()}" entity_type = str(entry.get("entity_type") or "").strip() entity_id = str(entry.get("entity_id") or "").strip() if entity_type == "groupByEmail" and entity_id: return f"{BIGQUERY_GROUP_MEMBER_PREFIX}{entity_id}" if ( entity_type == "iamMember" and entity_id.startswith(BIGQUERY_GROUP_MEMBER_PREFIX) ): return entity_id if entity_type == "userByEmail" and _is_service_account_email(entity_id): return f"{BIGQUERY_SERVICE_ACCOUNT_MEMBER_PREFIX}{entity_id.lower()}" if entity_type == "userByEmail" and _is_person_email(entity_id): return f"{BIGQUERY_USER_MEMBER_PREFIX}{entity_id.lower()}" if entity_type == "iamMember" and entity_id.startswith( BIGQUERY_SERVICE_ACCOUNT_MEMBER_PREFIX, ): email = entity_id.removeprefix(BIGQUERY_SERVICE_ACCOUNT_MEMBER_PREFIX) if _is_service_account_email(email): return f"{BIGQUERY_SERVICE_ACCOUNT_MEMBER_PREFIX}{email}" if entity_type == "iamMember" and entity_id.startswith(BIGQUERY_USER_MEMBER_PREFIX): email = entity_id.removeprefix(BIGQUERY_USER_MEMBER_PREFIX) if _is_person_email(email): return f"{BIGQUERY_USER_MEMBER_PREFIX}{email}" return None def _access_entry_role(entry: dict[str, Any]) -> str | None: role = str(entry.get("role") or "").strip() if role == "READER" or role == BIGQUERY_DATA_VIEWER_ROLE: return BIGQUERY_DATA_VIEWER_ROLE return None def _is_unconditional_access_entry(entry: dict[str, Any]) -> bool: return "condition" not in entry def _provider_grant_key( role: str, member: str, ) -> tuple[str, str]: if role != BIGQUERY_DATA_VIEWER_ROLE: raise ValueError(f"DAX can only patch {BIGQUERY_DATA_VIEWER_ROLE!r}") if member.startswith(BIGQUERY_GROUP_MEMBER_PREFIX): email = member.removeprefix(BIGQUERY_GROUP_MEMBER_PREFIX).strip().lower() if not email: raise ValueError("DAX group member must include an email") return role, f"{BIGQUERY_GROUP_MEMBER_PREFIX}{email}" if member.startswith(BIGQUERY_SERVICE_ACCOUNT_MEMBER_PREFIX): email = member.removeprefix( BIGQUERY_SERVICE_ACCOUNT_MEMBER_PREFIX, ).strip().lower() if not _is_service_account_email(email): raise ValueError( "DAX serviceAccount: member must be a GCP service account email", ) return role, f"{BIGQUERY_SERVICE_ACCOUNT_MEMBER_PREFIX}{email}" if member.startswith(BIGQUERY_USER_MEMBER_PREFIX): email = member.removeprefix(BIGQUERY_USER_MEMBER_PREFIX).strip().lower() if not _is_person_email(email): raise ValueError("DAX user: member must be a human work email") return role, f"{BIGQUERY_USER_MEMBER_PREFIX}{email}" raise ValueError("DAX can only patch group:, serviceAccount:, or user: members") def _provider_grant_keys( bindings: tuple[DataAccessDatasetIamBinding, ...], ) -> set[tuple[str, str]]: keys: set[tuple[str, str]] = set() for binding in bindings: for member in binding.members: keys.add(_provider_grant_key(binding.role, member)) return keys def _access_entry_key(entry: dict[str, Any]) -> tuple[str, str] | None: if not _is_unconditional_access_entry(entry): return None role = _access_entry_role(entry) member = _access_entry_member(entry) if not role or not member: return None return role, member def _access_entry_for_provider_grant( role: str, member: str, ) -> dict[str, str]: _, canonical_member = _provider_grant_key(role, member) if canonical_member.startswith(BIGQUERY_GROUP_MEMBER_PREFIX): group_email = canonical_member.removeprefix(BIGQUERY_GROUP_MEMBER_PREFIX) return { "role": "READER", "groupByEmail": group_email, } if canonical_member.startswith(BIGQUERY_SERVICE_ACCOUNT_MEMBER_PREFIX): service_account_email = canonical_member.removeprefix( BIGQUERY_SERVICE_ACCOUNT_MEMBER_PREFIX, ) return { "role": "READER", "userByEmail": service_account_email, } user_email = canonical_member.removeprefix(BIGQUERY_USER_MEMBER_PREFIX) return { "role": "READER", "userByEmail": user_email, } def _is_service_account_email(value: object) -> bool: return bool(re.fullmatch( GCP_SERVICE_ACCOUNT_EMAIL_PATTERN, str(value or "").strip().lower(), )) def _is_person_email(value: object) -> bool: """Accept an ordinary work email but never blur it with a GCP SA.""" email = str(value or "").strip().lower() return ( len(email) <= 255 and re.fullmatch(r"[^@\s]+@[^@\s]+\.[^@\s]+", email) is not None and not _is_service_account_email(email) ) def _access_entry_to_rest(entry: dict[str, Any]) -> dict[str, Any]: if "entity_type" not in entry: return dict(entry) out = {"role": entry.get("role")} entity_type = str(entry.get("entity_type") or "") entity_id = entry.get("entity_id") if entity_type and entity_id is not None: out[entity_type] = entity_id if "condition" in entry: out["condition"] = entry["condition"] return {key: value for key, value in out.items() if value is not None} def _dataset_access_update_body( policy: DataAccessDatasetIamPolicy, policy_patch: DataAccessDatasetIamPolicyPatch, ) -> dict[str, Any]: remove_keys = _provider_grant_keys(policy_patch.remove_bindings) add_keys = _provider_grant_keys(policy_patch.add_bindings) merged_access: list[dict[str, Any]] = [] present_keys: set[tuple[str, str]] = set() for raw_entry in _raw_dataset_access_entries(policy): if not isinstance(raw_entry, dict): continue entry = dict(raw_entry) key = _access_entry_key(entry) if key in remove_keys: continue if key is not None: present_keys.add(key) merged_access.append(_access_entry_to_rest(entry)) # If one desired-state batch contains the same exact grant in both sets, # revocation wins. Re-adding it would make a confirmed revoke a no-op. for role, member in sorted(add_keys - remove_keys - present_keys): merged_access.append(_access_entry_for_provider_grant(role, member)) return { "etag": policy.etag, "access": merged_access, } def _access_arrays_equal(left: Any, right: Any) -> bool: return isinstance(left, list) and isinstance(right, list) and left == right def _policy_access_update_is_noop( policy: DataAccessDatasetIamPolicy, update_body: dict[str, Any], ) -> bool: current_access = [ _access_entry_to_rest(entry) for entry in _raw_dataset_access_entries(policy) if isinstance(entry, dict) ] return _access_arrays_equal(current_access, update_body.get("access")) def _table_iam_update_body( policy: DataAccessDatasetIamPolicy, policy_patch: DataAccessDatasetIamPolicyPatch, ) -> dict[str, Any]: """Merge one exact table policy while preserving unmanaged bindings.""" raw = dict(policy.raw_policy or {}) raw_bindings = raw.get("bindings") if raw_bindings is None: raw_bindings = [] if not isinstance(raw_bindings, list): raise RuntimeError("BigQuery table IAM policy bindings are invalid") remove_keys = _provider_grant_keys(policy_patch.remove_bindings) add_keys = _provider_grant_keys(policy_patch.add_bindings) present_keys: set[tuple[str, str]] = set() conditioned_bindings: list[dict[str, Any]] = [] members_by_role: dict[str, set[str]] = defaultdict(set) for raw_binding in raw_bindings: if not isinstance(raw_binding, dict): continue binding = dict(raw_binding) role = str(binding.get("role") or "") members = [str(member) for member in binding.get("members", [])] if binding.get("condition"): conditioned_bindings.append(binding) continue for member in members: try: key = _provider_grant_key(role, member) except ValueError: members_by_role[role].add(member) continue if key in remove_keys: continue present_keys.add(key) members_by_role[role].add(member) for role, member in add_keys - remove_keys - present_keys: members_by_role[role].add(member) merged = [ {"role": role, "members": sorted(members)} for role, members in sorted(members_by_role.items()) if members ] merged.extend(conditioned_bindings) return { "etag": policy.etag, "version": int(raw.get("version") or 1), "bindings": merged, } def _http_error_status(exc: BaseException) -> int | None: for attr in ("code", "status_code"): status = getattr(exc, attr, None) if status is not None: try: return int(status) except (TypeError, ValueError): pass response = getattr(exc, "response", None) or getattr(exc, "resp", None) status = getattr(response, "status_code", None) if status is None: status = getattr(response, "status", None) if status is None: return None try: return int(status) except (TypeError, ValueError): return None def _is_etag_conflict(exc: BaseException) -> bool: return _http_error_status(exc) in {409, 412} @register_driver class BigQueryDataAccessDriver(DataAccessDriver): provider = "bigquery" display_name = "BigQuery Data Access" description = ( "Dedicated data-access integration for BigQuery dataset and table IAM. " "Used by DAX reconciler paths and kept separate from CMDB cloud " "integrations." ) supports_multi_instance = True supports_test_connection = False config_schema = [ IntegrationFieldInfo( key="service_account_json", label="Service Account JSON", is_secret=True, required=True, is_textarea=True, placeholder='{"type":"service_account", ...}', hint=( "Dedicated data-access service account. Requires only the " "HelixDaxDatasetIamReconciler custom role documented in " "DATA_ACCESS.md. Used only by the data-access reconciler for " "BigQuery dataset/table IAM reads and applies." ), ), IntegrationFieldInfo( key=APPLY_MODE_CONFIG_KEY, label="Enable DAX apply mode", required=False, default="false", options=[ {"label": "Disabled", "value": "false"}, {"label": "Enabled", "value": "true"}, ], hint=( "When enabled, the data-access reconciler may mutate datasets " "permitted by the Terraform-managed BigQuery service account." ), ), IntegrationFieldInfo( key=PERSONAL_APPLY_ENABLED_CONFIG_KEY, label="Enable personal employee DAX apply", required=False, default="false", options=[ {"label": "Disabled", "value": "false"}, {"label": "Enabled", "value": "true"}, ], hint=( "Separate rollout gate for approved personal employee dataset " "read effects. Keep disabled until staging validation passes." ), ), ] binding_schema: list[IntegrationFieldInfo] = [] async def list_dataset_tables( self, project_id: str, dataset_id: str, ) -> tuple[DataAccessDatasetTableSummary, ...]: self._validate_dataset_locator(project_id, dataset_id) try: return await asyncio.to_thread( self._list_dataset_tables_sync, project_id, dataset_id, ) except Exception as exc: raise RuntimeError( "BigQuery table metadata list failed for " f"{project_id}.{dataset_id}: {self._public_error_message(exc)}", ) from exc async def list_dataset_tables_page( self, project_id: str, dataset_id: str, *, page_size: int, page_token: str | None = None, ) -> DataAccessDatasetTablePage: self._validate_dataset_locator(project_id, dataset_id) if page_size < 1 or page_size > 100: raise ValueError("BigQuery table metadata page_size must be between 1 and 100") try: return await asyncio.to_thread( self._list_dataset_tables_page_sync, project_id, dataset_id, page_size, page_token, ) except Exception as exc: raise RuntimeError( "BigQuery table metadata page failed for " f"{project_id}.{dataset_id}: {self._public_error_message(exc)}", ) from exc async def get_dataset_table_schema( self, project_id: str, dataset_id: str, table_id: str, ) -> DataAccessDatasetTableSchema: self._validate_dataset_locator(project_id, dataset_id) normalized_table_id = self._validate_table_id(table_id) try: return await asyncio.to_thread( self._get_dataset_table_schema_sync, project_id, dataset_id, normalized_table_id, ) except Exception as exc: raise RuntimeError( "BigQuery table schema read failed for " f"{project_id}.{dataset_id}.{normalized_table_id}: " f"{self._public_error_message(exc)}", ) from exc async def test_connection(self) -> IntegrationTestResult: return IntegrationTestResult( ok=False, message=( "BigQuery data-access connection tests require a configured " "test dataset or reconciler run and are not available yet." ), detail={"reason": "dax_no_test_dataset_probe_yet"}, ) async def get_dataset_iam( self, project_id: str, dataset_id: str, ) -> DataAccessDatasetIamPolicy: self._validate_dataset_locator(project_id, dataset_id) try: return await asyncio.to_thread( self._get_dataset_iam_sync, project_id, dataset_id, ) except Exception as exc: raise RuntimeError( "BigQuery dataset IAM read failed for " f"{project_id}.{dataset_id}: {self._public_error_message(exc)}", ) from exc async def set_dataset_iam( self, project_id: str, dataset_id: str, expected_etag: str | None, policy_patch: DataAccessDatasetIamPolicyPatch, ) -> DataAccessDatasetIamUpdateResult: self._validate_dataset_locator(project_id, dataset_id) gate_error = self._apply_gate_error(project_id, dataset_id) if gate_error is not None: return DataAccessDatasetIamUpdateResult( ok=False, changed=False, error_message=gate_error, ) if not expected_etag: return DataAccessDatasetIamUpdateResult( ok=False, changed=False, error_message=( "BigQuery DAX apply requires an expected dataset etag" ), ) return await asyncio.to_thread( self._set_dataset_iam_sync, project_id, dataset_id, expected_etag, policy_patch, ) async def get_table_iam( self, project_id: str, dataset_id: str, table_id: str, ) -> DataAccessDatasetIamPolicy: self._validate_dataset_locator(project_id, dataset_id) clean_table_id = self._validate_table_id(table_id) try: return await asyncio.to_thread( self._get_table_iam_sync, project_id, dataset_id, clean_table_id, ) except Exception as exc: raise RuntimeError( "BigQuery table IAM read failed for " f"{project_id}.{dataset_id}.{clean_table_id}: " f"{self._public_error_message(exc)}", ) from exc async def set_table_iam( self, project_id: str, dataset_id: str, table_id: str, expected_etag: str | None, policy_patch: DataAccessDatasetIamPolicyPatch, ) -> DataAccessDatasetIamUpdateResult: self._validate_dataset_locator(project_id, dataset_id) clean_table_id = self._validate_table_id(table_id) gate_error = self._apply_gate_error(project_id, dataset_id) if gate_error is not None: return DataAccessDatasetIamUpdateResult( ok=False, changed=False, error_message=gate_error, ) if not expected_etag: return DataAccessDatasetIamUpdateResult( ok=False, changed=False, error_message="BigQuery DAX apply requires an expected table etag", ) return await asyncio.to_thread( self._set_table_iam_sync, project_id, dataset_id, clean_table_id, expected_etag, policy_patch, ) def _apply_gate_error(self, project_id: str, dataset_id: str) -> str | None: if not data_access_apply_mode_enabled(self.config): return "BigQuery DAX apply mode is disabled" return None def _set_dataset_iam_sync( self, project_id: str, dataset_id: str, expected_etag: str, policy_patch: DataAccessDatasetIamPolicyPatch, ) -> DataAccessDatasetIamUpdateResult: client = self._bigquery_client() try: dataset = self._get_dataset_for_acl_apply(client, project_id, dataset_id) policy = self._dataset_to_iam_policy(dataset, project_id, dataset_id) if policy.etag != expected_etag: return DataAccessDatasetIamUpdateResult( ok=False, changed=False, etag=policy.etag, conflict=True, error_message=( "BigQuery dataset IAM etag changed before apply for " f"{project_id}.{dataset_id}: expected " f"{expected_etag!r}, found {policy.etag!r}" ), policy=policy, ) body = _dataset_access_update_body(policy, policy_patch) if _policy_access_update_is_noop(policy, body): return DataAccessDatasetIamUpdateResult( ok=True, changed=False, etag=policy.etag, already_applied=True, policy=policy, raw_response=policy.raw_policy, ) updated_dataset = self._patch_dataset_acl( client, dataset, body, expected_etag=expected_etag, ) updated_policy = self._dataset_to_iam_policy( updated_dataset, project_id, dataset_id, ) return DataAccessDatasetIamUpdateResult( ok=True, changed=True, etag=updated_policy.etag, policy=updated_policy, raw_response=updated_dataset.to_api_repr(), ) except Exception as exc: if _is_etag_conflict(exc): return DataAccessDatasetIamUpdateResult( ok=False, changed=False, conflict=True, error_message=( "BigQuery dataset IAM etag conflict for " f"{project_id}.{dataset_id}: " f"{self._public_error_message(exc)}" ), ) return DataAccessDatasetIamUpdateResult( ok=False, changed=False, error_message=( "BigQuery dataset IAM apply failed for " f"{project_id}.{dataset_id}: " f"{self._public_error_message(exc)}" ), ) finally: close = getattr(client, "close", None) if callable(close): close() def _get_dataset_iam_sync( self, project_id: str, dataset_id: str, ) -> DataAccessDatasetIamPolicy: from google.cloud import bigquery # type: ignore from google.cloud.bigquery.enums import DatasetView # type: ignore client = self._bigquery_client() try: dataset_ref = bigquery.DatasetReference(project_id, dataset_id) dataset = client.get_dataset(dataset_ref, dataset_view=DatasetView.ACL) return self._dataset_to_iam_policy(dataset, project_id, dataset_id) finally: close = getattr(client, "close", None) if callable(close): close() def _get_table_iam_sync( self, project_id: str, dataset_id: str, table_id: str, ) -> DataAccessDatasetIamPolicy: from google.cloud import bigquery # type: ignore client = self._bigquery_client() try: table_ref = bigquery.TableReference( bigquery.DatasetReference(project_id, dataset_id), table_id, ) policy = client.get_iam_policy( table_ref, requested_policy_version=_TABLE_IAM_POLICY_VERSION, ) return self._table_to_iam_policy( policy, project_id, dataset_id, table_id, ) finally: close = getattr(client, "close", None) if callable(close): close() def _set_table_iam_sync( self, project_id: str, dataset_id: str, table_id: str, expected_etag: str, policy_patch: DataAccessDatasetIamPolicyPatch, ) -> DataAccessDatasetIamUpdateResult: from google.api_core.iam import Policy # type: ignore from google.cloud import bigquery # type: ignore client = self._bigquery_client() table_ref = bigquery.TableReference( bigquery.DatasetReference(project_id, dataset_id), table_id, ) try: current_provider_policy = client.get_iam_policy( table_ref, requested_policy_version=_TABLE_IAM_POLICY_VERSION, ) current = self._table_to_iam_policy( current_provider_policy, project_id, dataset_id, table_id, ) if current.etag != expected_etag: return DataAccessDatasetIamUpdateResult( ok=False, changed=False, etag=current.etag, conflict=True, policy=current, error_message=( "BigQuery table IAM etag changed before apply for " f"{project_id}.{dataset_id}.{table_id}" ), ) body = _table_iam_update_body(current, policy_patch) if body == current.raw_policy: return DataAccessDatasetIamUpdateResult( ok=True, changed=False, etag=current.etag, already_applied=True, policy=current, raw_response=current.raw_policy, ) updated_provider_policy = client.set_iam_policy( table_ref, Policy.from_api_repr(body), fields=["bindings", "etag", "version"], ) updated = self._table_to_iam_policy( updated_provider_policy, project_id, dataset_id, table_id, ) return DataAccessDatasetIamUpdateResult( ok=True, changed=True, etag=updated.etag, policy=updated, raw_response=updated.raw_policy, ) except Exception as exc: if _is_etag_conflict(exc): return DataAccessDatasetIamUpdateResult( ok=False, changed=False, conflict=True, error_message=( "BigQuery table IAM etag conflict for " f"{project_id}.{dataset_id}.{table_id}: " f"{self._public_error_message(exc)}" ), ) return DataAccessDatasetIamUpdateResult( ok=False, changed=False, error_message=( "BigQuery table IAM apply failed for " f"{project_id}.{dataset_id}.{table_id}: " f"{self._public_error_message(exc)}" ), ) finally: close = getattr(client, "close", None) if callable(close): close() def _list_dataset_tables_sync( self, project_id: str, dataset_id: str, ) -> tuple[DataAccessDatasetTableSummary, ...]: from google.cloud import bigquery # type: ignore client = self._bigquery_client() try: dataset_ref = bigquery.DatasetReference(project_id, dataset_id) tables = ( DataAccessDatasetTableSummary( table_id=str(table.table_id), table_type=str(getattr(table, "table_type", "TABLE") or "TABLE"), ) for table in client.list_tables(dataset_ref) ) return tuple(sorted(tables, key=lambda table: table.table_id.casefold())) finally: close = getattr(client, "close", None) if callable(close): close() def _list_dataset_tables_page_sync( self, project_id: str, dataset_id: str, page_size: int, page_token: str | None, ) -> DataAccessDatasetTablePage: from google.cloud import bigquery # type: ignore client = self._bigquery_client() try: dataset_ref = bigquery.DatasetReference(project_id, dataset_id) iterator = client.list_tables( dataset_ref, max_results=page_size, page_token=page_token, ) page = next(iterator.pages, None) items = tuple( DataAccessDatasetTableSummary( table_id=str(table.table_id), table_type=str(getattr(table, "table_type", "TABLE") or "TABLE"), ) for table in (page or ()) ) raw_page = getattr(page, "raw_page", {}) if page is not None else {} total_rows = raw_page.get("totalItems") return DataAccessDatasetTablePage( items=items, next_page_token=getattr(iterator, "next_page_token", None), total_rows=int(total_rows) if total_rows is not None else None, ) finally: close = getattr(client, "close", None) if callable(close): close() def _get_dataset_table_schema_sync( self, project_id: str, dataset_id: str, table_id: str, ) -> DataAccessDatasetTableSchema: from google.cloud import bigquery # type: ignore client = self._bigquery_client() try: dataset_ref = bigquery.DatasetReference(project_id, dataset_id) table_ref = bigquery.TableReference(dataset_ref, table_id) table = client.get_table(table_ref) return DataAccessDatasetTableSchema( project_id=project_id, dataset_id=dataset_id, table_id=str(getattr(table, "table_id", table_id) or table_id), table_type=str(getattr(table, "table_type", "TABLE") or "TABLE"), description=self._optional_text(getattr(table, "description", None)), fields=self._flatten_schema_fields( tuple(getattr(table, "schema", ()) or ()), ), ) finally: close = getattr(client, "close", None) if callable(close): close() def _bigquery_client(self) -> Any: from google.cloud import bigquery # type: ignore info = self._service_account_info() credentials = self._credentials_from_info(info) return bigquery.Client(credentials=credentials, project=info["project_id"]) def _credentials_from_info(self, info: dict[str, Any]) -> Any: from google.oauth2 import service_account # type: ignore return service_account.Credentials.from_service_account_info( info, scopes=[_BQ_SCOPE], ) def _get_dataset_for_acl_apply( self, client: Any, project_id: str, dataset_id: str, ) -> Any: from google.cloud import bigquery # type: ignore from google.cloud.bigquery.client import DEFAULT_RETRY # type: ignore from google.cloud.bigquery.enums import DatasetView # type: ignore dataset_ref = bigquery.DatasetReference(project_id, dataset_id) response = client._call_api( DEFAULT_RETRY, span_name="BigQuery.getDataset", span_attributes={"path": dataset_ref.path}, method="GET", path=dataset_ref.path, query_params={ "datasetView": DatasetView.ACL.value, "accessPolicyVersion": _ACCESS_POLICY_VERSION_WITH_CONDITIONS, }, ) return bigquery.Dataset.from_api_repr(response) def _patch_dataset_acl( self, client: Any, dataset: Any, body: dict[str, Any], *, expected_etag: str, ) -> Any: from google.cloud import bigquery # type: ignore from google.cloud.bigquery.client import DEFAULT_RETRY # type: ignore dataset_ref = bigquery.DatasetReference( dataset.project, dataset.dataset_id, ) response = client._call_api( DEFAULT_RETRY, span_name="BigQuery.updateDataset", span_attributes={"path": dataset_ref.path, "fields": ("access",)}, method="PATCH", path=dataset_ref.path, data=body, headers={"If-Match": expected_etag}, query_params={ "updateMode": _DATASET_UPDATE_MODE_ACL, "accessPolicyVersion": _ACCESS_POLICY_VERSION_WITH_CONDITIONS, }, ) return bigquery.Dataset.from_api_repr(response) def _service_account_info(self) -> dict[str, Any]: raw = str(self.config.get("service_account_json") or "").strip() if not raw: raise RuntimeError("service_account_json not configured") try: info = json.loads(raw) except ValueError as exc: raise ValueError("service_account_json is not valid JSON") from exc if not isinstance(info, dict): raise ValueError("service_account_json must decode to an object") missing = [key for key in _REQUIRED_SERVICE_ACCOUNT_FIELDS if not info.get(key)] if missing: raise ValueError( "service_account_json is missing fields: " + ", ".join(missing), ) return info @classmethod def _dataset_to_iam_policy( cls, dataset: Any, project_id: str, dataset_id: str, ) -> DataAccessDatasetIamPolicy: grouped: dict[str, set[str]] = defaultdict(set) raw_entries: list[dict[str, Any]] = [] for entry in tuple(getattr(dataset, "access_entries", ()) or ()): raw_entry = cls._access_entry_api_repr(entry) raw_entries.append(raw_entry) if cls._access_entry_has_condition(entry, raw_entry): continue role = cls._provider_role_for_access_entry(entry) member = cls._provider_member_for_access_entry(entry) if role and member: grouped[role].add(member) bindings = tuple( DataAccessDatasetIamBinding( role=role, members=tuple(sorted(members)), ) for role, members in sorted(grouped.items()) ) return DataAccessDatasetIamPolicy( project_id=project_id, dataset_id=dataset_id, etag=cls._dataset_etag(dataset), bindings=bindings, raw_policy={"access_entries": raw_entries}, ) @staticmethod def _table_to_iam_policy( policy: Any, project_id: str, dataset_id: str, table_id: str, ) -> DataAccessDatasetIamPolicy: raw = policy.to_api_repr() bindings = tuple( DataAccessDatasetIamBinding( role=str(binding.get("role") or ""), members=tuple(sorted(str(member) for member in binding.get("members", []))), ) for binding in raw.get("bindings", []) if isinstance(binding, dict) and not binding.get("condition") ) return DataAccessDatasetIamPolicy( project_id=project_id, dataset_id=dataset_id, table_id=table_id, etag=str(getattr(policy, "etag", "") or "") or None, bindings=bindings, raw_policy=raw, ) @staticmethod def _access_entry_api_repr(entry: Any) -> dict[str, Any]: to_api_repr = getattr(entry, "to_api_repr", None) if callable(to_api_repr): raw = to_api_repr() return dict(raw) if isinstance(raw, dict) else {"value": raw} out: dict[str, Any] = {} for attr in ("role", "entity_type", "entity_id"): value = getattr(entry, attr, None) if value is not None: out[attr] = value return out @staticmethod def _access_entry_has_condition(entry: Any, raw_entry: dict[str, Any]) -> bool: if raw_entry.get("condition"): return True return getattr(entry, "condition", None) is not None @staticmethod def _provider_role_for_access_entry(entry: Any) -> str | None: role = str(getattr(entry, "role", "") or "").strip() if not role: return None return _DATASET_ACCESS_ROLE_TO_PROVIDER_ROLE.get(role.upper(), role) @staticmethod def _provider_member_for_access_entry(entry: Any) -> str | None: entity_type = str(getattr(entry, "entity_type", "") or "").strip() entity_id = str(getattr(entry, "entity_id", "") or "").strip() if not entity_id: return None if entity_type == "groupByEmail": return f"group:{entity_id}" if entity_type == "userByEmail": if _is_service_account_email(entity_id): return f"serviceAccount:{entity_id.lower()}" return f"user:{entity_id}" if entity_type == "domain": return f"domain:{entity_id}" if entity_type == "specialGroup": return f"specialGroup:{entity_id}" if entity_type == "iamMember": return entity_id return None @staticmethod def _dataset_etag(dataset: Any) -> str | None: etag = getattr(dataset, "etag", None) if etag: return str(etag) props = getattr(dataset, "_properties", None) if isinstance(props, dict) and props.get("etag"): return str(props["etag"]) return None @staticmethod def _validate_dataset_locator(project_id: str, dataset_id: str) -> None: _validate_dataset_locator(project_id, dataset_id) @staticmethod def _validate_table_id(table_id: str) -> str: normalized = str(table_id or "").strip() if not normalized or len(normalized) > 1024: raise ValueError("table_id must be between 1 and 1024 characters") if "/" in normalized: raise ValueError("table_id must stay within one dataset") return normalized @classmethod def _flatten_schema_fields( cls, fields: tuple[Any, ...], *, parent_path: str = "", depth: int = 0, ) -> tuple[DataAccessDatasetFieldSchema, ...]: flattened: list[DataAccessDatasetFieldSchema] = [] for ordinal_position, field in enumerate(fields, start=1): name = str(getattr(field, "name", "") or "") field_path = f"{parent_path}.{name}" if parent_path else name flattened.append(DataAccessDatasetFieldSchema( name=name, field_path=field_path, field_type=str(getattr(field, "field_type", "") or "UNKNOWN"), mode=str(getattr(field, "mode", "") or "NULLABLE"), description=cls._optional_text(getattr(field, "description", None)), depth=depth, ordinal_position=ordinal_position, )) nested = tuple(getattr(field, "fields", ()) or ()) if nested: flattened.extend(cls._flatten_schema_fields( nested, parent_path=field_path, depth=depth + 1, )) return tuple(flattened) @staticmethod def _optional_text(value: Any) -> str | None: text = str(value or "").strip() return text or None @staticmethod def _public_error_message(exc: BaseException) -> str: msg = _PRIVATE_KEY_RE.sub("", str(exc)) return msg[:500]