This is an automated email from the ASF dual-hosted git repository.
potiuk pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/airflow.git
The following commit(s) were added to refs/heads/main by this push:
new d1d5b1af26e Adapt GCP BigtableHook to run in private cloud (#73181)
d1d5b1af26e is described below
commit d1d5b1af26e9308034a8c3116e7ef865d3381e5a
Author: olegkachur-e <[email protected]>
AuthorDate: Sun Oct 4 16:35:22 2026 +0000
Adapt GCP BigtableHook to run in private cloud (#73181)
* Adapt GCP BigtableHook to run in private cloud
- Customize ClientOptions generation,
for the non-default universal domain.
* Explain the Bigtable admin endpoint override and log the ignored override
The override exists because google.cloud.bigtable.Client ignores
universe_domain when it builds the admin channel; without saying so it is easy
to simplify away. Dropping api_endpoint_override silently also differed from
the base hook, which logs it.
Generated-by: Claude Opus 5
---------
Co-authored-by: Oleg Kachur <[email protected]>
Co-authored-by: Jarek Potiuk <[email protected]>
---
.../providers/google/cloud/hooks/bigtable.py | 24 ++++++++++++++++++++
.../tests/unit/google/cloud/hooks/test_bigtable.py | 26 ++++++++++++++++++++++
2 files changed, 50 insertions(+)
diff --git
a/providers/google/src/airflow/providers/google/cloud/hooks/bigtable.py
b/providers/google/src/airflow/providers/google/cloud/hooks/bigtable.py
index a336381ee24..0e79bdbce27 100644
--- a/providers/google/src/airflow/providers/google/cloud/hooks/bigtable.py
+++ b/providers/google/src/airflow/providers/google/cloud/hooks/bigtable.py
@@ -19,10 +19,12 @@
from __future__ import annotations
+import os
from collections.abc import Sequence
from typing import TYPE_CHECKING
import google.api_core.exceptions
+from google.api_core.client_options import ClientOptions
from google.cloud.bigtable import Client, enums
from google.cloud.bigtable.cluster import Cluster
from google.cloud.bigtable.instance import Instance
@@ -59,6 +61,28 @@ class BigtableHook(GoogleBaseHook):
)
self._client: Client | None = None
+ def get_client_options(
+ self,
+ api_endpoint_override: str | None = None,
+ ) -> ClientOptions:
+ """
+ Return the ClientOptions object for Google Bigtable Admin API.
+
+ The returned options point at the Bigtable Admin API. ``Client``
passes the same options to its
+ data client, so data-plane methods added to this hook would need their
own endpoint.
+ """
+ if not self.is_default_universe():
+ if api_endpoint_override:
+ self.log.info(
+ "Ignoring api_endpoint_override because the universe
domain is not Google default universe."
+ )
+ global_universe_domain = os.getenv("GOOGLE_CLOUD_UNIVERSE_DOMAIN")
+ # google.cloud.bigtable.Client builds the admin channel from
api_endpoint only and ignores
+ # universe_domain, so the base hook's
ClientOptions(universe_domain=...) would still reach
+ # bigtableadmin.googleapis.com.
+ return
ClientOptions(api_endpoint=f"bigtableadmin.{global_universe_domain}")
+ return
super().get_client_options(api_endpoint_override=api_endpoint_override)
+
def _get_client(self, project_id: str) -> Client:
if not self._client:
self._client = Client(
diff --git a/providers/google/tests/unit/google/cloud/hooks/test_bigtable.py
b/providers/google/tests/unit/google/cloud/hooks/test_bigtable.py
index b5c5420bf2c..b30d82cae27 100644
--- a/providers/google/tests/unit/google/cloud/hooks/test_bigtable.py
+++ b/providers/google/tests/unit/google/cloud/hooks/test_bigtable.py
@@ -59,6 +59,32 @@ class TestBigtableHookNoDefaultProjectId:
):
self.bigtable_hook_no_default_project_id =
BigtableHook(gcp_conn_id="test")
+ @pytest.mark.parametrize("api_endpoint_override", [None,
"custom-override-api_endpoint"])
+ def test_get_client_options_non_default_universe(self, monkeypatch,
api_endpoint_override):
+ monkeypatch.setenv("GOOGLE_CLOUD_UNIVERSE_DOMAIN",
"custom-universe.com")
+ with mock.patch.object(BigtableHook, "log", new_callable=PropertyMock)
as mock_log_property:
+ result =
self.bigtable_hook_no_default_project_id.get_client_options(
+ api_endpoint_override=api_endpoint_override
+ )
+
+ assert result.api_endpoint == "bigtableadmin.custom-universe.com"
+ assert mock_log_property.return_value.info.called is
(api_endpoint_override is not None)
+
+
@mock.patch("airflow.providers.google.common.hooks.base_google.GoogleBaseHook.get_client_options")
+
@mock.patch("airflow.providers.google.cloud.hooks.bigtable.BigtableHook.is_default_universe")
+ def test_get_client_options_default_universe(
+ self, mock_is_default_universe, mock_base_get_client_options
+ ):
+ mock_is_default_universe.return_value = True
+ test_api_endpoint = "custom-override-api_endpoint"
+ result = self.bigtable_hook_no_default_project_id.get_client_options(
+ api_endpoint_override=test_api_endpoint
+ )
+
+ mock_is_default_universe.assert_called_once()
+
mock_base_get_client_options.assert_called_once_with(api_endpoint_override=test_api_endpoint)
+ assert result == mock_base_get_client_options.return_value
+
@mock.patch("airflow.providers.google.cloud.hooks.bigtable.BigtableHook.get_client_options")
@mock.patch("airflow.providers.google.cloud.hooks.bigtable.BigtableHook.get_credentials")
@mock.patch("airflow.providers.google.cloud.hooks.bigtable.Client")
