From a95bc51699760d0d283ae4993540aa47ccfc0d74 Mon Sep 17 00:00:00 2001 From: Mattie Fu Date: Tue, 28 Jul 2026 21:13:56 +0000 Subject: [PATCH 1/5] feat(bigtable): pass Python client version to accelerator daemon for User-Agent prefixing Pass --caller-user-agent python-v when spawning the daemon so it can construct the full User-Agent prefix (e.g. python-v3.0.0a0-go-accelerator-v3.x.x) for accurate audit logging and traffic attribution. Change-Id: I621217335ad54377a84d44ab4b9af88770dd2958 --- .../google/cloud/bigtable/data/_async/client.py | 8 ++++++++ .../cloud/bigtable/data/_sync_autogen/client.py | 4 ++++ .../unit/data/test_accelerator_credentials.py | 15 +++++++++++++++ 3 files changed, 27 insertions(+) diff --git a/packages/google-cloud-bigtable/google/cloud/bigtable/data/_async/client.py b/packages/google-cloud-bigtable/google/cloud/bigtable/data/_async/client.py index 48c06e91196a..be86ecdd9f32 100644 --- a/packages/google-cloud-bigtable/google/cloud/bigtable/data/_async/client.py +++ b/packages/google-cloud-bigtable/google/cloud/bigtable/data/_async/client.py @@ -457,6 +457,14 @@ def _init_accelerator_config( credentials_file ) + # Identify this client to the daemon so it can construct the full + # User-Agent prefix (e.g. "python-v3.0.0a0-go-accelerator-v3.x.x"). + # The daemon appends its own "-go-accelerator-vX.Y.Z" suffix. + self._accelerator_flags += [ + "--caller-user-agent", + f"python-v{google.cloud.bigtable.__version__}", + ] + def _resolve_principal(self) -> str | None: """Resolve this client's identity to a principal. diff --git a/packages/google-cloud-bigtable/google/cloud/bigtable/data/_sync_autogen/client.py b/packages/google-cloud-bigtable/google/cloud/bigtable/data/_sync_autogen/client.py index 517052189031..ffe42eadc505 100644 --- a/packages/google-cloud-bigtable/google/cloud/bigtable/data/_sync_autogen/client.py +++ b/packages/google-cloud-bigtable/google/cloud/bigtable/data/_sync_autogen/client.py @@ -340,6 +340,10 @@ def _init_accelerator_config( self._accelerator_env["GOOGLE_APPLICATION_CREDENTIALS"] = ( credentials_file ) + self._accelerator_flags += [ + "--caller-user-agent", + f"python-v{google.cloud.bigtable.__version__}", + ] def _resolve_principal(self) -> str | None: """Resolve this client's identity to a principal. diff --git a/packages/google-cloud-bigtable/tests/unit/data/test_accelerator_credentials.py b/packages/google-cloud-bigtable/tests/unit/data/test_accelerator_credentials.py index f3ff885106f9..02306f47d840 100644 --- a/packages/google-cloud-bigtable/tests/unit/data/test_accelerator_credentials.py +++ b/packages/google-cloud-bigtable/tests/unit/data/test_accelerator_credentials.py @@ -22,6 +22,7 @@ from google.api_core.client_options import ClientOptions from google.auth import compute_engine +import google.cloud.bigtable from google.cloud.bigtable.data._async.client import ( BigtableDataClientAsync, _AcceleratorIdentityError, @@ -113,6 +114,20 @@ def test_credentials_file_forwarded_as_env_path(self): # The path is not a "secret" that blocks acceleration. assert c._accelerator_blocked_reason is None + def test_caller_user_agent_always_forwarded(self): + c = _bare_client() + c._init_accelerator_config(explicit_credentials=False, client_options=None) + flags = c._accelerator_flags + expected = f"python-v{google.cloud.bigtable.__version__}" + assert flags[flags.index("--caller-user-agent") + 1] == expected + + def test_caller_user_agent_forwarded_even_when_blocked(self): + c = _bare_client() + c._init_accelerator_config(explicit_credentials=True, client_options=None) + flags = c._accelerator_flags + expected = f"python-v{google.cloud.bigtable.__version__}" + assert flags[flags.index("--caller-user-agent") + 1] == expected + class TestResolvePrincipal: def test_service_account_email(self): From 19aedd02e3a8c97a751942c1d137fba81497d580 Mon Sep 17 00:00:00 2001 From: Mattie Fu Date: Tue, 13 Oct 2026 15:26:57 +0000 Subject: [PATCH 2/5] test(bigtable): add accelerator unit tests for fallback, routing, dispatch, and daemon lifecycle --- .../bigtable/data/_accelerator/_daemon.py | 50 +- .../cloud/bigtable/data/_async/client.py | 13 + .../bigtable/data/_sync_autogen/client.py | 7 + .../unit/data/test_accelerator_daemon.py | 440 +++++++++++- .../unit/data/test_accelerator_dispatch.py | 667 ++++++++++++++++++ .../unit/data/test_accelerator_fallback.py | 248 +++++++ .../unit/data/test_accelerator_routing.py | 55 ++ 7 files changed, 1468 insertions(+), 12 deletions(-) create mode 100644 packages/google-cloud-bigtable/tests/unit/data/test_accelerator_dispatch.py create mode 100644 packages/google-cloud-bigtable/tests/unit/data/test_accelerator_fallback.py create mode 100644 packages/google-cloud-bigtable/tests/unit/data/test_accelerator_routing.py diff --git a/packages/google-cloud-bigtable/google/cloud/bigtable/data/_accelerator/_daemon.py b/packages/google-cloud-bigtable/google/cloud/bigtable/data/_accelerator/_daemon.py index 9958df4b1a04..524f2ea90648 100644 --- a/packages/google-cloud-bigtable/google/cloud/bigtable/data/_accelerator/_daemon.py +++ b/packages/google-cloud-bigtable/google/cloud/bigtable/data/_accelerator/_daemon.py @@ -33,6 +33,14 @@ import time from typing import Any, Mapping, Sequence + +# Environment variable that overrides the daemon startup readiness timeout, in +# seconds (float). Lets a deployment tune the cold-start budget without a code +# change or rebuild — e.g. raise it on a very constrained/gVisor host, or (when +# a Cloud Run startup probe already gates traffic on daemon readiness) it can be +# left at the default since the platform absorbs the wait. +_STARTUP_TIMEOUT_ENV_VAR = "BIGTABLE_ACCELERATOR_STARTUP_TIMEOUT" + # Wheels ship the binary at this path relative to the `_accelerator/` package. _DEFAULT_BIN_RELATIVE_PATH = "bin/accelerator" @@ -42,8 +50,13 @@ _IDENTITY_FILENAME = "identity.json" # How long to wait for the daemon to start listening on its UDS before giving -# up at startup. -_DEFAULT_STARTUP_TIMEOUT = 10.0 +# up at startup. Cold, constrained environments (e.g. Cloud Run's gVisor +# sandbox) spend ~10s in Go runtime init + ADC token fetch from the metadata +# server before the daemon binds, so a 10s budget misses and falls back to the +# native client; 20s covers that cold start with margin. Override per-deployment +# via the ``BIGTABLE_ACCELERATOR_STARTUP_TIMEOUT`` env var (see +# ``_resolve_startup_timeout``). +_DEFAULT_STARTUP_TIMEOUT = 20.0 # Sequence: close stdin, wait this long; SIGTERM, wait again; SIGKILL. _STDIN_GRACE_SECONDS = 2.0 @@ -61,6 +74,30 @@ def _resolve_binary_path() -> str: return bundled +def _resolve_startup_timeout() -> float: + """Resolve the daemon startup timeout in seconds. + + Precedence: an explicit constructor argument wins; otherwise the + ``BIGTABLE_ACCELERATOR_STARTUP_TIMEOUT`` env var; otherwise the default. + A malformed or non-positive env value is a deployment misconfiguration, so + it is raised loudly rather than silently ignored. + """ + raw = os.environ.get(_STARTUP_TIMEOUT_ENV_VAR) + if raw is None or raw == "": + return _DEFAULT_STARTUP_TIMEOUT + try: + value = float(raw) + except ValueError as exc: + raise ValueError( + f"{_STARTUP_TIMEOUT_ENV_VAR}={raw!r} is not a valid number of seconds" + ) from exc + if value <= 0: + raise ValueError( + f"{_STARTUP_TIMEOUT_ENV_VAR}={raw!r} must be a positive number of seconds" + ) + return value + + class AcceleratorDaemon: """Manages the Go accelerator daemon subprocess. @@ -83,7 +120,7 @@ def __init__( self, cli_flags: Sequence[str] = (), *, - startup_timeout: float = _DEFAULT_STARTUP_TIMEOUT, + binary_path: str | None = None, extra_env: Mapping[str, str] | None = None, ): """Resolve the binary and pick the UDS path (does not spawn anything). @@ -97,9 +134,12 @@ def __init__( Raises: FileNotFoundError: the bundled binary is not present in this wheel. """ - self._binary_path = _resolve_binary_path() + self._binary_path = binary_path if binary_path is not None else _resolve_binary_path() self._cli_flags = list(cli_flags) - self._startup_timeout = startup_timeout + # ``None`` (the default) defers to the env var / module default; an + # explicit float overrides both. Resolved here, not as a signature + # default, so the env var is read at instantiation time. + self._startup_timeout = _resolve_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. diff --git a/packages/google-cloud-bigtable/google/cloud/bigtable/data/_async/client.py b/packages/google-cloud-bigtable/google/cloud/bigtable/data/_async/client.py index be86ecdd9f32..cbaf1f187593 100644 --- a/packages/google-cloud-bigtable/google/cloud/bigtable/data/_async/client.py +++ b/packages/google-cloud-bigtable/google/cloud/bigtable/data/_async/client.py @@ -200,6 +200,11 @@ class _AcceleratorIdentityError(RuntimeError): this client's — signals the caller to fall back to the native client.""" +# Alias used in _start_accelerator's except clause; identity and "unverified" +# refer to the same failure mode. +_AcceleratorUnverified = _AcceleratorIdentityError + + def _normalize_scopes(scopes: Any) -> frozenset[str]: """Normalize a scope value to a comparable, order-insensitive set. @@ -419,6 +424,14 @@ def _init_accelerator_config( "client_cert_source (mTLS) cannot be forwarded to the " "accelerator daemon" ) + # User-Agent prefix always forwarded so the daemon can construct the full + # "python-v{ver}-go-accelerator-v{ver}" string even when the daemon is + # ultimately blocked (the flag is harmless when the daemon never starts). + self._accelerator_flags += [ + "--caller-user-agent", + f"python-v{google.cloud.bigtable.__version__}", + ] + if self._accelerator_blocked_reason: return diff --git a/packages/google-cloud-bigtable/google/cloud/bigtable/data/_sync_autogen/client.py b/packages/google-cloud-bigtable/google/cloud/bigtable/data/_sync_autogen/client.py index ffe42eadc505..a714449339c8 100644 --- a/packages/google-cloud-bigtable/google/cloud/bigtable/data/_sync_autogen/client.py +++ b/packages/google-cloud-bigtable/google/cloud/bigtable/data/_sync_autogen/client.py @@ -148,6 +148,9 @@ class _AcceleratorIdentityError(RuntimeError): this client's — signals the caller to fall back to the native client.""" +_AcceleratorUnverified = _AcceleratorIdentityError + + def _normalize_scopes(scopes: Any) -> frozenset[str]: """Normalize a scope value to a comparable, order-insensitive set. @@ -317,6 +320,10 @@ def _init_accelerator_config( client_options, "client_cert_source", None ): self._accelerator_blocked_reason = "client_cert_source (mTLS) cannot be forwarded to the accelerator daemon" + self._accelerator_flags += [ + "--caller-user-agent", + f"python-v{google.cloud.bigtable.__version__}", + ] if self._accelerator_blocked_reason: return scopes = getattr(client_options, "scopes", None) if client_options else None diff --git a/packages/google-cloud-bigtable/tests/unit/data/test_accelerator_daemon.py b/packages/google-cloud-bigtable/tests/unit/data/test_accelerator_daemon.py index e295997e6f50..5792bc89dfa2 100644 --- a/packages/google-cloud-bigtable/tests/unit/data/test_accelerator_daemon.py +++ b/packages/google-cloud-bigtable/tests/unit/data/test_accelerator_daemon.py @@ -21,21 +21,21 @@ from google.cloud.bigtable.data._accelerator import _daemon from google.cloud.bigtable.data._accelerator._daemon import ( + _DEFAULT_STARTUP_TIMEOUT, _IDENTITY_FILENAME, + _STARTUP_TIMEOUT_ENV_VAR, AcceleratorDaemon, + _resolve_startup_timeout, # noqa: F401 (used indirectly by AcceleratorDaemon) ) def _make_daemon(tmp_path, **kwargs): - # Patch _resolve_binary_path so the constructor doesn't require the bundled - # binary; tests that call start() monkeypatch Popen separately. + # Use binary_path= so the constructor doesn't require the bundled binary; + # tests that call start() monkeypatch Popen separately. fake_binary = tmp_path / "fake-binary" fake_binary.touch() - with unittest.mock.patch( - "google.cloud.bigtable.data._accelerator._daemon._resolve_binary_path", - return_value=str(fake_binary), - ): - return AcceleratorDaemon(**kwargs) + kwargs.setdefault("binary_path", str(fake_binary)) + return AcceleratorDaemon(**kwargs) class TestExtraEnv: @@ -195,3 +195,429 @@ def fake_mkdtemp(prefix=""): monkeypatch.setattr(_daemon.tempfile, "mkdtemp", fake_mkdtemp) yield + + +# --------------------------------------------------------------------------- +# Properties before/after start +# --------------------------------------------------------------------------- + +class TestPropertiesBeforeStart: + def test_uds_path_before_start_raises(self, tmp_path): + d = _make_daemon(tmp_path) + with pytest.raises(RuntimeError, match="has not been started"): + _ = d.uds_path + + def test_log_path_before_start_raises(self, tmp_path): + d = _make_daemon(tmp_path) + with pytest.raises(RuntimeError, match="has not been started"): + _ = d.log_path + + def test_pid_before_start_raises(self, tmp_path): + d = _make_daemon(tmp_path) + with pytest.raises(RuntimeError, match="has not been started"): + _ = d.pid + + def test_is_running_before_start_is_false(self, tmp_path): + d = _make_daemon(tmp_path) + assert d.is_running is False + + def test_is_running_with_live_proc(self, tmp_path): + d = _make_daemon(tmp_path) + proc_mock = unittest.mock.Mock() + proc_mock.poll.return_value = None # still running + d._proc = proc_mock + assert d.is_running is True + + def test_is_running_with_dead_proc(self, tmp_path): + d = _make_daemon(tmp_path) + proc_mock = unittest.mock.Mock() + proc_mock.poll.return_value = 0 # exited + d._proc = proc_mock + assert d.is_running is False + + +# --------------------------------------------------------------------------- +# Binary resolution +# --------------------------------------------------------------------------- + +class TestResolveBinaryPath: + def test_missing_binary_raises_file_not_found(self, tmp_path): + from google.cloud.bigtable.data._accelerator._daemon import _resolve_binary_path + + with unittest.mock.patch("os.path.isfile", return_value=False): + with pytest.raises(FileNotFoundError, match="Accelerator binary not found"): + _resolve_binary_path() + + def test_binary_path_param_bypasses_lookup(self, tmp_path): + """Passing binary_path= directly skips _resolve_binary_path entirely.""" + fake_bin = tmp_path / "my-binary" + fake_bin.touch() + d = AcceleratorDaemon(binary_path=str(fake_bin)) + assert d._binary_path == str(fake_bin) + + +# --------------------------------------------------------------------------- +# start() failure modes +# --------------------------------------------------------------------------- + +class TestStartFailures: + def test_called_twice_raises(self, tmp_path, monkeypatch): + def fake_popen(argv, **kwargs): + raise _StopStart() + + monkeypatch.setattr(_daemon.subprocess, "Popen", fake_popen) + d = _make_daemon(tmp_path) + # Manually set _proc so it looks like it was already started. + d._proc = unittest.mock.Mock() + with pytest.raises(RuntimeError, match="called twice"): + d.start() + + def test_popen_oserror_raises_runtime_error(self, tmp_path, monkeypatch): + def bad_popen(argv, **kwargs): + raise OSError("permission denied") + + monkeypatch.setattr(_daemon.subprocess, "Popen", bad_popen) + d = _make_daemon(tmp_path) + with pytest.raises(RuntimeError, match="Failed to spawn"): + d.start() + + def test_process_exits_during_startup_raises(self, tmp_path, monkeypatch): + import subprocess + + proc_mock = unittest.mock.Mock(spec=subprocess.Popen) + proc_mock.poll.return_value = 1 # exited immediately + proc_mock.stdin = unittest.mock.Mock() + + monkeypatch.setattr( + _daemon.subprocess, "Popen", lambda *a, **kw: proc_mock + ) + # _cleanup_tempdir is already a no-op via autouse fixture, but + # _force_kill is too — that's fine, we just need the start to raise. + monkeypatch.setenv(_STARTUP_TIMEOUT_ENV_VAR, "0.1") + d = _make_daemon(tmp_path) + with pytest.raises(RuntimeError, match="exited during startup"): + d.start() + + def test_uds_never_appears_raises_timeout(self, tmp_path, monkeypatch): + import subprocess + + proc_mock = unittest.mock.Mock(spec=subprocess.Popen) + proc_mock.poll.return_value = None # alive but UDS never arrives + proc_mock.stdin = unittest.mock.Mock() + + monkeypatch.setattr( + _daemon.subprocess, "Popen", lambda *a, **kw: proc_mock + ) + # Make os.path.exists always False so UDS never appears. + monkeypatch.setattr(_daemon.os.path, "exists", lambda p: False) + monkeypatch.setenv(_STARTUP_TIMEOUT_ENV_VAR, "0.05") + d = _make_daemon(tmp_path) + with pytest.raises(RuntimeError, match="did not become ready"): + d.start() + + def test_no_af_unix_raises_os_error(self, tmp_path, monkeypatch): + import socket as _socket + monkeypatch.delattr(_socket, "AF_UNIX", raising=False) + # Reload the attribute check in _daemon + monkeypatch.setattr(_daemon.socket, "AF_UNIX", None, raising=False) + # Simpler: just check that hasattr(socket, "AF_UNIX") guards the call. + with unittest.mock.patch.object( + _daemon.socket, "AF_UNIX", None, create=True + ), unittest.mock.patch( + "google.cloud.bigtable.data._accelerator._daemon.hasattr", + return_value=False, + ): + d = _make_daemon(tmp_path) + with pytest.raises(OSError, match="AF_UNIX"): + d.start() + + +# --------------------------------------------------------------------------- +# close() lifecycle +# --------------------------------------------------------------------------- + +class TestClose: + def test_close_before_start_is_noop(self, tmp_path): + d = _make_daemon(tmp_path) + d.close() # should not raise + + def test_close_graceful_via_stdin_eof(self, tmp_path): + import subprocess + + proc_mock = unittest.mock.Mock(spec=subprocess.Popen) + proc_mock.poll.return_value = None + proc_mock.stdin = unittest.mock.Mock() + # First wait call returns True (exited after stdin close). + d = _make_daemon(tmp_path) + d._proc = proc_mock + + with unittest.mock.patch.object(d, "_wait_for_exit", return_value=True), \ + unittest.mock.patch.object(d, "_cleanup_tempdir"): + d.close() + + proc_mock.stdin.close.assert_called_once() + proc_mock.terminate.assert_not_called() + proc_mock.kill.assert_not_called() + + def test_close_escalates_to_sigterm(self, tmp_path): + import subprocess + + proc_mock = unittest.mock.Mock(spec=subprocess.Popen) + proc_mock.poll.return_value = None + proc_mock.stdin = unittest.mock.Mock() + d = _make_daemon(tmp_path) + d._proc = proc_mock + + wait_results = iter([False, True]) # stdin wait fails, sigterm wait succeeds + + with unittest.mock.patch.object(d, "_wait_for_exit", side_effect=wait_results), \ + unittest.mock.patch.object(d, "_cleanup_tempdir"): + d.close() + + proc_mock.terminate.assert_called_once() + proc_mock.kill.assert_not_called() + + def test_close_escalates_to_sigkill(self, tmp_path): + import subprocess + + proc_mock = unittest.mock.Mock(spec=subprocess.Popen) + proc_mock.poll.return_value = None + proc_mock.stdin = unittest.mock.Mock() + d = _make_daemon(tmp_path) + d._proc = proc_mock + + # Both stdin and sigterm waits fail → sigkill. + wait_results = iter([False, False, True]) + + with unittest.mock.patch.object(d, "_wait_for_exit", side_effect=wait_results), \ + unittest.mock.patch.object(d, "_cleanup_tempdir"): + d.close() + + proc_mock.terminate.assert_called_once() + proc_mock.kill.assert_called_once() + + def test_close_clears_proc_reference(self, tmp_path): + import subprocess + + proc_mock = unittest.mock.Mock(spec=subprocess.Popen) + proc_mock.poll.return_value = None + proc_mock.stdin = unittest.mock.Mock() + d = _make_daemon(tmp_path) + d._proc = proc_mock + + with unittest.mock.patch.object(d, "_wait_for_exit", return_value=True), \ + unittest.mock.patch.object(d, "_cleanup_tempdir"): + d.close() + + assert d._proc is None + + +# --------------------------------------------------------------------------- +# Context manager +# --------------------------------------------------------------------------- + +class TestContextManager: + def test_enter_calls_start_and_returns_self(self, tmp_path, monkeypatch): + def fake_popen(argv, **kwargs): + raise _StopStart() + + monkeypatch.setattr(_daemon.subprocess, "Popen", fake_popen) + d = _make_daemon(tmp_path) + with unittest.mock.patch.object(d, "start") as start_mock, \ + unittest.mock.patch.object(d, "close"): + result = d.__enter__() + start_mock.assert_called_once() + assert result is d + + def test_exit_calls_close(self, tmp_path): + d = _make_daemon(tmp_path) + with unittest.mock.patch.object(d, "close") as close_mock: + d.__exit__(None, None, None) + close_mock.assert_called_once() + + +# --------------------------------------------------------------------------- +# _read_log_tail +# --------------------------------------------------------------------------- + +class TestReadLogTail: + def test_reads_tail_of_log(self, tmp_path): + d = _make_daemon(tmp_path) + log = tmp_path / "daemon.log" + log.write_bytes(b"hello world") + d._log_path = str(log) + tail = d._read_log_tail() + assert "hello world" in tail + + def test_truncates_to_max_bytes(self, tmp_path): + d = _make_daemon(tmp_path) + log = tmp_path / "daemon.log" + log.write_bytes(b"x" * 8192) + d._log_path = str(log) + tail = d._read_log_tail(max_bytes=100) + assert len(tail) <= 100 + + def test_missing_log_returns_empty(self, tmp_path): + d = _make_daemon(tmp_path) + d._log_path = str(tmp_path / "nonexistent.log") + assert d._read_log_tail() == "" + + def test_no_log_path_returns_empty(self, tmp_path): + d = _make_daemon(tmp_path) + d._log_path = None + assert d._read_log_tail() == "" + + +# --------------------------------------------------------------------------- +# _wait_for_exit +# --------------------------------------------------------------------------- + +class TestWaitForExit: + def test_returns_true_when_no_proc(self, tmp_path): + d = _make_daemon(tmp_path) + assert d._wait_for_exit(1.0) is True + + def test_returns_true_when_proc_exits_in_time(self, tmp_path): + import subprocess + proc_mock = unittest.mock.Mock(spec=subprocess.Popen) + proc_mock.wait.return_value = 0 # exited + d = _make_daemon(tmp_path) + d._proc = proc_mock + assert d._wait_for_exit(1.0) is True + + def test_returns_false_on_timeout(self, tmp_path): + import subprocess + proc_mock = unittest.mock.Mock(spec=subprocess.Popen) + proc_mock.wait.side_effect = subprocess.TimeoutExpired(cmd="x", timeout=0.1) + d = _make_daemon(tmp_path) + d._proc = proc_mock + assert d._wait_for_exit(0.1) is False + + +# --------------------------------------------------------------------------- +# read_identity: corrupt JSON +# --------------------------------------------------------------------------- + +class TestReadIdentityCorruptJson: + def test_corrupt_json_raises_runtime_error(self, tmp_path): + d = _make_daemon(tmp_path) + dd = tmp_path / "accel-tmp" + dd.mkdir(exist_ok=True) + d._tempdir = str(dd) + identity_path = dd / _IDENTITY_FILENAME + identity_path.write_text("{not valid json}") + with pytest.raises(RuntimeError, match="Failed to read"): + d.read_identity() + + def test_read_identity_unlinks_even_on_error(self, tmp_path): + d = _make_daemon(tmp_path) + dd = tmp_path / "accel-tmp2" + dd.mkdir(exist_ok=True) + d._tempdir = str(dd) + identity_path = dd / _IDENTITY_FILENAME + identity_path.write_text("{bad json}") + try: + d.read_identity() + except RuntimeError: + pass + # File should be unlinked regardless of the parse error. + assert not identity_path.exists() + + +# --------------------------------------------------------------------------- +# CLI flags passed to Popen +# --------------------------------------------------------------------------- + +class TestCliFlags: + def test_cli_flags_appended_to_argv(self, tmp_path, monkeypatch): + captured = {} + + def fake_popen(argv, **kwargs): + captured["argv"] = argv + raise _StopStart() + + monkeypatch.setattr(_daemon.subprocess, "Popen", fake_popen) + d = _make_daemon(tmp_path, cli_flags=["--project", "proj", "--instance", "inst"]) + with pytest.raises(_StopStart): + d.start() + + argv = captured["argv"] + assert "--uds-path" in argv + assert "--project" in argv + assert argv[argv.index("--project") + 1] == "proj" + assert "--instance" in argv + assert argv[argv.index("--instance") + 1] == "inst" + + +# --------------------------------------------------------------------------- +# Startup timeout env variable (_resolve_startup_timeout) +# --------------------------------------------------------------------------- + +class TestResolveStartupTimeout: + """Tests for the BIGTABLE_ACCELERATOR_STARTUP_TIMEOUT env-var override. + + The remote implementation raises ValueError for invalid values rather than + warning and falling back — invalid env vars are deployment misconfigurations + and should surface loudly. + """ + + def test_no_env_var_returns_default(self, monkeypatch): + monkeypatch.delenv(_STARTUP_TIMEOUT_ENV_VAR, raising=False) + assert _resolve_startup_timeout() == _DEFAULT_STARTUP_TIMEOUT + + def test_empty_string_returns_default(self, monkeypatch): + monkeypatch.setenv(_STARTUP_TIMEOUT_ENV_VAR, "") + assert _resolve_startup_timeout() == _DEFAULT_STARTUP_TIMEOUT + + def test_valid_integer_string_overrides(self, monkeypatch): + monkeypatch.setenv(_STARTUP_TIMEOUT_ENV_VAR, "30") + assert _resolve_startup_timeout() == 30.0 + + def test_valid_float_string_overrides(self, monkeypatch): + monkeypatch.setenv(_STARTUP_TIMEOUT_ENV_VAR, "0.5") + assert _resolve_startup_timeout() == 0.5 + + def test_non_numeric_string_raises_value_error(self, monkeypatch): + monkeypatch.setenv(_STARTUP_TIMEOUT_ENV_VAR, "not-a-number") + with pytest.raises(ValueError, match="not a valid number of seconds"): + _resolve_startup_timeout() + + def test_whitespace_only_raises_value_error(self, monkeypatch): + monkeypatch.setenv(_STARTUP_TIMEOUT_ENV_VAR, " ") + with pytest.raises(ValueError, match="not a valid number of seconds"): + _resolve_startup_timeout() + + def test_negative_value_raises_value_error(self, monkeypatch): + monkeypatch.setenv(_STARTUP_TIMEOUT_ENV_VAR, "-1") + with pytest.raises(ValueError, match="must be a positive number of seconds"): + _resolve_startup_timeout() + + def test_zero_raises_value_error(self, monkeypatch): + monkeypatch.setenv(_STARTUP_TIMEOUT_ENV_VAR, "0") + with pytest.raises(ValueError, match="must be a positive number of seconds"): + _resolve_startup_timeout() + + def test_zero_float_raises_value_error(self, monkeypatch): + monkeypatch.setenv(_STARTUP_TIMEOUT_ENV_VAR, "0.0") + with pytest.raises(ValueError, match="must be a positive number of seconds"): + _resolve_startup_timeout() + + def test_error_message_includes_env_var_name_and_value(self, monkeypatch): + monkeypatch.setenv(_STARTUP_TIMEOUT_ENV_VAR, "bad") + with pytest.raises(ValueError) as exc_info: + _resolve_startup_timeout() + msg = str(exc_info.value) + assert _STARTUP_TIMEOUT_ENV_VAR in msg + assert "bad" in msg + + def test_daemon_constructor_picks_up_env_var(self, tmp_path, monkeypatch): + """End-to-end: env var flows into AcceleratorDaemon._startup_timeout.""" + monkeypatch.setenv(_STARTUP_TIMEOUT_ENV_VAR, "42") + d = _make_daemon(tmp_path) + assert d._startup_timeout == 42.0 + + def test_daemon_constructor_invalid_env_var_raises(self, tmp_path, monkeypatch): + """Invalid env var causes constructor to raise ValueError.""" + monkeypatch.setenv(_STARTUP_TIMEOUT_ENV_VAR, "xyz") + with pytest.raises(ValueError, match="not a valid number of seconds"): + _make_daemon(tmp_path) diff --git a/packages/google-cloud-bigtable/tests/unit/data/test_accelerator_dispatch.py b/packages/google-cloud-bigtable/tests/unit/data/test_accelerator_dispatch.py new file mode 100644 index 000000000000..b73245e17a9a --- /dev/null +++ b/packages/google-cloud-bigtable/tests/unit/data/test_accelerator_dispatch.py @@ -0,0 +1,667 @@ +# 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 accelerator request routing/failover and _start_accelerator +lifecycle. Tests cover: + - _use_accelerator() routing predicate + - _read_row_via_accelerator() and _mutate_row_via_accelerator() call paths + - read_row() / mutate_row() fallback to native when accelerator raises + - _start_accelerator() success and failure with cleanup +""" + +from __future__ import annotations + +import asyncio +from types import SimpleNamespace +from unittest import mock + +import pytest +from grpc import StatusCode + +from google.api_core import exceptions as core_exceptions + +from google.cloud.bigtable.data._accelerator._fallback import ( + AcceleratorBreaker, + _AcceleratorFallback, +) +from google.cloud.bigtable.data._accelerator._routing import is_supported +import grpc as _grpc + +from google.cloud.bigtable.data._async.client import ( + _AcceleratorIdentityError, + _DataApiTargetAsync, +) + + +class _FakeRpcError(_grpc.RpcError): + """grpc.RpcError subclass that exposes code()/details()/trailing_metadata().""" + + def __init__(self, status_code, details=""): + self._code = status_code + self._details = details + + def code(self): + return self._code + + def details(self): + return self._details + + def trailing_metadata(self): + return None + + +# --------------------------------------------------------------------------- +# Minimal concrete _DataApiTargetAsync that skips __init__ +# --------------------------------------------------------------------------- + +class _ConcreteTarget(_DataApiTargetAsync): + """Minimal subclass to make the abstract class instantiable.""" + + @property + def _request_path(self): + return {"table_name": "projects/p/instances/i/tables/t"} + + +def _bare_target( + *, + accel_client=None, + accel_daemon=None, + breaker=None, +): + """Build a _DataApiTargetAsync that skips __init__. + + Wired with mock client, metrics, and configurable accelerator state. + """ + from google.cloud.bigtable.data._metrics import BigtableClientSideMetricsController + + t = object.__new__(_ConcreteTarget) + t.client = SimpleNamespace( + project="p", + _emulator_host=None, + _accelerator_blocked_reason=None, + _accelerator_flags=[], + _accelerator_env={}, + _use_accelerator=None, + _metrics=BigtableClientSideMetricsController(handlers=[]), + _gapic_client=mock.Mock(), + ) + t.instance_id = "i" + t.app_profile_id = None + t.default_operation_timeout = 60.0 + t.default_attempt_timeout = 20.0 + t.default_read_rows_operation_timeout = 600.0 + t.default_read_rows_attempt_timeout = 20.0 + t.default_mutate_rows_operation_timeout = 600.0 + t.default_mutate_rows_attempt_timeout = 60.0 + t.default_read_rows_retryable_errors = () + t.default_mutate_rows_retryable_errors = () + t.default_retryable_errors = () + t._accelerator_client = accel_client + t._accelerator_daemon = accel_daemon + t._accelerator_breaker = breaker if breaker is not None else AcceleratorBreaker() + return t + + +# --------------------------------------------------------------------------- +# _use_accelerator predicate +# --------------------------------------------------------------------------- + +class TestUseAccelerator: + def test_returns_true_when_all_conditions_met(self): + t = _bare_target(accel_client=mock.Mock()) + assert t._use_accelerator("read_row") is True + assert t._use_accelerator("mutate_row") is True + + def test_returns_false_when_client_is_none(self): + t = _bare_target(accel_client=None) + assert t._use_accelerator("read_row") is False + + def test_returns_false_for_unsupported_method(self): + t = _bare_target(accel_client=mock.Mock()) + assert t._use_accelerator("read_rows") is False + assert t._use_accelerator("mutate_rows") is False + assert t._use_accelerator("sample_row_keys") is False + + def test_returns_false_when_breaker_tripped(self): + breaker = AcceleratorBreaker() + breaker.trip() + t = _bare_target(accel_client=mock.Mock(), breaker=breaker) + assert t._use_accelerator("read_row") is False + + def test_becomes_false_after_trip(self): + t = _bare_target(accel_client=mock.Mock()) + assert t._use_accelerator("read_row") is True + t._accelerator_breaker.trip() + assert t._use_accelerator("read_row") is False + + +# --------------------------------------------------------------------------- +# _mutate_row_via_accelerator +# --------------------------------------------------------------------------- + +class TestMutateRowViaAccelerator: + def _run(self, coro): + return asyncio.get_event_loop().run_until_complete(coro) + + def test_raises_fallback_when_client_is_none(self): + t = _bare_target() + # accel_client is None by default + + from google.cloud.bigtable_v2.types.bigtable import MutateRowRequest + + req = MutateRowRequest() + with pytest.raises(_AcceleratorFallback): + self._run(t._mutate_row_via_accelerator(req, 30.0)) + + def test_returns_response_on_success(self): + from google.cloud.bigtable_v2.types.bigtable import MutateRowResponse, MutateRowRequest + + resp = MutateRowResponse() + client_mock = mock.AsyncMock() + client_mock.mutate_row.return_value = resp + t = _bare_target(accel_client=client_mock, accel_daemon=mock.Mock(is_running=True)) + + req = MutateRowRequest() + result = self._run(t._mutate_row_via_accelerator(req, 30.0)) + assert result == resp + client_mock.mutate_row.assert_awaited_once_with(req, timeout=30.0) + + def test_unimplemented_trips_breaker_and_raises_fallback(self): + from google.cloud.bigtable_v2.types.bigtable import MutateRowRequest + + rpc_exc = _FakeRpcError(StatusCode.UNIMPLEMENTED, "no sessions") + client_mock = mock.AsyncMock() + client_mock.mutate_row.side_effect = rpc_exc + daemon = mock.Mock(is_running=True) + t = _bare_target(accel_client=client_mock, accel_daemon=daemon) + + req = MutateRowRequest() + with pytest.raises(_AcceleratorFallback): + self._run(t._mutate_row_via_accelerator(req, 30.0)) + assert t._accelerator_breaker.bypass() + + def test_dead_daemon_trips_breaker_and_raises_fallback(self): + from google.cloud.bigtable_v2.types.bigtable import MutateRowRequest + + rpc_exc = _FakeRpcError(StatusCode.INTERNAL, "channel closed") + client_mock = mock.AsyncMock() + client_mock.mutate_row.side_effect = rpc_exc + daemon = mock.Mock(is_running=False) + t = _bare_target(accel_client=client_mock, accel_daemon=daemon) + + req = MutateRowRequest() + with pytest.raises(_AcceleratorFallback): + self._run(t._mutate_row_via_accelerator(req, 30.0)) + assert t._accelerator_breaker.bypass() + + def test_not_found_translates_and_raises(self): + from google.cloud.bigtable_v2.types.bigtable import MutateRowRequest + + rpc_exc = _FakeRpcError(StatusCode.NOT_FOUND, "row not found") + client_mock = mock.AsyncMock() + client_mock.mutate_row.side_effect = rpc_exc + daemon = mock.Mock(is_running=True) + t = _bare_target(accel_client=client_mock, accel_daemon=daemon) + + req = MutateRowRequest() + with pytest.raises(core_exceptions.GoogleAPICallError): + self._run(t._mutate_row_via_accelerator(req, 30.0)) + # NOT_FOUND is a real server result, not a fallback — breaker stays clear. + assert not t._accelerator_breaker.bypass() + + def test_non_grpc_exception_reraises(self): + from google.cloud.bigtable_v2.types.bigtable import MutateRowRequest + + client_mock = mock.AsyncMock() + client_mock.mutate_row.side_effect = RuntimeError("unexpected") + daemon = mock.Mock(is_running=True) + t = _bare_target(accel_client=client_mock, accel_daemon=daemon) + + req = MutateRowRequest() + with pytest.raises(RuntimeError, match="unexpected"): + self._run(t._mutate_row_via_accelerator(req, 30.0)) + + +# --------------------------------------------------------------------------- +# _read_row_via_accelerator +# --------------------------------------------------------------------------- + +class TestReadRowViaAccelerator: + def _run(self, coro): + return asyncio.get_event_loop().run_until_complete(coro) + + def test_raises_fallback_when_client_is_none(self): + from google.cloud.bigtable.data.read_rows_query import ReadRowsQuery + + t = _bare_target() + query = ReadRowsQuery(row_keys=b"key1", limit=1) + + with pytest.raises(_AcceleratorFallback): + self._run(t._read_row_via_accelerator(query, 30.0, 20.0)) + + def test_unimplemented_trips_breaker_and_raises_fallback(self): + from google.cloud.bigtable.data.read_rows_query import ReadRowsQuery + + rpc_exc = _FakeRpcError(StatusCode.UNIMPLEMENTED, "no sessions") + client_mock = mock.AsyncMock() + client_mock.read_rows.side_effect = rpc_exc + daemon = mock.Mock(is_running=True) + t = _bare_target(accel_client=client_mock, accel_daemon=daemon) + + query = ReadRowsQuery(row_keys=b"key1", limit=1) + with pytest.raises(_AcceleratorFallback): + self._run(t._read_row_via_accelerator(query, 30.0, 20.0)) + assert t._accelerator_breaker.bypass() + + def test_dead_daemon_trips_breaker_and_raises_fallback(self): + from google.cloud.bigtable.data.read_rows_query import ReadRowsQuery + + rpc_exc = _FakeRpcError(StatusCode.UNAVAILABLE, "channel dead") + client_mock = mock.AsyncMock() + client_mock.read_rows.side_effect = rpc_exc + daemon = mock.Mock(is_running=False) + t = _bare_target(accel_client=client_mock, accel_daemon=daemon) + + query = ReadRowsQuery(row_keys=b"key1", limit=1) + with pytest.raises(_AcceleratorFallback): + self._run(t._read_row_via_accelerator(query, 30.0, 20.0)) + assert t._accelerator_breaker.bypass() + + def test_permission_denied_translates_and_raises(self): + from google.cloud.bigtable.data.read_rows_query import ReadRowsQuery + + rpc_exc = _FakeRpcError(StatusCode.PERMISSION_DENIED, "caller lacks permission") + client_mock = mock.AsyncMock() + client_mock.read_rows.side_effect = rpc_exc + daemon = mock.Mock(is_running=True) + t = _bare_target(accel_client=client_mock, accel_daemon=daemon) + + query = ReadRowsQuery(row_keys=b"key1", limit=1) + with pytest.raises(core_exceptions.GoogleAPICallError): + self._run(t._read_row_via_accelerator(query, 30.0, 20.0)) + assert not t._accelerator_breaker.bypass() + + +# --------------------------------------------------------------------------- +# read_row: accelerator dispatch and fallback to native +# --------------------------------------------------------------------------- + +class TestReadRowAcceleratorDispatch: + """Test read_row() routing without spinning up a real client. + + We patch _use_accelerator and _read_row_via_accelerator so these tests are + decoupled from gRPC and the row-merging machinery. + """ + + def _run(self, coro): + return asyncio.get_event_loop().run_until_complete(coro) + + def test_routes_to_accelerator_when_enabled(self): + from google.cloud.bigtable.data.row import Row + + t = _bare_target() + expected_row = Row(b"key1", cells={}) + + async def fake_accel(query, op_timeout, att_timeout): + return expected_row + + with mock.patch.object(t, "_use_accelerator", return_value=True), \ + mock.patch.object(t, "_read_row_via_accelerator", side_effect=fake_accel) as accel_mock: + result = self._run(t.read_row(b"key1")) + + assert result == expected_row + accel_mock.assert_called_once() + + def test_falls_back_to_native_on_accelerator_fallback(self): + from google.cloud.bigtable.data.row import Row + + t = _bare_target() + expected_row = Row(b"key1", cells={}) + + async def fake_accel_fallback(query, op_timeout, att_timeout): + raise _AcceleratorFallback() + + async def _row_gen(): + yield expected_row + + # start_operation must return the async generator directly (not a coroutine). + native_row_merger = mock.Mock() + native_row_merger.start_operation.return_value = _row_gen() + + with mock.patch.object(t, "_use_accelerator", return_value=True), \ + mock.patch.object(t, "_read_row_via_accelerator", side_effect=fake_accel_fallback), \ + mock.patch( + "google.cloud.bigtable.data._async.client.CrossSync._ReadRowsOperation", + return_value=native_row_merger, + ): + result = self._run(t.read_row(b"key1")) + + assert result == expected_row + + def test_uses_native_directly_when_accelerator_disabled(self): + from google.cloud.bigtable.data.row import Row + + t = _bare_target() # _accelerator_client is None → _use_accelerator returns False + expected_row = Row(b"key1", cells={}) + + async def _row_gen(): + yield expected_row + + native_row_merger = mock.Mock() + native_row_merger.start_operation.return_value = _row_gen() + + with mock.patch.object(t, "_use_accelerator", return_value=False), \ + mock.patch( + "google.cloud.bigtable.data._async.client.CrossSync._ReadRowsOperation", + return_value=native_row_merger, + ): + result = self._run(t.read_row(b"key1")) + + assert result == expected_row + + def test_raises_value_error_for_none_key(self): + t = _bare_target() + with pytest.raises(ValueError, match="row_key must be string or bytes"): + self._run(t.read_row(None)) + + def test_returns_none_for_missing_row(self): + t = _bare_target() + + async def _empty_gen(): + return + yield # make it a proper async generator + + native_row_merger = mock.Mock() + native_row_merger.start_operation.return_value = _empty_gen() + + with mock.patch.object(t, "_use_accelerator", return_value=False), \ + mock.patch( + "google.cloud.bigtable.data._async.client.CrossSync._ReadRowsOperation", + return_value=native_row_merger, + ): + result = self._run(t.read_row(b"missing")) + + assert result is None + + +# --------------------------------------------------------------------------- +# mutate_row: accelerator dispatch and fallback to native +# --------------------------------------------------------------------------- + +class TestMutateRowAcceleratorDispatch: + def _run(self, coro): + return asyncio.get_event_loop().run_until_complete(coro) + + def test_routes_to_accelerator_when_enabled(self): + from google.cloud.bigtable.data.mutations import SetCell + from google.cloud.bigtable_v2.types.bigtable import MutateRowResponse + + t = _bare_target() + resp = MutateRowResponse() + + async def fake_accel(request, op_timeout): + return resp + + mutations = [SetCell("cf", "col", b"val", timestamp_micros=1000)] + + with mock.patch.object(t, "_use_accelerator", return_value=True), \ + mock.patch.object(t, "_mutate_row_via_accelerator", side_effect=fake_accel) as accel_mock: + result = self._run(t.mutate_row(b"key1", mutations)) + + assert result == resp + accel_mock.assert_called_once() + + def test_falls_back_to_native_on_accelerator_fallback(self): + from google.cloud.bigtable.data.mutations import SetCell + from google.cloud.bigtable_v2.types.bigtable import MutateRowResponse + + t = _bare_target() + native_resp = MutateRowResponse() + + async def fake_accel_fallback(request, op_timeout): + raise _AcceleratorFallback() + + async def fake_native_mutate(*args, **kwargs): + return native_resp + + mutations = [SetCell("cf", "col", b"val", timestamp_micros=1000)] + + with mock.patch.object(t, "_use_accelerator", return_value=True), \ + mock.patch.object(t, "_mutate_row_via_accelerator", side_effect=fake_accel_fallback), \ + mock.patch( + "google.cloud.bigtable.data._async.client.tracked_retry", + side_effect=fake_native_mutate, + ): + result = self._run(t.mutate_row(b"key1", mutations)) + + assert result == native_resp + + def test_uses_native_directly_when_accelerator_disabled(self): + from google.cloud.bigtable.data.mutations import SetCell + from google.cloud.bigtable_v2.types.bigtable import MutateRowResponse + + t = _bare_target() + native_resp = MutateRowResponse() + mutations = [SetCell("cf", "col", b"val", timestamp_micros=1000)] + + async def fake_native_mutate(*args, **kwargs): + return native_resp + + with mock.patch.object(t, "_use_accelerator", return_value=False), \ + mock.patch( + "google.cloud.bigtable.data._async.client.tracked_retry", + side_effect=fake_native_mutate, + ): + result = self._run(t.mutate_row(b"key1", mutations)) + + assert result == native_resp + + def test_raises_value_error_for_empty_mutations(self): + t = _bare_target() + with pytest.raises(ValueError, match="No mutations provided"): + self._run(t.mutate_row(b"key1", [])) + + +# --------------------------------------------------------------------------- +# _start_accelerator: startup lifecycle +# --------------------------------------------------------------------------- + +class TestStartAccelerator: + """Tests for _start_accelerator() — the method that spawns the daemon and + wires up the AcceleratorClient. These use mocks to avoid real subprocesses.""" + + def test_success_sets_daemon_and_client(self): + from google.cloud.bigtable.data._accelerator._daemon import AcceleratorDaemon + + t = _bare_target() + t._accelerator_daemon = None + t._accelerator_client = None + + fake_daemon = mock.Mock(spec=AcceleratorDaemon) + fake_daemon.uds_path = "/tmp/fake-sock" + fake_daemon.auth_secret = "secret-token" + fake_client = mock.Mock() + + with mock.patch( + "google.cloud.bigtable.data._async.client.AcceleratorDaemon", + return_value=fake_daemon, + ), mock.patch( + "google.cloud.bigtable.data._async.client.AcceleratorClientType", + return_value=fake_client, + ), mock.patch.object(t, "_verify_daemon_identity"): + t._start_accelerator() + + fake_daemon.start.assert_called_once() + assert t._accelerator_daemon is fake_daemon + assert t._accelerator_client is fake_client + + def test_daemon_start_failure_calls_close_and_reraises(self): + from google.cloud.bigtable.data._accelerator._daemon import AcceleratorDaemon + + t = _bare_target() + t._accelerator_daemon = None + t._accelerator_client = None + + fake_daemon = mock.Mock(spec=AcceleratorDaemon) + fake_daemon.start.side_effect = RuntimeError("binary missing") + + with mock.patch( + "google.cloud.bigtable.data._async.client.AcceleratorDaemon", + return_value=fake_daemon, + ), pytest.raises(RuntimeError, match="binary missing"): + t._start_accelerator() + + fake_daemon.close.assert_called_once() + assert t._accelerator_daemon is None + assert t._accelerator_client is None + + def test_identity_verify_failure_closes_daemon_and_warns(self): + """Identity mismatches are now caught in _start_accelerator: daemon is + closed, a RuntimeWarning is issued, and the method returns (no re-raise). + """ + from google.cloud.bigtable.data._accelerator._daemon import AcceleratorDaemon + + t = _bare_target() + t._accelerator_daemon = None + t._accelerator_client = None + + fake_daemon = mock.Mock(spec=AcceleratorDaemon) + fake_daemon.uds_path = "/tmp/fake-sock" + fake_daemon.auth_secret = "secret-token" + + with mock.patch( + "google.cloud.bigtable.data._async.client.AcceleratorDaemon", + return_value=fake_daemon, + ), mock.patch.object( + t, + "_verify_daemon_identity", + side_effect=_AcceleratorIdentityError("mismatch"), + ), pytest.warns(RuntimeWarning, match="could not verify"): + t._start_accelerator() + + fake_daemon.close.assert_called_once() + # After a warn-and-return the target has no daemon or client set. + assert t._accelerator_daemon is None + assert t._accelerator_client is None + + def test_passes_project_instance_flags(self): + from google.cloud.bigtable.data._accelerator._daemon import AcceleratorDaemon + + t = _bare_target() + t.client.project = "my-project" + t.instance_id = "my-instance" + t.app_profile_id = None + t.client._accelerator_flags = [] + t.client._accelerator_env = {} + t._accelerator_daemon = None + t._accelerator_client = None + + fake_daemon = mock.Mock(spec=AcceleratorDaemon) + fake_daemon.uds_path = "/tmp/fake-sock" + captured = {} + + def capture_daemon(cli_flags=(), **kwargs): + captured["cli_flags"] = list(cli_flags) + captured["extra_env"] = kwargs.get("extra_env", {}) + return fake_daemon + + with mock.patch( + "google.cloud.bigtable.data._async.client.AcceleratorDaemon", + side_effect=capture_daemon, + ), mock.patch( + "google.cloud.bigtable.data._async.client.AcceleratorClientType", + ), mock.patch.object(t, "_verify_daemon_identity"): + t._start_accelerator() + + flags = captured["cli_flags"] + assert "--project" in flags + assert flags[flags.index("--project") + 1] == "my-project" + assert "--instance" in flags + assert flags[flags.index("--instance") + 1] == "my-instance" + + def test_passes_app_profile_flag_when_set(self): + from google.cloud.bigtable.data._accelerator._daemon import AcceleratorDaemon + + t = _bare_target() + t.client.project = "p" + t.instance_id = "i" + t.app_profile_id = "my-app-profile" + t.client._accelerator_flags = [] + t.client._accelerator_env = {} + t._accelerator_daemon = None + t._accelerator_client = None + + fake_daemon = mock.Mock(spec=AcceleratorDaemon) + fake_daemon.uds_path = "/tmp/fake-sock" + captured = {} + + def capture_daemon(cli_flags=(), **kwargs): + captured["cli_flags"] = list(cli_flags) + return fake_daemon + + with mock.patch( + "google.cloud.bigtable.data._async.client.AcceleratorDaemon", + side_effect=capture_daemon, + ), mock.patch( + "google.cloud.bigtable.data._async.client.AcceleratorClientType", + ), mock.patch.object(t, "_verify_daemon_identity"): + t._start_accelerator() + + flags = captured["cli_flags"] + assert "--app-profile" in flags + assert flags[flags.index("--app-profile") + 1] == "my-app-profile" + + def test_no_app_profile_flag_when_none(self): + from google.cloud.bigtable.data._accelerator._daemon import AcceleratorDaemon + + t = _bare_target() + t.client.project = "p" + t.instance_id = "i" + t.app_profile_id = None + t.client._accelerator_flags = [] + t.client._accelerator_env = {} + t._accelerator_daemon = None + t._accelerator_client = None + + fake_daemon = mock.Mock(spec=AcceleratorDaemon) + fake_daemon.uds_path = "/tmp/fake-sock" + captured = {} + + def capture_daemon(cli_flags=(), **kwargs): + captured["cli_flags"] = list(cli_flags) + return fake_daemon + + with mock.patch( + "google.cloud.bigtable.data._async.client.AcceleratorDaemon", + side_effect=capture_daemon, + ), mock.patch( + "google.cloud.bigtable.data._async.client.AcceleratorClientType", + ), mock.patch.object(t, "_verify_daemon_identity"): + t._start_accelerator() + + assert "--app-profile" not in captured["cli_flags"] + + +# --------------------------------------------------------------------------- +# Event loop fixture so run_until_complete works +# --------------------------------------------------------------------------- + +@pytest.fixture(autouse=True) +def _event_loop(): + loop = asyncio.new_event_loop() + asyncio.set_event_loop(loop) + try: + yield loop + finally: + loop.close() + asyncio.set_event_loop(None) diff --git a/packages/google-cloud-bigtable/tests/unit/data/test_accelerator_fallback.py b/packages/google-cloud-bigtable/tests/unit/data/test_accelerator_fallback.py new file mode 100644 index 000000000000..f0856204f239 --- /dev/null +++ b/packages/google-cloud-bigtable/tests/unit/data/test_accelerator_fallback.py @@ -0,0 +1,248 @@ +# 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 AcceleratorBreaker, _grpc_code, and handle_accelerator_error +in google.cloud.bigtable.data._accelerator._fallback.""" + +import pytest +from unittest import mock + +import grpc +from grpc import StatusCode + +from google.api_core import exceptions as core_exceptions + +from google.cloud.bigtable.data._accelerator._fallback import ( + AcceleratorBreaker, + _AcceleratorFallback, + _grpc_code, + handle_accelerator_error, +) + + +# --------------------------------------------------------------------------- +# AcceleratorBreaker +# --------------------------------------------------------------------------- + +class TestAcceleratorBreaker: + def test_initially_not_bypassed(self): + b = AcceleratorBreaker() + assert not b.bypass() + + def test_trip_sets_bypass(self): + b = AcceleratorBreaker() + b.trip() + assert b.bypass() + + def test_trip_is_idempotent(self): + b = AcceleratorBreaker() + b.trip() + b.trip() + assert b.bypass() + + def test_bypass_is_sticky(self): + b = AcceleratorBreaker() + b.trip() + # Multiple reads are all True after tripping. + assert b.bypass() + assert b.bypass() + + def test_thread_safety(self): + import threading + b = AcceleratorBreaker() + errors = [] + + def trip_and_check(): + try: + b.trip() + assert b.bypass() + except Exception as e: + errors.append(e) + + threads = [threading.Thread(target=trip_and_check) for _ in range(20)] + for t in threads: + t.start() + for t in threads: + t.join() + assert not errors + + +# --------------------------------------------------------------------------- +# _grpc_code +# --------------------------------------------------------------------------- + +class TestGrpcCode: + def test_rpc_error_with_code_method(self): + # grpc.RpcError doesn't define code() in its class spec, so use plain Mock. + exc = mock.Mock() + exc.code.return_value = StatusCode.NOT_FOUND + assert _grpc_code(exc) == StatusCode.NOT_FOUND + + def test_google_api_call_error_with_grpc_status_code_attr(self): + exc = mock.Mock(spec=[]) + exc.grpc_status_code = StatusCode.UNAVAILABLE + # No callable code() method, but has grpc_status_code attribute. + assert _grpc_code(exc) == StatusCode.UNAVAILABLE + + def test_plain_exception_returns_none(self): + assert _grpc_code(ValueError("boom")) is None + + def test_code_method_raises_returns_none(self): + exc = mock.Mock() + exc.code.side_effect = RuntimeError("broken") + assert _grpc_code(exc) is None + + def test_grpc_status_code_wrong_type_returns_none(self): + exc = mock.Mock(spec=[]) + exc.grpc_status_code = "NOT_FOUND" # string, not StatusCode + assert _grpc_code(exc) is None + + +# --------------------------------------------------------------------------- +# handle_accelerator_error +# --------------------------------------------------------------------------- + +class _FakeRpcError(grpc.RpcError): + """Minimal grpc.RpcError subclass for tests — passes isinstance checks. + + Implements the methods that google.api_core.exceptions.from_grpc_error + inspects when translating a gRPC error. + """ + + def __init__(self, status_code: StatusCode, details: str = ""): + self._code = status_code + self._details = details + + def code(self) -> StatusCode: + return self._code + + def details(self) -> str: + return self._details + + def trailing_metadata(self): + return None + + +def _make_rpc_error(code: StatusCode) -> grpc.RpcError: + return _FakeRpcError(code, f"mocked {code}") + + +class TestHandleAcceleratorError: + # -- Dead daemon --------------------------------------------------------- + + def test_dead_daemon_trips_breaker_and_falls_back(self): + daemon = mock.Mock() + daemon.is_running = False + breaker = AcceleratorBreaker() + + exc = _make_rpc_error(StatusCode.INTERNAL) + with pytest.raises(_AcceleratorFallback): + handle_accelerator_error(exc, daemon=daemon, breaker=breaker) + assert breaker.bypass() + + def test_dead_daemon_wins_over_unimplemented(self): + """Daemon liveness check runs before status code check.""" + daemon = mock.Mock() + daemon.is_running = False + breaker = AcceleratorBreaker() + + exc = _make_rpc_error(StatusCode.UNIMPLEMENTED) + with pytest.raises(_AcceleratorFallback): + handle_accelerator_error(exc, daemon=daemon, breaker=breaker) + assert breaker.bypass() + + # -- UNIMPLEMENTED ------------------------------------------------------- + + def test_unimplemented_trips_breaker_and_falls_back(self): + daemon = mock.Mock() + daemon.is_running = True + breaker = AcceleratorBreaker() + + exc = _make_rpc_error(StatusCode.UNIMPLEMENTED) + with pytest.raises(_AcceleratorFallback): + handle_accelerator_error(exc, daemon=daemon, breaker=breaker) + assert breaker.bypass() + + def test_unimplemented_without_daemon_object_falls_back(self): + """daemon=None means the client has no daemon reference; UNIMPLEMENTED still trips.""" + breaker = AcceleratorBreaker() + exc = _make_rpc_error(StatusCode.UNIMPLEMENTED) + with pytest.raises(_AcceleratorFallback): + handle_accelerator_error(exc, daemon=None, breaker=breaker) + assert breaker.bypass() + + # -- Other gRPC errors are translated ------------------------------------ + + def test_not_found_translated_and_raised(self): + daemon = mock.Mock() + daemon.is_running = True + breaker = AcceleratorBreaker() + + exc = _make_rpc_error(StatusCode.NOT_FOUND) + with pytest.raises(core_exceptions.GoogleAPICallError): + handle_accelerator_error(exc, daemon=daemon, breaker=breaker) + # Breaker must NOT have been tripped for a real server error. + assert not breaker.bypass() + + def test_permission_denied_translated_and_raised(self): + daemon = mock.Mock() + daemon.is_running = True + breaker = AcceleratorBreaker() + + exc = _make_rpc_error(StatusCode.PERMISSION_DENIED) + with pytest.raises(core_exceptions.GoogleAPICallError): + handle_accelerator_error(exc, daemon=daemon, breaker=breaker) + assert not breaker.bypass() + + def test_internal_error_translated_and_raised(self): + daemon = mock.Mock() + daemon.is_running = True + breaker = AcceleratorBreaker() + + exc = _make_rpc_error(StatusCode.INTERNAL) + with pytest.raises(core_exceptions.GoogleAPICallError): + handle_accelerator_error(exc, daemon=daemon, breaker=breaker) + assert not breaker.bypass() + + # -- Non-gRPC exceptions are re-raised unchanged ------------------------- + + def test_non_grpc_exception_reraises_unchanged(self): + daemon = mock.Mock() + daemon.is_running = True + breaker = AcceleratorBreaker() + + exc = ValueError("bad merge") + with pytest.raises(ValueError, match="bad merge"): + handle_accelerator_error(exc, daemon=daemon, breaker=breaker) + assert not breaker.bypass() + + def test_runtime_error_reraises_unchanged(self): + daemon = mock.Mock() + daemon.is_running = True + breaker = AcceleratorBreaker() + + exc = RuntimeError("internal logic bug") + with pytest.raises(RuntimeError, match="internal logic bug"): + handle_accelerator_error(exc, daemon=daemon, breaker=breaker) + assert not breaker.bypass() + + # -- daemon=None with alive-ish context ---------------------------------- + + def test_none_daemon_unavailable_translates(self): + """When no daemon reference is supplied and the error is UNAVAILABLE, + translate and raise (not UNIMPLEMENTED, so no fallback).""" + breaker = AcceleratorBreaker() + exc = _make_rpc_error(StatusCode.UNAVAILABLE) + with pytest.raises(core_exceptions.GoogleAPICallError): + handle_accelerator_error(exc, daemon=None, breaker=breaker) + assert not breaker.bypass() diff --git a/packages/google-cloud-bigtable/tests/unit/data/test_accelerator_routing.py b/packages/google-cloud-bigtable/tests/unit/data/test_accelerator_routing.py new file mode 100644 index 000000000000..d37e9b88b65e --- /dev/null +++ b/packages/google-cloud-bigtable/tests/unit/data/test_accelerator_routing.py @@ -0,0 +1,55 @@ +# 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 routing table in _routing.py.""" + +import pytest + +from google.cloud.bigtable.data._accelerator._routing import ( + _ACCELERATOR_SUPPORTED, + is_supported, +) + + +class TestIsSupported: + def test_read_row_supported(self): + assert is_supported("read_row") is True + + def test_mutate_row_supported(self): + assert is_supported("mutate_row") is True + + def test_read_rows_not_supported(self): + # read_rows (plural) is not in the supported set. + assert is_supported("read_rows") is False + + def test_mutate_rows_not_supported(self): + assert is_supported("mutate_rows") is False + + def test_sample_row_keys_not_supported(self): + assert is_supported("sample_row_keys") is False + + def test_bulk_mutate_rows_not_supported(self): + assert is_supported("bulk_mutate_rows") is False + + def test_empty_string_not_supported(self): + assert is_supported("") is False + + def test_unknown_method_not_supported(self): + assert is_supported("some_unknown_rpc") is False + + def test_supported_set_is_frozenset(self): + assert isinstance(_ACCELERATOR_SUPPORTED, frozenset) + + def test_supported_set_contains_expected_methods(self): + assert "read_row" in _ACCELERATOR_SUPPORTED + assert "mutate_row" in _ACCELERATOR_SUPPORTED From 5e1d30dcd1758320eb8cd0d592edb005485d5328 Mon Sep 17 00:00:00 2001 From: Mattie Fu Date: Wed, 14 Oct 2026 01:20:03 +0000 Subject: [PATCH 3/5] review(bigtable): address PR #18686 reviewer comments - Remove _AcceleratorUnverified alias (no longer needed) - Remove duplicate User-Agent flags block from _init_accelerator_config; keep only the early block so the flag is forwarded even when the daemon is blocked - Move _AcceleratorIdentityError handling into _start_accelerator: close daemon, warn 'could not verify daemon identity', return without re-raising - Remove test_binary_path_param_bypasses_lookup (binary_path is a test-only knob already exercised implicitly) - Add test_close_clears_all_state to verify _proc/_tempdir/_uds_path/_log_path/_auth_secret are all None after close() - Add --app-profile to TestCliFlags.test_cli_flags_appended_to_argv - Add assert not t._accelerator_breaker.bypass() to test_non_grpc_exception_reraises in both TestMutateRowViaAccelerator and TestReadRowViaAccelerator Change-Id: Iff4489acc63a6e71211939c23db04ff1a5df10ba review(bigtable): address PR #18686 reviewer comments - Remove _AcceleratorUnverified alias (no longer needed) - Remove duplicate User-Agent flags block from _init_accelerator_config; keep only the early block so the flag is forwarded even when the daemon is blocked - Add local import google.cloud.bigtable inside _init_accelerator_config to avoid potential circular import at module level - Move _AcceleratorIdentityError handling into _start_accelerator: close daemon, warn 'could not verify daemon identity', return without re-raising - Remove test_binary_path_param_bypasses_lookup (binary_path is a test-only knob already exercised implicitly) - Add --app-profile to TestCliFlags.test_cli_flags_appended_to_argv - Add assert not t._accelerator_breaker.bypass() to test_non_grpc_exception_reraises in both TestMutateRowViaAccelerator and TestReadRowViaAccelerator Change-Id: I43b354ad6517a2c0b9f6a2294cdd2d1639d7fe8c test(bigtable): add TestCleanupTempdir to verify _cleanup_tempdir clears all state Change-Id: I6d6ec51cc039c78d884451502b368862fe8d4ae3 --- .../cloud/bigtable/data/_async/client.py | 29 ++++------ .../bigtable/data/_sync_autogen/client.py | 24 ++++---- .../unit/data/test_accelerator_daemon.py | 58 ++++++++++++++++--- .../unit/data/test_accelerator_dispatch.py | 14 +++++ 4 files changed, 86 insertions(+), 39 deletions(-) diff --git a/packages/google-cloud-bigtable/google/cloud/bigtable/data/_async/client.py b/packages/google-cloud-bigtable/google/cloud/bigtable/data/_async/client.py index cbaf1f187593..495e46fc796a 100644 --- a/packages/google-cloud-bigtable/google/cloud/bigtable/data/_async/client.py +++ b/packages/google-cloud-bigtable/google/cloud/bigtable/data/_async/client.py @@ -200,11 +200,6 @@ class _AcceleratorIdentityError(RuntimeError): this client's — signals the caller to fall back to the native client.""" -# Alias used in _start_accelerator's except clause; identity and "unverified" -# refer to the same failure mode. -_AcceleratorUnverified = _AcceleratorIdentityError - - def _normalize_scopes(scopes: Any) -> frozenset[str]: """Normalize a scope value to a comparable, order-insensitive set. @@ -427,6 +422,8 @@ def _init_accelerator_config( # User-Agent prefix always forwarded so the daemon can construct the full # "python-v{ver}-go-accelerator-v{ver}" string even when the daemon is # ultimately blocked (the flag is harmless when the daemon never starts). + import google.cloud.bigtable # local import avoids circular dependency + self._accelerator_flags += [ "--caller-user-agent", f"python-v{google.cloud.bigtable.__version__}", @@ -470,14 +467,6 @@ def _init_accelerator_config( credentials_file ) - # Identify this client to the daemon so it can construct the full - # User-Agent prefix (e.g. "python-v3.0.0a0-go-accelerator-v3.x.x"). - # The daemon appends its own "-go-accelerator-vX.Y.Z" suffix. - self._accelerator_flags += [ - "--caller-user-agent", - f"python-v{google.cloud.bigtable.__version__}", - ] - def _resolve_principal(self) -> str | None: """Resolve this client's identity to a principal. @@ -1423,9 +1412,6 @@ def _maybe_start_accelerator(self, *, explicit: bool) -> None: except _AcceleratorIdentityError as exc: if explicit: raise - # Callers using gcloud user ADC (_resolve_principal returns None) - # consistently land here; fail-closed to the native client is - # intentional — user ADC is almost never what the daemon resolved. warnings.warn( f"Accelerator disabled: {exc}", RuntimeWarning, @@ -1472,9 +1458,16 @@ def _start_accelerator(self) -> None: self._accelerator_client = AcceleratorClientType( server.uds_path, server.auth_secret ) - except _AcceleratorIdentityError: + except _AcceleratorIdentityError as exc: server.close() - raise + warnings.warn( + f"Accelerator disabled: could not verify daemon identity: {exc}", + RuntimeWarning, + stacklevel=2, + ) + self._accelerator_daemon = None + self._accelerator_client = None + return except BaseException: server.close() raise diff --git a/packages/google-cloud-bigtable/google/cloud/bigtable/data/_sync_autogen/client.py b/packages/google-cloud-bigtable/google/cloud/bigtable/data/_sync_autogen/client.py index a714449339c8..b470ab06abda 100644 --- a/packages/google-cloud-bigtable/google/cloud/bigtable/data/_sync_autogen/client.py +++ b/packages/google-cloud-bigtable/google/cloud/bigtable/data/_sync_autogen/client.py @@ -148,9 +148,6 @@ class _AcceleratorIdentityError(RuntimeError): this client's — signals the caller to fall back to the native client.""" -_AcceleratorUnverified = _AcceleratorIdentityError - - def _normalize_scopes(scopes: Any) -> frozenset[str]: """Normalize a scope value to a comparable, order-insensitive set. @@ -320,6 +317,8 @@ def _init_accelerator_config( client_options, "client_cert_source", None ): self._accelerator_blocked_reason = "client_cert_source (mTLS) cannot be forwarded to the accelerator daemon" + import google.cloud.bigtable # local import avoids circular dependency + self._accelerator_flags += [ "--caller-user-agent", f"python-v{google.cloud.bigtable.__version__}", @@ -347,11 +346,6 @@ def _init_accelerator_config( self._accelerator_env["GOOGLE_APPLICATION_CREDENTIALS"] = ( credentials_file ) - self._accelerator_flags += [ - "--caller-user-agent", - f"python-v{google.cloud.bigtable.__version__}", - ] - def _resolve_principal(self) -> str | None: """Resolve this client's identity to a principal. @@ -1113,9 +1107,6 @@ def _maybe_start_accelerator(self, *, explicit: bool) -> None: except _AcceleratorIdentityError as exc: if explicit: raise - # Callers using gcloud user ADC (_resolve_principal returns None) - # consistently land here; fail-closed to the native client is - # intentional — user ADC is almost never what the daemon resolved. warnings.warn( f"Accelerator disabled: {exc}", RuntimeWarning, @@ -1155,9 +1146,16 @@ def _start_accelerator(self) -> None: self._accelerator_client = AcceleratorClientType( server.uds_path, server.auth_secret ) - except _AcceleratorIdentityError: + except _AcceleratorIdentityError as exc: server.close() - raise + warnings.warn( + f"Accelerator disabled: could not verify daemon identity: {exc}", + RuntimeWarning, + stacklevel=2, + ) + self._accelerator_daemon = None + self._accelerator_client = None + return except BaseException: server.close() raise diff --git a/packages/google-cloud-bigtable/tests/unit/data/test_accelerator_daemon.py b/packages/google-cloud-bigtable/tests/unit/data/test_accelerator_daemon.py index 5792bc89dfa2..b713bb04c3e7 100644 --- a/packages/google-cloud-bigtable/tests/unit/data/test_accelerator_daemon.py +++ b/packages/google-cloud-bigtable/tests/unit/data/test_accelerator_daemon.py @@ -38,6 +38,12 @@ def _make_daemon(tmp_path, **kwargs): return AcceleratorDaemon(**kwargs) +# Captured at import time, before the autouse fixture patches _cleanup_tempdir +# to a no-op. Tests that need to exercise the real implementation call this +# directly as _real_cleanup_tempdir(daemon_instance). +_real_cleanup_tempdir = AcceleratorDaemon._cleanup_tempdir + + class TestExtraEnv: def test_default_env_none(self, tmp_path, monkeypatch): captured = {} @@ -248,13 +254,6 @@ def test_missing_binary_raises_file_not_found(self, tmp_path): with pytest.raises(FileNotFoundError, match="Accelerator binary not found"): _resolve_binary_path() - def test_binary_path_param_bypasses_lookup(self, tmp_path): - """Passing binary_path= directly skips _resolve_binary_path entirely.""" - fake_bin = tmp_path / "my-binary" - fake_bin.touch() - d = AcceleratorDaemon(binary_path=str(fake_bin)) - assert d._binary_path == str(fake_bin) - # --------------------------------------------------------------------------- # start() failure modes @@ -412,6 +411,40 @@ def test_close_clears_proc_reference(self, tmp_path): assert d._proc is None +# --------------------------------------------------------------------------- +# _cleanup_tempdir: state clearing +# --------------------------------------------------------------------------- + +class TestCleanupTempdir: + def test_clears_all_state_fields(self, tmp_path): + d = _make_daemon(tmp_path) + d._tempdir = str(tmp_path / "accel-tmp") + d._uds_path = str(tmp_path / "sock") + d._log_path = str(tmp_path / "daemon.log") + d._auth_secret = "super-secret" + + _real_cleanup_tempdir(d) + + assert d._tempdir is None + assert d._uds_path is None + assert d._log_path is None + assert d._auth_secret is None + + def test_removes_tempdir_from_disk(self, tmp_path): + d = _make_daemon(tmp_path) + td = tmp_path / "accel-tmp" + td.mkdir() + d._tempdir = str(td) + + _real_cleanup_tempdir(d) + + assert not td.exists() + + def test_noop_when_tempdir_is_none(self, tmp_path): + d = _make_daemon(tmp_path) + _real_cleanup_tempdir(d) # should not raise + + # --------------------------------------------------------------------------- # Context manager # --------------------------------------------------------------------------- @@ -537,7 +570,14 @@ def fake_popen(argv, **kwargs): raise _StopStart() monkeypatch.setattr(_daemon.subprocess, "Popen", fake_popen) - d = _make_daemon(tmp_path, cli_flags=["--project", "proj", "--instance", "inst"]) + d = _make_daemon( + tmp_path, + cli_flags=[ + "--project", "proj", + "--instance", "inst", + "--app-profile", "prof", + ], + ) with pytest.raises(_StopStart): d.start() @@ -547,6 +587,8 @@ def fake_popen(argv, **kwargs): assert argv[argv.index("--project") + 1] == "proj" assert "--instance" in argv assert argv[argv.index("--instance") + 1] == "inst" + assert "--app-profile" in argv + assert argv[argv.index("--app-profile") + 1] == "prof" # --------------------------------------------------------------------------- diff --git a/packages/google-cloud-bigtable/tests/unit/data/test_accelerator_dispatch.py b/packages/google-cloud-bigtable/tests/unit/data/test_accelerator_dispatch.py index b73245e17a9a..36862aa36e56 100644 --- a/packages/google-cloud-bigtable/tests/unit/data/test_accelerator_dispatch.py +++ b/packages/google-cloud-bigtable/tests/unit/data/test_accelerator_dispatch.py @@ -230,6 +230,7 @@ def test_non_grpc_exception_reraises(self): req = MutateRowRequest() with pytest.raises(RuntimeError, match="unexpected"): self._run(t._mutate_row_via_accelerator(req, 30.0)) + assert not t._accelerator_breaker.bypass() # --------------------------------------------------------------------------- @@ -291,6 +292,19 @@ def test_permission_denied_translates_and_raises(self): self._run(t._read_row_via_accelerator(query, 30.0, 20.0)) assert not t._accelerator_breaker.bypass() + def test_non_grpc_exception_reraises(self): + from google.cloud.bigtable.data.read_rows_query import ReadRowsQuery + + client_mock = mock.AsyncMock() + client_mock.read_rows.side_effect = RuntimeError("unexpected") + daemon = mock.Mock(is_running=True) + t = _bare_target(accel_client=client_mock, accel_daemon=daemon) + + query = ReadRowsQuery(row_keys=b"key1", limit=1) + with pytest.raises(RuntimeError, match="unexpected"): + self._run(t._read_row_via_accelerator(query, 30.0, 20.0)) + assert not t._accelerator_breaker.bypass() + # --------------------------------------------------------------------------- # read_row: accelerator dispatch and fallback to native From 43223b2231f656bc522786edfb81abe63d0b0b66 Mon Sep 17 00:00:00 2001 From: Mattie Fu Date: Wed, 14 Oct 2026 15:00:04 +0000 Subject: [PATCH 4/5] fix(bigtable): propagate identity error from _start_accelerator; close auth session Change-Id: I78f226004fdffe98d83f7d24b04557a24e9fc7a7 --- .../cloud/bigtable/data/_async/client.py | 25 ++++++--------- .../bigtable/data/_sync_autogen/client.py | 25 ++++++--------- .../unit/data/test_accelerator_dispatch.py | 12 +++---- .../unit/data/test_accelerator_enablement.py | 32 +++++++++++++++++++ 4 files changed, 55 insertions(+), 39 deletions(-) diff --git a/packages/google-cloud-bigtable/google/cloud/bigtable/data/_async/client.py b/packages/google-cloud-bigtable/google/cloud/bigtable/data/_async/client.py index 495e46fc796a..51aabe3eed0d 100644 --- a/packages/google-cloud-bigtable/google/cloud/bigtable/data/_async/client.py +++ b/packages/google-cloud-bigtable/google/cloud/bigtable/data/_async/client.py @@ -494,14 +494,17 @@ def _resolve_principal(self) -> str | None: # Request.__call__ passes timeout=120 to session.request, so # inject our bound at the session level to always override it. _auth_req = google_auth_requests.Request() - _http = _auth_req.session.request + try: + _http = _auth_req.session.request - def _with_timeout(*args, **kwargs): - kwargs["timeout"] = _PRINCIPAL_RESOLVE_TIMEOUT - return _http(*args, **kwargs) + def _with_timeout(*args, **kwargs): + kwargs["timeout"] = _PRINCIPAL_RESOLVE_TIMEOUT + return _http(*args, **kwargs) - _auth_req.session.request = _with_timeout - creds.refresh(_auth_req) + _auth_req.session.request = _with_timeout + creds.refresh(_auth_req) + finally: + _auth_req.session.close() except Exception as exc: _LOGGER.debug("Compute Engine credential refresh failed: %s", exc) email = getattr(creds, "service_account_email", None) @@ -1458,16 +1461,6 @@ def _start_accelerator(self) -> None: self._accelerator_client = AcceleratorClientType( server.uds_path, server.auth_secret ) - except _AcceleratorIdentityError as exc: - server.close() - warnings.warn( - f"Accelerator disabled: could not verify daemon identity: {exc}", - RuntimeWarning, - stacklevel=2, - ) - self._accelerator_daemon = None - self._accelerator_client = None - return except BaseException: server.close() raise diff --git a/packages/google-cloud-bigtable/google/cloud/bigtable/data/_sync_autogen/client.py b/packages/google-cloud-bigtable/google/cloud/bigtable/data/_sync_autogen/client.py index b470ab06abda..296017b6ec25 100644 --- a/packages/google-cloud-bigtable/google/cloud/bigtable/data/_sync_autogen/client.py +++ b/packages/google-cloud-bigtable/google/cloud/bigtable/data/_sync_autogen/client.py @@ -367,14 +367,17 @@ def _resolve_principal(self) -> str | None: ): try: _auth_req = google_auth_requests.Request() - _http = _auth_req.session.request + try: + _http = _auth_req.session.request - def _with_timeout(*args, **kwargs): - kwargs["timeout"] = _PRINCIPAL_RESOLVE_TIMEOUT - return _http(*args, **kwargs) + def _with_timeout(*args, **kwargs): + kwargs["timeout"] = _PRINCIPAL_RESOLVE_TIMEOUT + return _http(*args, **kwargs) - _auth_req.session.request = _with_timeout - creds.refresh(_auth_req) + _auth_req.session.request = _with_timeout + creds.refresh(_auth_req) + finally: + _auth_req.session.close() except Exception as exc: _LOGGER.debug("Compute Engine credential refresh failed: %s", exc) email = getattr(creds, "service_account_email", None) @@ -1146,16 +1149,6 @@ def _start_accelerator(self) -> None: self._accelerator_client = AcceleratorClientType( server.uds_path, server.auth_secret ) - except _AcceleratorIdentityError as exc: - server.close() - warnings.warn( - f"Accelerator disabled: could not verify daemon identity: {exc}", - RuntimeWarning, - stacklevel=2, - ) - self._accelerator_daemon = None - self._accelerator_client = None - return except BaseException: server.close() raise diff --git a/packages/google-cloud-bigtable/tests/unit/data/test_accelerator_dispatch.py b/packages/google-cloud-bigtable/tests/unit/data/test_accelerator_dispatch.py index 36862aa36e56..ce71d84d938c 100644 --- a/packages/google-cloud-bigtable/tests/unit/data/test_accelerator_dispatch.py +++ b/packages/google-cloud-bigtable/tests/unit/data/test_accelerator_dispatch.py @@ -539,9 +539,10 @@ def test_daemon_start_failure_calls_close_and_reraises(self): assert t._accelerator_daemon is None assert t._accelerator_client is None - def test_identity_verify_failure_closes_daemon_and_warns(self): - """Identity mismatches are now caught in _start_accelerator: daemon is - closed, a RuntimeWarning is issued, and the method returns (no re-raise). + def test_identity_verify_failure_closes_daemon_and_propagates(self): + """Identity mismatches propagate from _start_accelerator so + _maybe_start_accelerator can decide whether to warn (auto) or raise + (explicit). The daemon is closed before the error surfaces. """ from google.cloud.bigtable.data._accelerator._daemon import AcceleratorDaemon @@ -560,13 +561,10 @@ def test_identity_verify_failure_closes_daemon_and_warns(self): t, "_verify_daemon_identity", side_effect=_AcceleratorIdentityError("mismatch"), - ), pytest.warns(RuntimeWarning, match="could not verify"): + ), pytest.raises(_AcceleratorIdentityError, match="mismatch"): t._start_accelerator() fake_daemon.close.assert_called_once() - # After a warn-and-return the target has no daemon or client set. - assert t._accelerator_daemon is None - assert t._accelerator_client is None def test_passes_project_instance_flags(self): from google.cloud.bigtable.data._accelerator._daemon import AcceleratorDaemon diff --git a/packages/google-cloud-bigtable/tests/unit/data/test_accelerator_enablement.py b/packages/google-cloud-bigtable/tests/unit/data/test_accelerator_enablement.py index a9b88e171cff..30aca06e0aaa 100644 --- a/packages/google-cloud-bigtable/tests/unit/data/test_accelerator_enablement.py +++ b/packages/google-cloud-bigtable/tests/unit/data/test_accelerator_enablement.py @@ -152,3 +152,35 @@ def _start(): t = _bare_target(emulator=None, start=_start) with pytest.raises(_AcceleratorIdentityError, match="principal mismatch"): t._maybe_start_accelerator(explicit=True) + + def test_start_accelerator_propagates_identity_error(self, monkeypatch): + """_start_accelerator must not catch _AcceleratorIdentityError internally. + + Previously, _start_accelerator absorbed the error and returned normally, + making _maybe_start_accelerator's explicit=True re-raise unreachable. + """ + import google.cloud.bigtable.data._async.client as _mod + + mock_server = unittest.mock.MagicMock() + mock_server.start.return_value = None + + t = object.__new__(_ConcreteTarget) + t.client = SimpleNamespace( + project="p", + _accelerator_flags=[], + _accelerator_env={}, + ) + t.instance_id = "i" + t.app_profile_id = None + t._accelerator_daemon = None + t._accelerator_client = None + t._verify_daemon_identity = unittest.mock.MagicMock( + side_effect=_AcceleratorIdentityError("identity mismatch") + ) + + monkeypatch.setattr(_mod, "AcceleratorDaemon", lambda **_: mock_server) + + with pytest.raises(_AcceleratorIdentityError, match="identity mismatch"): + t._start_accelerator() + + mock_server.close.assert_called_once() From 60a2601481ef961d16f15252e0b02406201af0d1 Mon Sep 17 00:00:00 2001 From: Mattie Fu Date: Wed, 14 Oct 2026 15:35:15 +0000 Subject: [PATCH 5/5] fix(bigtable): move _accelerator_daemon assignment inside try block to prevent orphaned daemon on BaseException Change-Id: I6f4c18f5b03da26efa9b70ed8c6b4d20f9d4a8df --- .../google/cloud/bigtable/data/_async/client.py | 2 +- .../google/cloud/bigtable/data/_sync_autogen/client.py | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/packages/google-cloud-bigtable/google/cloud/bigtable/data/_async/client.py b/packages/google-cloud-bigtable/google/cloud/bigtable/data/_async/client.py index 51aabe3eed0d..a7aa03b1ea31 100644 --- a/packages/google-cloud-bigtable/google/cloud/bigtable/data/_async/client.py +++ b/packages/google-cloud-bigtable/google/cloud/bigtable/data/_async/client.py @@ -1461,10 +1461,10 @@ def _start_accelerator(self) -> None: self._accelerator_client = AcceleratorClientType( server.uds_path, server.auth_secret ) + self._accelerator_daemon = server except BaseException: server.close() raise - self._accelerator_daemon = server def _use_accelerator(self, method_name: str) -> bool: """Whether this call should be routed through the accelerator daemon. diff --git a/packages/google-cloud-bigtable/google/cloud/bigtable/data/_sync_autogen/client.py b/packages/google-cloud-bigtable/google/cloud/bigtable/data/_sync_autogen/client.py index 296017b6ec25..711021230127 100644 --- a/packages/google-cloud-bigtable/google/cloud/bigtable/data/_sync_autogen/client.py +++ b/packages/google-cloud-bigtable/google/cloud/bigtable/data/_sync_autogen/client.py @@ -1149,10 +1149,10 @@ def _start_accelerator(self) -> None: self._accelerator_client = AcceleratorClientType( server.uds_path, server.auth_secret ) + self._accelerator_daemon = server except BaseException: server.close() raise - self._accelerator_daemon = server def _use_accelerator(self, method_name: str) -> bool: """Whether this call should be routed through the accelerator daemon."""