Skip to content

feat(bigtable): forward auth/identity config to accelerator daemon #3

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
mutianf wants to merge 1 commit into from
Open

feat(bigtable): forward auth/identity config to accelerator daemon #3

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
11 changes: 10 additions & 1 deletion packages/google-cloud-bigtable/google/cloud/bigtable/data/_accelerator/_daemon.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -31,7 +31,7 @@
import tempfile
import time
import typing
from typing import Sequence
from typing import Mapping, Sequence

# Environment variable that overrides the bundled binary location. Primarily
# for development against a locally-built daemon, and for tests pointing at a
Expand Down Expand Up @@ -97,10 +97,15 @@ def __init__(
*,
binary_path: str | None = None,
startup_timeout: float = _DEFAULT_STARTUP_TIMEOUT,
extra_env: Mapping[str, str] | None = None,
):
self._binary_path = binary_path or _resolve_binary_path()
self._cli_flags = list(cli_flags)
self._startup_timeout = startup_timeout
# Extra environment for the subprocess, merged over the inherited env.
# Used to forward GOOGLE_APPLICATION_CREDENTIALS (path only) so the
# daemon's ADC resolves the caller's credentials_file.
self._extra_env = dict(extra_env) if extra_env else {}
self._tempdir: str | None = None
self._uds_path: str | None = None
self._log_path: str | None = None
Expand Down Expand Up @@ -146,13 +151,17 @@ def start(self) -> None:
self._log_path = os.path.join(self._tempdir, "daemon.log")
self._log_file = open(self._log_path, "wb")
argv = [self._binary_path, "--uds-path", self._uds_path, *self._cli_flags]
env = None
if self._extra_env:
env = {**os.environ, **self._extra_env}
Comment thread
mutianf marked this conversation as resolved.
try:
self._proc = subprocess.Popen(
argv,
stdin=subprocess.PIPE,
stdout=self._log_file,
stderr=subprocess.STDOUT,
close_fds=True,
env=env,
Comment thread
mutianf marked this conversation as resolved.
)
except OSError as exc:
self._close_log_file()
Expand Down
110 changes: 101 additions & 9 deletions packages/google-cloud-bigtable/google/cloud/bigtable/data/_async/client.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -50,6 +50,13 @@
from grpc import Channel

from google.cloud.bigtable.client import _DEFAULT_BIGTABLE_EMULATOR_CLIENT
from google.cloud.bigtable.data._accelerator._daemon import AcceleratorDaemon
from google.cloud.bigtable.data._accelerator._fallback import (
AcceleratorBreaker,
_AcceleratorFallback,
handle_accelerator_error,
)
from google.cloud.bigtable.data._accelerator._routing import is_supported
from google.cloud.bigtable.data._cross_sync import CrossSync
from google.cloud.bigtable.data._helpers import (
_CONCURRENCY_LIMIT,
Expand Down Expand Up @@ -104,14 +111,6 @@
SampleRowKeysRequest,
)

from google.cloud.bigtable.data._accelerator._daemon import AcceleratorDaemon
from google.cloud.bigtable.data._accelerator._fallback import (
AcceleratorBreaker,
_AcceleratorFallback,
handle_accelerator_error,
)
from google.cloud.bigtable.data._accelerator._routing import is_supported

if CrossSync.is_async:
from grpc.aio import insecure_channel

Expand Down Expand Up @@ -243,6 +242,10 @@ def __init__(
client_options = cast(
Optional[client_options_lib.ClientOptions], client_options
)
# Compute the accelerator forward-config now, while both credentials and
# client_options are still in scope and before the emulator block below
# can reassign credentials. Reads only client_options fields.
self._init_accelerator_config(credentials is not None, client_options)
self._emulator_host = os.getenv(BIGTABLE_EMULATOR)
if self._emulator_host is not None:
warnings.warn(
Expand Down Expand Up @@ -310,6 +313,80 @@ def __init__(
stacklevel=2,
)

def _init_accelerator_config(
self,
explicit_credentials: bool,
client_options: "client_options_lib.ClientOptions | None",
) -> None:
"""Compute the accelerator eligibility and forward-config from the
caller's auth/identity configuration.

Sets three attributes read later by ``_DataApiTarget``:

* ``_accelerator_blocked_reason``: non-None means an in-memory secret
was supplied that the daemon cannot reproduce (``credentials=``,
``api_key``, ``client_cert_source``) -> fall back to the native
client with a warning.
* ``_accelerator_flags``: extra daemon CLI flags for the non-secret,
reproducible knobs (scopes, quota project, endpoint, universe).
* ``_accelerator_env``: extra subprocess env; ``credentials_file`` is
forwarded here as ``GOOGLE_APPLICATION_CREDENTIALS`` so only the path
crosses, never the key bytes.
"""
self._accelerator_blocked_reason: str | None = None
self._accelerator_flags: list[str] = []
self._accelerator_env: dict[str, str] = {}

# In-memory secrets that cannot be forwarded to a separate process.
if explicit_credentials:
self._accelerator_blocked_reason = (
"an explicit credentials= object cannot be forwarded to the "
"accelerator daemon"
)
elif client_options is not None and getattr(client_options, "api_key", None):
self._accelerator_blocked_reason = (
"api_key cannot be forwarded to the accelerator daemon"
)
elif client_options is not None and getattr(
client_options, "client_cert_source", None
):
self._accelerator_blocked_reason = (
"client_cert_source (mTLS) cannot be forwarded to the "
"accelerator daemon"
)

# Non-secret / path-style knobs the daemon can reproduce.
scopes = getattr(client_options, "scopes", None) if client_options else None
# Forward the effective auth scopes so the daemon's token audience
# matches the native client (whose default is TransportType.AUTH_SCOPES),
# rather than the daemon's single hardcoded data scope.
effective_scopes = list(scopes) if scopes else list(TransportType.AUTH_SCOPES)
if effective_scopes:
self._accelerator_flags += ["--scopes", ",".join(effective_scopes)]

if client_options is not None:
quota_project_id = getattr(client_options, "quota_project_id", None)
if quota_project_id:
self._accelerator_flags += ["--quota-project", quota_project_id]

api_endpoint = getattr(client_options, "api_endpoint", None)
if api_endpoint:
# Normalize to host[:port]; daemon takes a bare endpoint.
normalized = api_endpoint.split("://", 1)[-1]
self._accelerator_flags += ["--data-endpoint", normalized]

universe_domain = getattr(client_options, "universe_domain", None)
if universe_domain and universe_domain != "googleapis.com":
self._accelerator_flags += ["--universe-domain", universe_domain]

credentials_file = getattr(client_options, "credentials_file", None)
if credentials_file:
# Forward the path only; the daemon's ADC + dial honor
# GOOGLE_APPLICATION_CREDENTIALS automatically.
self._accelerator_env["GOOGLE_APPLICATION_CREDENTIALS"] = (
credentials_file
)

def _build_grpc_channel(self, *args, **kwargs) -> SwappableChannelType:
"""
This method is called by the gapic transport to create a grpc channel.
Expand Down Expand Up @@ -1187,6 +1264,18 @@ def _maybe_start_accelerator(self, *, explicit: bool) -> None:
stacklevel=2,
)
return
if self.client._accelerator_blocked_reason is not None:
Comment thread
mutianf marked this conversation as resolved.
# An in-memory secret we cannot reproduce in the daemon was
# supplied; stay on the native client rather than run the
# accelerated path under a different identity.
warnings.warn(
"Accelerator disabled: "
f"{self.client._accelerator_blocked_reason}; using the "
"native client instead.",
RuntimeWarning,
stacklevel=2,
)
return
try:
self._start_accelerator()
except Exception as exc:
Expand All @@ -1212,7 +1301,10 @@ def _start_accelerator(self) -> None:
]
if self.app_profile_id:
flags.extend(["--app-profile", self.app_profile_id])
server = AcceleratorDaemon(cli_flags=flags)
flags.extend(self.client._accelerator_flags)
server = AcceleratorDaemon(
cli_flags=flags, extra_env=self.client._accelerator_env
)
try:
server.start()
self._accelerator_client = AcceleratorClientType(server.uds_path)
Expand Down
66 changes: 65 additions & 1 deletion packages/google-cloud-bigtable/google/cloud/bigtable/data/_sync_autogen/client.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -178,6 +178,7 @@ def __init__(
client_options = cast(
Optional[client_options_lib.ClientOptions], client_options
)
self._init_accelerator_config(credentials is not None, client_options)
self._emulator_host = os.getenv(BIGTABLE_EMULATOR)
if self._emulator_host is not None:
warnings.warn(
Expand Down Expand Up @@ -234,6 +235,59 @@ def __init__(
stacklevel=2,
)

def _init_accelerator_config(
self,
explicit_credentials: bool,
client_options: "client_options_lib.ClientOptions | None",
) -> None:
"""Compute the accelerator eligibility and forward-config from the
caller's auth/identity configuration.

Sets three attributes read later by ``_DataApiTarget``:

* ``_accelerator_blocked_reason``: non-None means an in-memory secret
was supplied that the daemon cannot reproduce (``credentials=``,
``api_key``, ``client_cert_source``) -> fall back to the native
client with a warning.
* ``_accelerator_flags``: extra daemon CLI flags for the non-secret,
reproducible knobs (scopes, quota project, endpoint, universe).
* ``_accelerator_env``: extra subprocess env; ``credentials_file`` is
forwarded here as ``GOOGLE_APPLICATION_CREDENTIALS`` so only the path
crosses, never the key bytes."""
self._accelerator_blocked_reason: str | None = None
self._accelerator_flags: list[str] = []
self._accelerator_env: dict[str, str] = {}
if explicit_credentials:
self._accelerator_blocked_reason = "an explicit credentials= object cannot be forwarded to the accelerator daemon"
elif client_options is not None and getattr(client_options, "api_key", None):
self._accelerator_blocked_reason = (
"api_key cannot be forwarded to the accelerator daemon"
)
elif client_options is not None and getattr(
client_options, "client_cert_source", None
):
self._accelerator_blocked_reason = "client_cert_source (mTLS) cannot be forwarded to the accelerator daemon"
scopes = getattr(client_options, "scopes", None) if client_options else None
effective_scopes = list(scopes) if scopes else list(TransportType.AUTH_SCOPES)
if effective_scopes:
self._accelerator_flags += ["--scopes", ",".join(effective_scopes)]
if client_options is not None:
quota_project_id = getattr(client_options, "quota_project_id", None)
if quota_project_id:
self._accelerator_flags += ["--quota-project", quota_project_id]
api_endpoint = getattr(client_options, "api_endpoint", None)
if api_endpoint:
normalized = api_endpoint.split("://", 1)[-1]
self._accelerator_flags += ["--data-endpoint", normalized]
universe_domain = getattr(client_options, "universe_domain", None)
if universe_domain and universe_domain != "googleapis.com":
self._accelerator_flags += ["--universe-domain", universe_domain]
credentials_file = getattr(client_options, "credentials_file", None)
if credentials_file:
self._accelerator_env["GOOGLE_APPLICATION_CREDENTIALS"] = (
credentials_file
)

def _build_grpc_channel(self, *args, **kwargs) -> SwappableChannelType:
"""This method is called by the gapic transport to create a grpc channel.

Expand Down Expand Up @@ -934,6 +988,13 @@ def _maybe_start_accelerator(self, *, explicit: bool) -> None:
stacklevel=2,
)
return
if self.client._accelerator_blocked_reason is not None:
warnings.warn(
f"Accelerator disabled: {self.client._accelerator_blocked_reason}; using the native client instead.",
RuntimeWarning,
stacklevel=2,
)
return
try:
self._start_accelerator()
except Exception as exc:
Expand All @@ -952,7 +1013,10 @@ def _start_accelerator(self) -> None:
flags = ["--project", self.client.project, "--instance", self.instance_id]
if self.app_profile_id:
flags.extend(["--app-profile", self.app_profile_id])
server = AcceleratorDaemon(cli_flags=flags)
flags.extend(self.client._accelerator_flags)
server = AcceleratorDaemon(
cli_flags=flags, extra_env=self.client._accelerator_env
)
try:
server.start()
self._accelerator_client = AcceleratorClientType(server.uds_path)
Expand Down
95 changes: 95 additions & 0 deletions packages/google-cloud-bigtable/tests/unit/data/test_accelerator_credentials.py
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -0,0 +1,95 @@
# Copyright 2026 Google LLC
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
"""Unit tests for the accelerator credential-safety guardrail: which auth
configuration the client forwards to the daemon vs. falls back to native for.
These exercise pure helpers on the async client without spinning up a full
client or event loop."""

from google.api_core.client_options import ClientOptions

from google.cloud.bigtable.data._async.client import BigtableDataClientAsync


def _bare_client():
"""A client instance that skips __init__ (no network / event loop)."""
return object.__new__(BigtableDataClientAsync)


class TestInitAcceleratorConfig:
def test_no_config_forwards_default_scopes(self):
c = _bare_client()
c._init_accelerator_config(explicit_credentials=False, client_options=None)
assert c._accelerator_blocked_reason is None
assert "--scopes" in c._accelerator_flags
assert c._accelerator_env == {}

def test_explicit_credentials_blocks(self):
c = _bare_client()
c._init_accelerator_config(explicit_credentials=True, client_options=None)
assert c._accelerator_blocked_reason is not None
assert "credentials" in c._accelerator_blocked_reason

def test_api_key_blocks(self):
c = _bare_client()
c._init_accelerator_config(
explicit_credentials=False, client_options=ClientOptions(api_key="AIzaKEY")
)
assert "api_key" in c._accelerator_blocked_reason

def test_client_cert_source_blocks(self):
c = _bare_client()
c._init_accelerator_config(
explicit_credentials=False,
client_options=ClientOptions(client_cert_source=lambda: (b"", b"")),
)
assert "client_cert_source" in c._accelerator_blocked_reason

def test_forwards_non_secret_knobs(self):
c = _bare_client()
opts = ClientOptions(
scopes=["https://www.googleapis.com/auth/custom"],
quota_project_id="quota-proj",
api_endpoint="https://custom.example.com:443",
universe_domain="my-universe.example.com",
)
c._init_accelerator_config(explicit_credentials=False, client_options=opts)
flags = c._accelerator_flags
assert flags[flags.index("--scopes") + 1] == (
"https://www.googleapis.com/auth/custom"
)
assert flags[flags.index("--quota-project") + 1] == "quota-proj"
# api_endpoint is normalized to host:port (scheme stripped).
assert flags[flags.index("--data-endpoint") + 1] == "custom.example.com:443"
assert flags[flags.index("--universe-domain") + 1] == "my-universe.example.com"
assert c._accelerator_blocked_reason is None

def test_default_universe_not_forwarded(self):
c = _bare_client()
c._init_accelerator_config(
explicit_credentials=False,
client_options=ClientOptions(universe_domain="googleapis.com"),
)
assert "--universe-domain" not in c._accelerator_flags

def test_credentials_file_forwarded_as_env_path(self):
c = _bare_client()
c._init_accelerator_config(
explicit_credentials=False,
client_options=ClientOptions(credentials_file="/path/to/key.json"),
)
assert c._accelerator_env["GOOGLE_APPLICATION_CREDENTIALS"] == (
"/path/to/key.json"
)
# The path is not a "secret" that blocks acceleration.
assert c._accelerator_blocked_reason is None
Loading
Loading