From 5090a4ef794130e17f41e593ed9632c31b6ba74f Mon Sep 17 00:00:00 2001 From: Konstantinos St Date: Thu, 20 Aug 2026 16:06:48 +0200 Subject: [PATCH 1/5] fix: harden private Storage control-plane authentication (#359) * feat: Create dedicated control plane client for Storage --- src/tower/_storage.py | 236 ++++++++++++++----- src/tower/exceptions.py | 20 ++ tests/tower/test_storage.py | 435 ++++++++++++++++++++++++++++++------ 3 files changed, 572 insertions(+), 119 deletions(-) diff --git a/src/tower/_storage.py b/src/tower/_storage.py index bff78bc1..9975a171 100644 --- a/src/tower/_storage.py +++ b/src/tower/_storage.py @@ -5,11 +5,20 @@ import time from dataclasses import dataclass from datetime import datetime, timedelta, timezone -from http import HTTPStatus -from typing import Any, Optional +from typing import TYPE_CHECKING + +import httpx + +if TYPE_CHECKING: + from pyiceberg.catalog import Catalog -from ._client import _env_client from ._context import TowerContext +from .exceptions import ( + StorageConnectionError, + StorageInvalidCredentialError, + StorageMissingAuthenticationError, +) +from .tower_api_client import AuthenticatedClient from .tower_api_client.api.default import describe_catalog as describe_catalog_api from .tower_api_client.api.default import ( describe_default_catalog as describe_default_catalog_api, @@ -33,13 +42,137 @@ # only cache failed catalog type describe requests for this long # retry only after this period CATALOG_TYPE_FAILURE_CACHE_TTL_SECONDS = 30.0 +DEFAULT_STORAGE_TIMEOUT_SECONDS = 30.0 DEFAULT_CATALOG_PROVISION_RETRY_DELAYS = (0.25, 0.5, 1.0, 2.0) DEFAULT_CATALOG_NAME = "default" DEFAULT_ENVIRONMENT_NAME = "default" TOWER_CATALOG_TYPE = "tower-catalog" +INVALID_CREDENTIAL_SENTINEL = "" logger = logging.getLogger("tower.storage") +def _auth_from_context(context: TowerContext) -> tuple[str, str, str]: + if context.jwt is not None: + token = context.jwt + auth_header_name = "Authorization" + prefix = "Bearer" + source = "TOWER_JWT" + elif context.api_key is not None: + token = context.api_key + auth_header_name = "X-API-Key" + prefix = "" + source = "TOWER_API_KEY" + else: + raise StorageMissingAuthenticationError( + "No Tower authentication found. Set TOWER_API_KEY or TOWER_JWT." + ) + + if token.strip() == INVALID_CREDENTIAL_SENTINEL: + raise StorageInvalidCredentialError( + f"{source} contains the {INVALID_CREDENTIAL_SENTINEL!r} placeholder, " + "not a usable Tower credential." + ) + if not token.strip(): + raise StorageMissingAuthenticationError( + "No Tower authentication found. Set TOWER_API_KEY or TOWER_JWT." + ) + + return token, auth_header_name, prefix + + +def _new_tower_control_plane_client( + *, + base_url: str, + token: str, + auth_header_name: str, + prefix: str, + timeout: float, +) -> AuthenticatedClient: + return AuthenticatedClient( + verify_ssl=True, + base_url=base_url, + token=token, + auth_header_name=auth_header_name, + prefix=prefix, + timeout=httpx.Timeout(timeout), + raise_on_unexpected_status=True, + ) + + +def _auth_hash(token: str, auth_header_name: str, prefix: str) -> str: + presented_auth = f"{auth_header_name}\0{prefix}\0{token}" + return hashlib.sha256(presented_auth.encode("utf-8")).hexdigest() + + +class _StorageResolver: + """Private Tower configuration and authentication for catalog resolution.""" + + def __init__( + self, + *, + environment: str | None = None, + ) -> None: + context = TowerContext.build() + + if environment is not None and not isinstance(environment, str): + raise TypeError("environment must be a string or None") + + if environment is not None and not environment.strip(): + raise ValueError("environment must not be blank") + + self._target_environment = ( + environment or context.environment or DEFAULT_ENVIRONMENT_NAME + ) + self._base_url = _api_base_url(context.tower_url) + self._token, self._auth_header_name, self._auth_prefix = _auth_from_context( + context + ) + self._auth_hash = _auth_hash( + self._token, + self._auth_header_name, + self._auth_prefix, + ) + + def _new_client(self) -> AuthenticatedClient: + return _new_tower_control_plane_client( + base_url=self._base_url, + token=self._token, + auth_header_name=self._auth_header_name, + prefix=self._auth_prefix, + timeout=DEFAULT_STORAGE_TIMEOUT_SECONDS, + ) + + def _request_catalog_credentials( + self, + name: str, + mode: str, + ) -> ErrorModel | VendCatalogCredentialsResponse | None: + body = VendCatalogCredentialsBody(mode=_vend_mode(mode)) + + try: + with self._new_client() as client: + return vend_catalog_credentials_api.sync( + name=name, + client=client, + environment=self._target_environment, + body=body, + ) + except httpx.RequestError as error: + raise StorageConnectionError( + f"Could not connect to Tower at {self._base_url}." + ) from error + + +def _api_base_url(tower_url: str) -> str: + try: + url = httpx.URL(tower_url) + except (TypeError, httpx.InvalidURL) as error: + raise ValueError(f"Invalid Tower URL: {tower_url}") from error + if not url.is_absolute_url or url.scheme not in ("http", "https"): + raise ValueError(f"Invalid Tower URL: {tower_url}") + return str(url.copy_with(path="/v1", query=None, fragment=None)) + + @dataclass class _CachedCredentials: credentials: CatalogCredentials @@ -56,14 +189,14 @@ class _CachedCatalogType: _credential_cache: dict[tuple[str, str, str, str, str], _CachedCredentials] = {} -_catalog_type_cache: dict[tuple[str, str, str], _CachedCatalogType] = {} +_catalog_type_cache: dict[tuple[str, str, str, str], _CachedCatalogType] = {} def get_tower_catalog( name: str = DEFAULT_CATALOG_NAME, - environment: Optional[str] = None, + environment: str | None = None, mode: str = "read", -) -> Any: +) -> Catalog: """ Load a PyIceberg REST catalog using short-lived credentials vended by Tower. """ @@ -73,13 +206,12 @@ def get_tower_catalog( def get_tower_catalog_credentials( name: str = DEFAULT_CATALOG_NAME, - environment: Optional[str] = None, + environment: str | None = None, mode: str = "read", ) -> CatalogCredentials: - ctx = TowerContext.build() - environment = environment or ctx.environment or DEFAULT_ENVIRONMENT_NAME + storage_resolver = _StorageResolver(environment=environment) mode = _normalize_mode(mode) - cache_key = _cache_key(ctx, name, environment, mode) + cache_key = _cache_key(storage_resolver, name, mode) now = datetime.now(timezone.utc) _prune_credential_cache(now) @@ -87,12 +219,12 @@ def get_tower_catalog_credentials( if cached is not None and cached.is_usable(now): return cached.credentials - credentials = _vend_with_default_catalog_fallback(ctx, name, environment, mode) + credentials = _vend_with_default_catalog_fallback(storage_resolver, name, mode) _credential_cache[cache_key] = _CachedCredentials(credentials) return credentials -def load_vended_catalog(name: str, credentials: CatalogCredentials) -> Any: +def load_vended_catalog(name: str, credentials: CatalogCredentials) -> Catalog: from pyiceberg.catalog import load_catalog return load_catalog( @@ -105,20 +237,23 @@ def load_vended_catalog(name: str, credentials: CatalogCredentials) -> Any: def _vend_with_default_catalog_fallback( - ctx: TowerContext, name: str, environment: str, mode: str + storage_resolver: _StorageResolver, + name: str, + mode: str, ) -> CatalogCredentials: - result = _vend_catalog_credentials(ctx, name, environment, mode) + environment = storage_resolver._target_environment + result = storage_resolver._request_catalog_credentials(name, mode) if not _is_not_found(result): return _unwrap_vend_result(result, name, environment) if name == DEFAULT_CATALOG_NAME and environment == DEFAULT_ENVIRONMENT_NAME: - _ensure_legacy_default_catalog(ctx) + _ensure_legacy_default_catalog(storage_resolver) for delay in DEFAULT_CATALOG_PROVISION_RETRY_DELAYS: time.sleep(delay) - result = _vend_catalog_credentials(ctx, name, environment, mode) + result = storage_resolver._request_catalog_credentials(name, mode) if not _is_not_found(result): return _unwrap_vend_result(result, name, environment) - _ensure_legacy_default_catalog(ctx) + _ensure_legacy_default_catalog(storage_resolver) return _unwrap_vend_result(result, name, environment) @@ -127,26 +262,16 @@ def _vend_with_default_catalog_fallback( ) -def _vend_catalog_credentials( - ctx: TowerContext, name: str, environment: str, mode: str -) -> ErrorModel | VendCatalogCredentialsResponse | None: - _ensure_tower_auth(ctx) - body = VendCatalogCredentialsBody(mode=_vend_mode(mode)) - return vend_catalog_credentials_api.sync( - name=name, - client=_env_client(ctx), - environment=environment, - body=body, - ) - - def _describe_tower_catalog_type( ctx: TowerContext, name: str, environment: str ) -> str | None: - if not (ctx.api_key or ctx.jwt): + if ctx.jwt is None and ctx.api_key is None: return None - cache_key = (ctx.tower_url, name, environment) + token, auth_header_name, prefix = _auth_from_context(ctx) + base_url = _api_base_url(ctx.tower_url) + auth_hash = _auth_hash(token, auth_header_name, prefix) + cache_key = (base_url, auth_hash, name, environment) cached = _catalog_type_cache.get(cache_key) if cached is not None: if cached.retry_at is None: @@ -157,12 +282,21 @@ def _describe_tower_catalog_type( _catalog_type_cache.pop(cache_key, None) + tower_client = _new_tower_control_plane_client( + base_url=base_url, + token=token, + auth_header_name=auth_header_name, + prefix=prefix, + timeout=CATALOG_TYPE_DESCRIBE_TIMEOUT_SECONDS, + ) + try: - result = describe_catalog_api.sync( - name=name, - client=_env_client(ctx, timeout=CATALOG_TYPE_DESCRIBE_TIMEOUT_SECONDS), - environment=environment, - ) + with tower_client: + result = describe_catalog_api.sync( + name=name, + client=tower_client, + environment=environment, + ) except Exception: logger.debug( "Failed to describe Tower catalog %r in environment %r; " @@ -199,11 +333,10 @@ def _failed_catalog_type_cache_entry() -> _CachedCatalogType: ) -def _ensure_legacy_default_catalog(ctx: TowerContext) -> None: +def _ensure_legacy_default_catalog(storage_resolver: _StorageResolver) -> None: try: - response = describe_default_catalog_api.sync_detailed(client=_env_client(ctx)) - if response.status_code not in (HTTPStatus.OK, HTTPStatus.ACCEPTED): - return + with storage_resolver._new_client() as client: + describe_default_catalog_api.sync(client=client) except Exception: # The following vend retry will surface the actionable backend/auth error. return @@ -230,19 +363,18 @@ def _unwrap_vend_result( ) -def _ensure_tower_auth(ctx: TowerContext) -> None: - if ctx.api_key or ctx.jwt: - return - - raise RuntimeError("No Tower authentication found. Set TOWER_API_KEY or TOWER_JWT.") - - def _cache_key( - ctx: TowerContext, name: str, environment: str, mode: str + storage_resolver: _StorageResolver, + name: str, + mode: str, ) -> tuple[str, str, str, str, str]: - token = ctx.api_key or ctx.jwt or "" - principal_hash = hashlib.sha256(token.encode("utf-8")).hexdigest() - return (ctx.tower_url, principal_hash, name, environment, mode) + return ( + storage_resolver._base_url, + storage_resolver._auth_hash, + name, + storage_resolver._target_environment, + mode, + ) def _prune_credential_cache(now: datetime) -> None: diff --git a/src/tower/exceptions.py b/src/tower/exceptions.py index 0035f952..0f0bfc9b 100644 --- a/src/tower/exceptions.py +++ b/src/tower/exceptions.py @@ -43,3 +43,23 @@ def __init__(self): "replace [a, b] with a & b. You can also pass a PyIceberg " "BooleanExpression or a SQL-like filter string." ) + + +class StorageError(RuntimeError): + """Base error for Tower Storage control-plane operations.""" + + +class StorageAuthenticationError(StorageError): + """Base error for Storage authentication failures.""" + + +class StorageMissingAuthenticationError(StorageAuthenticationError): + """No supported Tower API key or JWT was available.""" + + +class StorageInvalidCredentialError(StorageAuthenticationError): + """A configured credential is a known placeholder rather than a secret.""" + + +class StorageConnectionError(StorageError): + """Tower's control-plane API could not be reached.""" diff --git a/tests/tower/test_storage.py b/tests/tower/test_storage.py index 2f29c6e1..7f4307fd 100644 --- a/tests/tower/test_storage.py +++ b/tests/tower/test_storage.py @@ -1,7 +1,16 @@ from datetime import datetime, timedelta, timezone +from http import HTTPStatus + +import httpx +import pytest -from tower._context import TowerContext from tower import _storage +from tower._context import TowerContext +from tower.exceptions import ( + StorageConnectionError, + StorageInvalidCredentialError, + StorageMissingAuthenticationError, +) from tower.tower_api_client.models import ( Catalog, CatalogCredentials, @@ -11,7 +20,8 @@ ) -def clear_tower_env(monkeypatch): +@pytest.fixture(autouse=True) +def isolate_tower_environment(monkeypatch, tmp_path): for name in ( "TOWER_URL", "TOWER_ENVIRONMENT", @@ -21,11 +31,10 @@ def clear_tower_env(monkeypatch): "TOWER__RUNTIME__ENVIRONMENT_NAME", ): monkeypatch.delenv(name, raising=False) + monkeypatch.setenv("HOME", str(tmp_path)) -def test_context_prefers_runtime_environment(monkeypatch, tmp_path): - clear_tower_env(monkeypatch) - monkeypatch.setenv("HOME", str(tmp_path)) +def test_context_prefers_runtime_environment(monkeypatch): monkeypatch.setenv("TOWER_ENVIRONMENT", "local-env") monkeypatch.setenv("TOWER__RUNTIME__ENVIRONMENT_NAME", "run-env") @@ -34,9 +43,7 @@ def test_context_prefers_runtime_environment(monkeypatch, tmp_path): assert ctx.environment == "run-env" -def test_context_treats_blank_auth_env_as_missing(monkeypatch, tmp_path): - clear_tower_env(monkeypatch) - monkeypatch.setenv("HOME", str(tmp_path)) +def test_context_treats_blank_auth_env_as_missing(monkeypatch): monkeypatch.setenv("TOWER_URL", "") monkeypatch.setenv("TOWER_API_KEY", "") monkeypatch.setenv("TOWER_JWT", "") @@ -48,41 +55,325 @@ def test_context_treats_blank_auth_env_as_missing(monkeypatch, tmp_path): assert ctx.jwt is None -def test_ensure_tower_auth_requires_explicit_credentials(): - ctx = TowerContext(tower_url="https://api.example.com", environment="production") +def test_storage_resolver_configuration_and_tls_defaults(monkeypatch): + monkeypatch.setenv("TOWER_URL", "https://tower.example.com/") + monkeypatch.setenv("TOWER_API_KEY", "ambient-key") + resolver = _storage._StorageResolver(environment="production") + client = resolver._new_client() + + assert resolver._target_environment == "production" + assert resolver._base_url == "https://tower.example.com/v1" + assert client._base_url == "https://tower.example.com/v1" + assert client._timeout == httpx.Timeout(_storage.DEFAULT_STORAGE_TIMEOUT_SECONDS) + assert client._verify_ssl is True + + http_client = client.get_httpx_client() + try: + assert http_client.headers["X-API-Key"] == "ambient-key" + assert "Authorization" not in http_client.headers + finally: + http_client.close() + +@pytest.mark.parametrize( + ("ambient_auth", "expected_header", "expected_value"), + [ + ( + {"TOWER_API_KEY": "ambient-key", "TOWER_JWT": "ambient-jwt"}, + "Authorization", + "Bearer ambient-jwt", + ), + ({"TOWER_API_KEY": "ambient-key"}, "X-API-Key", "ambient-key"), + ], + ids=("jwt-over-api-key", "api-key-fallback"), +) +def test_storage_resolver_static_auth_precedence( + monkeypatch, + ambient_auth, + expected_header, + expected_value, +): + for name, value in ambient_auth.items(): + monkeypatch.setenv(name, value) + + client = _storage._StorageResolver()._new_client() + http_client = client.get_httpx_client() try: - _storage._ensure_tower_auth(ctx) - except RuntimeError as error: - assert str(error) == ( - "No Tower authentication found. Set TOWER_API_KEY or TOWER_JWT." + assert http_client.headers[expected_header] == expected_value + other_header = ( + "Authorization" if expected_header == "X-API-Key" else "X-API-Key" + ) + assert other_header not in http_client.headers + finally: + http_client.close() + + +def test_auth_hash_includes_how_the_credential_is_presented(): + assert _storage._auth_hash("same-token", "Authorization", "Bearer") != ( + _storage._auth_hash("same-token", "X-API-Key", "") + ) + + +def test_missing_auth_fails_before_cache_or_vend(monkeypatch): + _storage._clear_credential_cache() + + monkeypatch.setattr( + _storage, + "_prune_credential_cache", + lambda now: pytest.fail("cache access must not run without authentication"), + ) + monkeypatch.setattr( + _storage.vend_catalog_credentials_api, + "sync", + lambda **kwargs: pytest.fail("vend must not run without authentication"), + ) + + with pytest.raises(StorageMissingAuthenticationError): + _storage.get_tower_catalog_credentials("analytics") + + +def test_storage_resolver_rejects_redacted_jwt_without_falling_back(monkeypatch): + monkeypatch.setenv("TOWER_JWT", " ") + monkeypatch.setenv("TOWER_API_KEY", "otherwise-valid-api-key") + + with pytest.raises(StorageInvalidCredentialError): + _storage._StorageResolver() + + +@pytest.mark.parametrize("status", [HTTPStatus.UNAUTHORIZED, HTTPStatus.FORBIDDEN]) +def test_static_auth_rejection_is_returned(monkeypatch, status): + monkeypatch.setenv("TOWER_JWT", "ambient-jwt") + rejected = ErrorModel(status=int(status), detail="rejected") + vend_calls = [] + + def vend(**kwargs): + vend_calls.append(kwargs) + return rejected + + monkeypatch.setattr(_storage.vend_catalog_credentials_api, "sync", vend) + + result = _storage._StorageResolver()._request_catalog_credentials( + "analytics", "read" + ) + + assert result is rejected + assert len(vend_calls) == 1 + + +def test_storage_resolver_vends_with_ambient_api_key(monkeypatch): + monkeypatch.setenv("TOWER_URL", "https://api.example.com") + monkeypatch.setenv("TOWER_API_KEY", "service-account-key") + captured = {} + clients = [] + response = ErrorModel(status=418, detail="captured") + + def vend(*, name, client, environment, body): + clients.append(client) + http_client = client.get_httpx_client() + captured.update( + name=name, + environment=environment, + api_key=http_client.headers.get("X-API-Key"), + authorization=http_client.headers.get("Authorization"), + mode=body.mode, + ) + return response + + monkeypatch.setattr(_storage.vend_catalog_credentials_api, "sync", vend) + + result = _storage._StorageResolver( + environment="production" + )._request_catalog_credentials("analytics", "read") + + assert result is response + assert captured == { + "name": "analytics", + "environment": "production", + "api_key": "service-account-key", + "authorization": None, + "mode": _storage.VendCatalogCredentialsBodyMode.READ, + } + assert clients[0]._client is not None + assert clients[0]._client.is_closed + + +def test_describe_and_vend_prefer_jwt_when_both_auth_vars_are_set(monkeypatch): + _storage._clear_credential_cache() + monkeypatch.setenv("TOWER_URL", "https://api.example.com") + monkeypatch.setenv("TOWER_ENVIRONMENT", "production") + monkeypatch.setenv("TOWER_API_KEY", "ambient-api-key") + monkeypatch.setenv("TOWER_JWT", "ambient-jwt") + captured_auth = [] + + def capture_auth(operation, client): + http_client = client.get_httpx_client() + captured_auth.append( + ( + operation, + http_client.headers.get("Authorization"), + http_client.headers.get("X-API-Key"), + ) ) - else: - raise AssertionError("expected missing-auth error") - - _storage._ensure_tower_auth( - TowerContext( - tower_url="https://api.example.com", - environment="production", - api_key="api-key", + + def describe(*, name, client, environment): + capture_auth("describe", client) + return DescribeCatalogResponse( + catalog=Catalog( + created_at=datetime.now(timezone.utc), + environment=environment, + name=name, + properties=[], + type_=_storage.TOWER_CATALOG_TYPE, + ) ) + + vended = ErrorModel(status=418, detail="captured") + + def vend(*, client, **kwargs): + capture_auth("vend", client) + return vended + + monkeypatch.setattr(_storage.describe_catalog_api, "sync", describe) + monkeypatch.setattr(_storage.vend_catalog_credentials_api, "sync", vend) + + ctx = TowerContext.build() + assert ( + _storage._describe_tower_catalog_type(ctx, "analytics", "production") + == _storage.TOWER_CATALOG_TYPE + ) + assert ( + _storage._StorageResolver()._request_catalog_credentials("analytics", "read") + is vended ) - _storage._ensure_tower_auth( - TowerContext( - tower_url="https://api.example.com", - environment="production", - jwt="jwt", + assert captured_auth == [ + ("describe", "Bearer ambient-jwt", None), + ("vend", "Bearer ambient-jwt", None), + ] + + +def test_storage_resolver_allows_explicit_http_tower_url(monkeypatch): + monkeypatch.setenv("TOWER_URL", "http://localhost:9000") + monkeypatch.setenv("TOWER_API_KEY", "key") + client = _storage._StorageResolver()._new_client() + + assert client._base_url == "http://localhost:9000/v1" + assert client._verify_ssl is True + + +def test_get_tower_catalog_credentials_allows_http_and_reaches_vend(monkeypatch): + _storage._clear_credential_cache() + monkeypatch.setenv("TOWER_URL", "http://localhost:9000") + monkeypatch.setenv("TOWER_ENVIRONMENT", "production") + monkeypatch.setenv("TOWER_API_KEY", "api-key") + credentials = CatalogCredentials( + catalog_uri="http://catalog.example.com", + expires_at=datetime.now(timezone.utc) + timedelta(hours=1), + mode="read", + oauth_token="oauth-token", + warehouse="warehouse-id", + ) + vend_calls = [] + + def vend(*, name, client, environment, body): + vend_calls.append( + (name, client._base_url, client._verify_ssl, environment, body.mode) + ) + return VendCatalogCredentialsResponse( + credentials=credentials, + environment=environment, + ) + + monkeypatch.setattr(_storage.vend_catalog_credentials_api, "sync", vend) + + result = _storage.get_tower_catalog_credentials("analytics") + + assert result is credentials + assert vend_calls == [ + ( + "analytics", + "http://localhost:9000/v1", + True, + "production", + _storage.VendCatalogCredentialsBodyMode.READ, ) + ] + + +def test_storage_resolver_normalizes_and_validates_tower_api_url(monkeypatch): + monkeypatch.setenv( + "TOWER_URL", + "https://TOWER.example.com:443/old/path?debug=true#fragment", + ) + monkeypatch.setenv("TOWER_API_KEY", "key") + resolver = _storage._StorageResolver() + assert resolver._base_url == "https://tower.example.com/v1" + + with pytest.raises(ValueError, match="Invalid Tower URL"): + _storage._api_base_url("not-a-url") + + +def test_get_tower_catalog_credentials_rejects_invalid_mode_before_cache_or_vend( + monkeypatch, +): + monkeypatch.setenv("TOWER_API_KEY", "api-key") + monkeypatch.setattr( + _storage, + "_prune_credential_cache", + lambda now: pytest.fail("cache access must not run for an invalid mode"), + ) + monkeypatch.setattr( + _storage.vend_catalog_credentials_api, + "sync", + lambda **kwargs: pytest.fail("vend must not run for an invalid mode"), + ) + + with pytest.raises(ValueError, match="mode must be 'read' or 'read-write'"): + _storage.get_tower_catalog_credentials("analytics", mode="write") + + +def test_storage_resolver_rejects_invalid_environment(): + with pytest.raises(TypeError, match="environment must be a string or None"): + _storage._StorageResolver(environment=123) + with pytest.raises(ValueError, match="environment must not be blank"): + _storage._StorageResolver(environment=" ") + + +def test_storage_resolver_maps_connection_errors_and_closes_client(monkeypatch): + monkeypatch.setenv("TOWER_API_KEY", "key") + cause = httpx.ConnectError("connection refused") + clients = [] + + def vend(*, client, **kwargs): + clients.append(client) + raise cause + + monkeypatch.setattr(_storage.vend_catalog_credentials_api, "sync", vend) + + with pytest.raises(StorageConnectionError) as error: + _storage._StorageResolver()._request_catalog_credentials("analytics", "read") + + assert error.value.__cause__ is cause + assert clients[0]._client is not None + assert clients[0]._client.is_closed + + +def test_storage_resolver_uses_runtime_environment_only_as_target_config(monkeypatch): + monkeypatch.setenv("TOWER_API_KEY", "key") + monkeypatch.setenv("TOWER_ENVIRONMENT", "ambient-env") + monkeypatch.setenv("TOWER__RUNTIME__ENVIRONMENT_NAME", "run-env") + + assert _storage._StorageResolver()._target_environment == "run-env" + assert ( + _storage._StorageResolver(environment="explicit-env")._target_environment + == "explicit-env" ) def test_get_tower_catalog_credentials_caches_vended_credentials(monkeypatch): _storage._clear_credential_cache() - ctx = TowerContext( - tower_url="https://api.example.com", - environment="production", - api_key="api-key", - ) + monkeypatch.setenv("TOWER_URL", "https://api.example.com") + monkeypatch.setenv("TOWER_ENVIRONMENT", "production") + monkeypatch.setenv("TOWER_API_KEY", "api-key") expires_at = datetime.now(timezone.utc) + timedelta(hours=1) credentials = CatalogCredentials( catalog_uri="https://catalog.example.com", @@ -93,15 +384,14 @@ def test_get_tower_catalog_credentials_caches_vended_credentials(monkeypatch): ) calls = [] - def vend(ctx, name, environment, mode): - calls.append((name, environment, mode)) + def vend(client, name, mode): + calls.append((name, client._target_environment, mode)) return VendCatalogCredentialsResponse( credentials=credentials, - environment=environment, + environment=client._target_environment, ) - monkeypatch.setattr(_storage.TowerContext, "build", staticmethod(lambda: ctx)) - monkeypatch.setattr(_storage, "_vend_catalog_credentials", vend) + monkeypatch.setattr(_storage._StorageResolver, "_request_catalog_credentials", vend) first = _storage.get_tower_catalog_credentials("default") second = _storage.get_tower_catalog_credentials("default") @@ -113,11 +403,9 @@ def vend(ctx, name, environment, mode): def test_get_tower_catalog_credentials_prunes_expired_cache_entries(monkeypatch): _storage._clear_credential_cache() - ctx = TowerContext( - tower_url="https://api.example.com", - environment="production", - api_key="api-key", - ) + monkeypatch.setenv("TOWER_URL", "https://api.example.com") + monkeypatch.setenv("TOWER_ENVIRONMENT", "production") + monkeypatch.setenv("TOWER_API_KEY", "api-key") expired_credentials = CatalogCredentials( catalog_uri="https://old-catalog.example.com", expires_at=datetime.now(timezone.utc) - timedelta(minutes=1), @@ -132,19 +420,19 @@ def test_get_tower_catalog_credentials_prunes_expired_cache_entries(monkeypatch) oauth_token="oauth-token", warehouse="warehouse-id", ) - expired_key = _storage._cache_key(ctx, "stale", "production", "read") + storage_resolver = _storage._StorageResolver() + expired_key = _storage._cache_key(storage_resolver, "stale", "read") _storage._credential_cache[expired_key] = _storage._CachedCredentials( expired_credentials ) - def vend(ctx, name, environment, mode): + def vend(client, name, mode): return VendCatalogCredentialsResponse( credentials=fresh_credentials, - environment=environment, + environment=client._target_environment, ) - monkeypatch.setattr(_storage.TowerContext, "build", staticmethod(lambda: ctx)) - monkeypatch.setattr(_storage, "_vend_catalog_credentials", vend) + monkeypatch.setattr(_storage._StorageResolver, "_request_catalog_credentials", vend) result = _storage.get_tower_catalog_credentials("default") @@ -152,17 +440,12 @@ def vend(ctx, name, environment, mode): assert expired_key not in _storage._credential_cache -def test_default_catalog_vend_retries_after_legacy_provisioning(monkeypatch): +def test_default_catalog_retries_reuse_client_auth_snapshot(monkeypatch): _storage._clear_credential_cache() - ctx = TowerContext( - tower_url="https://api.example.com", - environment="default", - api_key="api-key", - ) - expires_at = datetime.now(timezone.utc) + timedelta(hours=1) + monkeypatch.setenv("TOWER_API_KEY", "operation-token") credentials = CatalogCredentials( catalog_uri="https://catalog.example.com", - expires_at=expires_at, + expires_at=datetime.now(timezone.utc) + timedelta(hours=1), mode="read", oauth_token="oauth-token", warehouse="warehouse-id", @@ -175,23 +458,37 @@ def test_default_catalog_vend_retries_after_legacy_provisioning(monkeypatch): environment="default", ), ] - legacy_calls = [] - - def vend(ctx, name, environment, mode): + vend_tokens = [] + legacy_tokens = [] + clients = [] + + def vend(*, client, **kwargs): + clients.append(client) + vend_tokens.append(client.token) + monkeypatch.setenv("TOWER_API_KEY", "changed-after-client-construction") return responses.pop(0) - def legacy_default(ctx): - legacy_calls.append(ctx) + def legacy_default(*, client): + clients.append(client) + legacy_tokens.append(client.token) + return ErrorModel(status=404, detail="not provisioned") - monkeypatch.setattr(_storage.TowerContext, "build", staticmethod(lambda: ctx)) - monkeypatch.setattr(_storage, "_vend_catalog_credentials", vend) - monkeypatch.setattr(_storage, "_ensure_legacy_default_catalog", legacy_default) + monkeypatch.setattr(_storage.vend_catalog_credentials_api, "sync", vend) + monkeypatch.setattr( + _storage.describe_default_catalog_api, + "sync", + legacy_default, + ) monkeypatch.setattr(_storage.time, "sleep", lambda delay: None) result = _storage.get_tower_catalog_credentials("default") assert result is credentials - assert len(legacy_calls) == 2 + assert vend_tokens == ["operation-token"] * 3 + assert legacy_tokens == ["operation-token"] * 2 + assert len({id(client) for client in clients}) == 5 + assert all(client._client is not None for client in clients) + assert all(client._client.is_closed for client in clients) def test_describe_tower_catalog_type_uses_timeout_and_recovers_after_cooldown( @@ -205,8 +502,10 @@ def test_describe_tower_catalog_type_uses_timeout_and_recovers_after_cooldown( ) now = {"value": 100.0} calls = [] + transports = [] def describe_catalog_api_sync(name, client, environment): + transports.append(client) calls.append((name, environment, client._timeout)) if len(calls) == 1: raise TimeoutError("describe timed out") @@ -232,7 +531,7 @@ def describe_catalog_api_sync(name, client, environment): ( "s3-tables", "production", - _storage.CATALOG_TYPE_DESCRIBE_TIMEOUT_SECONDS, + httpx.Timeout(_storage.CATALOG_TYPE_DESCRIBE_TIMEOUT_SECONDS), ) ] @@ -250,11 +549,13 @@ def describe_catalog_api_sync(name, client, environment): ( "s3-tables", "production", - _storage.CATALOG_TYPE_DESCRIBE_TIMEOUT_SECONDS, + httpx.Timeout(_storage.CATALOG_TYPE_DESCRIBE_TIMEOUT_SECONDS), ), ( "s3-tables", "production", - _storage.CATALOG_TYPE_DESCRIBE_TIMEOUT_SECONDS, + httpx.Timeout(_storage.CATALOG_TYPE_DESCRIBE_TIMEOUT_SECONDS), ), ] + assert all(transport._client is not None for transport in transports) + assert all(transport._client.is_closed for transport in transports) From 4946b4319e338cdb756c6292be8917f034cfb140 Mon Sep 17 00:00:00 2001 From: Konstantinos St Date: Fri, 21 Aug 2026 16:09:08 +0200 Subject: [PATCH 2/5] feat: cache resolved catalog access per private resolver (#361) --- src/tower/_storage.py | 207 ++++++++++++++++-------- tests/tower/test_storage.py | 302 ++++++++++++++++++++++++++---------- tests/tower/test_tables.py | 10 +- 3 files changed, 372 insertions(+), 147 deletions(-) diff --git a/src/tower/_storage.py b/src/tower/_storage.py index 9975a171..c0db1286 100644 --- a/src/tower/_storage.py +++ b/src/tower/_storage.py @@ -3,8 +3,10 @@ import hashlib import logging import time -from dataclasses import dataclass +from concurrent.futures import Future +from dataclasses import dataclass, field from datetime import datetime, timedelta, timezone +from threading import Lock from typing import TYPE_CHECKING import httpx @@ -15,6 +17,7 @@ from ._context import TowerContext from .exceptions import ( StorageConnectionError, + StorageError, StorageInvalidCredentialError, StorageMissingAuthenticationError, ) @@ -51,6 +54,12 @@ logger = logging.getLogger("tower.storage") +@dataclass(frozen=True, slots=True) +class _AccessCacheKey: + name: str + mode: str + + def _auth_from_context(context: TowerContext) -> tuple[str, str, str]: if context.jwt is not None: token = context.jwt @@ -69,7 +78,7 @@ def _auth_from_context(context: TowerContext) -> tuple[str, str, str]: if token.strip() == INVALID_CREDENTIAL_SENTINEL: raise StorageInvalidCredentialError( - f"{source} contains the {INVALID_CREDENTIAL_SENTINEL!r} placeholder, " + f"{source} contains the {INVALID_CREDENTIAL_SENTINEL} placeholder, " "not a usable Tower credential." ) if not token.strip(): @@ -104,6 +113,33 @@ def _auth_hash(token: str, auth_header_name: str, prefix: str) -> str: return hashlib.sha256(presented_auth.encode("utf-8")).hexdigest() +@dataclass(frozen=True) +class _ResolvedCatalogAccess: + target_environment: str + catalog_environment: str + catalog_name: str + catalog_uri: str + warehouse: str + mode: str + oauth_token: str = field(repr=False) + expires_at: datetime + + def is_inherited(self) -> bool: + return self.catalog_environment != self.target_environment + + def is_usable(self, now: datetime) -> bool: + return self.expires_at - now > CREDENTIAL_REFRESH_WINDOW + + def to_credentials(self) -> CatalogCredentials: + return CatalogCredentials( + catalog_uri=self.catalog_uri, + expires_at=self.expires_at, + mode=self.mode, + oauth_token=self.oauth_token, + warehouse=self.warehouse, + ) + + class _StorageResolver: """Private Tower configuration and authentication for catalog resolution.""" @@ -120,18 +156,20 @@ def __init__( if environment is not None and not environment.strip(): raise ValueError("environment must not be blank") - self._target_environment = ( + self._target_environment: str = ( environment or context.environment or DEFAULT_ENVIRONMENT_NAME ) - self._base_url = _api_base_url(context.tower_url) + self._base_url: str = _api_base_url(context.tower_url) + + self._token: str + self._auth_header_name: str + self._auth_prefix: str self._token, self._auth_header_name, self._auth_prefix = _auth_from_context( context ) - self._auth_hash = _auth_hash( - self._token, - self._auth_header_name, - self._auth_prefix, - ) + self._access_cache: dict[_AccessCacheKey, _ResolvedCatalogAccess] = {} + self._access_flights: dict[_AccessCacheKey, Future[_ResolvedCatalogAccess]] = {} + self._access_lock: Lock = Lock() def _new_client(self) -> AuthenticatedClient: return _new_tower_control_plane_client( @@ -142,6 +180,79 @@ def _new_client(self) -> AuthenticatedClient: timeout=DEFAULT_STORAGE_TIMEOUT_SECONDS, ) + def _resolve_catalog_access( + self, + name: str, + mode: str, + ) -> _ResolvedCatalogAccess: + mode = _normalize_mode(mode) + target_environment = self._target_environment + # Host, authentication, and target are fixed for this resolver. + cache_key = _AccessCacheKey(name=name, mode=mode) + + with self._access_lock: + cached = self._access_cache.get(cache_key) + if cached is not None and cached.is_usable(datetime.now(timezone.utc)): + return cached + _ = self._access_cache.pop(cache_key, None) + + flight = self._access_flights.get(cache_key) + if flight is None: + flight = Future() + self._access_flights[cache_key] = flight + should_vend = True + else: + should_vend = False + + if not should_vend: + return flight.result() + + try: + response = _vend_with_default_catalog_fallback( + self, + name, + mode, + ) + credentials = response.credentials + if credentials.mode != mode: + raise StorageError( + f"Tower returned {credentials.mode} credentials after " + f"{mode} access was requested." + ) + if response.environment not in ( + target_environment, + DEFAULT_ENVIRONMENT_NAME, + ): + raise StorageError( + f"Tower resolved catalog {name} from unexpected environment " + f"{response.environment}." + ) + + access = _ResolvedCatalogAccess( + target_environment=target_environment, + catalog_environment=response.environment, + catalog_name=name, + catalog_uri=credentials.catalog_uri, + warehouse=credentials.warehouse, + mode=credentials.mode, + oauth_token=credentials.oauth_token, + expires_at=_ensure_aware(credentials.expires_at), + ) + cacheable = access.is_usable(datetime.now(timezone.utc)) + except BaseException as error: + with self._access_lock: + flight.set_exception(error) + _ = self._access_flights.pop(cache_key, None) + raise + + with self._access_lock: + if cacheable: + self._access_cache[cache_key] = access + flight.set_result(access) + _ = self._access_flights.pop(cache_key, None) + + return access + def _request_catalog_credentials( self, name: str, @@ -173,13 +284,12 @@ def _api_base_url(tower_url: str) -> str: return str(url.copy_with(path="/v1", query=None, fragment=None)) -@dataclass -class _CachedCredentials: - credentials: CatalogCredentials - - def is_usable(self, now: datetime) -> bool: - expires_at = _ensure_aware(self.credentials.expires_at) - return now < expires_at - CREDENTIAL_REFRESH_WINDOW +@dataclass(frozen=True, slots=True) +class _CatalogTypeCacheKey: + base_url: str + auth_hash: str + name: str + environment: str @dataclass @@ -188,8 +298,7 @@ class _CachedCatalogType: retry_at: float | None = None -_credential_cache: dict[tuple[str, str, str, str, str], _CachedCredentials] = {} -_catalog_type_cache: dict[tuple[str, str, str, str], _CachedCatalogType] = {} +_catalog_type_cache: dict[_CatalogTypeCacheKey, _CachedCatalogType] = {} def get_tower_catalog( @@ -210,18 +319,8 @@ def get_tower_catalog_credentials( mode: str = "read", ) -> CatalogCredentials: storage_resolver = _StorageResolver(environment=environment) - mode = _normalize_mode(mode) - cache_key = _cache_key(storage_resolver, name, mode) - - now = datetime.now(timezone.utc) - _prune_credential_cache(now) - cached = _credential_cache.get(cache_key) - if cached is not None and cached.is_usable(now): - return cached.credentials - - credentials = _vend_with_default_catalog_fallback(storage_resolver, name, mode) - _credential_cache[cache_key] = _CachedCredentials(credentials) - return credentials + access = storage_resolver._resolve_catalog_access(name, mode) + return access.to_credentials() def load_vended_catalog(name: str, credentials: CatalogCredentials) -> Catalog: @@ -240,7 +339,7 @@ def _vend_with_default_catalog_fallback( storage_resolver: _StorageResolver, name: str, mode: str, -) -> CatalogCredentials: +) -> VendCatalogCredentialsResponse: environment = storage_resolver._target_environment result = storage_resolver._request_catalog_credentials(name, mode) if not _is_not_found(result): @@ -258,7 +357,7 @@ def _vend_with_default_catalog_fallback( return _unwrap_vend_result(result, name, environment) raise RuntimeError( - f"Tower catalog {name!r} does not exist in environment {environment!r}." + f"Tower catalog {name} does not exist in environment {environment}." ) @@ -271,7 +370,12 @@ def _describe_tower_catalog_type( token, auth_header_name, prefix = _auth_from_context(ctx) base_url = _api_base_url(ctx.tower_url) auth_hash = _auth_hash(token, auth_header_name, prefix) - cache_key = (base_url, auth_hash, name, environment) + cache_key = _CatalogTypeCacheKey( + base_url=base_url, + auth_hash=auth_hash, + name=name, + environment=environment, + ) cached = _catalog_type_cache.get(cache_key) if cached is not None: if cached.retry_at is None: @@ -346,45 +450,23 @@ def _unwrap_vend_result( result: ErrorModel | VendCatalogCredentialsResponse | None, name: str, environment: str, -) -> CatalogCredentials: +) -> VendCatalogCredentialsResponse: if isinstance(result, VendCatalogCredentialsResponse): - return result.credentials + return result if isinstance(result, ErrorModel): detail = _error_text(result) raise RuntimeError( - f"Failed to vend credentials for Tower catalog {name!r} " - f"in environment {environment!r}: {detail}" + f"Failed to vend credentials for Tower catalog {name} " + f"in environment {environment}: {detail}" ) raise RuntimeError( - f"Failed to vend credentials for Tower catalog {name!r} " - f"in environment {environment!r}." + f"Failed to vend credentials for Tower catalog {name} " + f"in environment {environment}." ) -def _cache_key( - storage_resolver: _StorageResolver, - name: str, - mode: str, -) -> tuple[str, str, str, str, str]: - return ( - storage_resolver._base_url, - storage_resolver._auth_hash, - name, - storage_resolver._target_environment, - mode, - ) - - -def _prune_credential_cache(now: datetime) -> None: - expired_keys = [ - key for key, cached in _credential_cache.items() if not cached.is_usable(now) - ] - for key in expired_keys: - _credential_cache.pop(key, None) - - def _normalize_mode(mode: str) -> str: if mode not in ("read", "read-write"): raise ValueError("mode must be 'read' or 'read-write'") @@ -416,6 +498,5 @@ def _ensure_aware(value: datetime) -> datetime: return value.astimezone(timezone.utc) -def _clear_credential_cache() -> None: - _credential_cache.clear() +def _clear_catalog_type_cache() -> None: _catalog_type_cache.clear() diff --git a/tests/tower/test_storage.py b/tests/tower/test_storage.py index 7f4307fd..2a0a136f 100644 --- a/tests/tower/test_storage.py +++ b/tests/tower/test_storage.py @@ -1,5 +1,7 @@ +from concurrent.futures import ThreadPoolExecutor from datetime import datetime, timedelta, timezone from http import HTTPStatus +from threading import Event import httpx import pytest @@ -8,6 +10,7 @@ from tower._context import TowerContext from tower.exceptions import ( StorageConnectionError, + StorageError, StorageInvalidCredentialError, StorageMissingAuthenticationError, ) @@ -20,6 +23,39 @@ ) +def make_vend_response( + environment, + token, + expires_at, + mode="read", +): + return VendCatalogCredentialsResponse( + credentials=CatalogCredentials( + catalog_uri="https://catalog.example.com", + expires_at=expires_at, + mode=mode, + oauth_token=token, + warehouse="warehouse-id", + ), + environment=environment, + ) + + +def script_vend(monkeypatch, results): + results = iter(results) + calls = [] + + def vend(client, name, mode): + calls.append((name, mode)) + result = next(results) + if isinstance(result, BaseException): + raise result + return result + + monkeypatch.setattr(_storage._StorageResolver, "_request_catalog_credentials", vend) + return calls + + @pytest.fixture(autouse=True) def isolate_tower_environment(monkeypatch, tmp_path): for name in ( @@ -115,13 +151,6 @@ def test_auth_hash_includes_how_the_credential_is_presented(): def test_missing_auth_fails_before_cache_or_vend(monkeypatch): - _storage._clear_credential_cache() - - monkeypatch.setattr( - _storage, - "_prune_credential_cache", - lambda now: pytest.fail("cache access must not run without authentication"), - ) monkeypatch.setattr( _storage.vend_catalog_credentials_api, "sync", @@ -152,9 +181,8 @@ def vend(**kwargs): monkeypatch.setattr(_storage.vend_catalog_credentials_api, "sync", vend) - result = _storage._StorageResolver()._request_catalog_credentials( - "analytics", "read" - ) + resolver = _storage._StorageResolver() + result = resolver._request_catalog_credentials("analytics", "read") assert result is rejected assert len(vend_calls) == 1 @@ -181,9 +209,8 @@ def vend(*, name, client, environment, body): monkeypatch.setattr(_storage.vend_catalog_credentials_api, "sync", vend) - result = _storage._StorageResolver( - environment="production" - )._request_catalog_credentials("analytics", "read") + resolver = _storage._StorageResolver(environment="production") + result = resolver._request_catalog_credentials("analytics", "read") assert result is response assert captured == { @@ -198,7 +225,7 @@ def vend(*, name, client, environment, body): def test_describe_and_vend_prefer_jwt_when_both_auth_vars_are_set(monkeypatch): - _storage._clear_credential_cache() + _storage._clear_catalog_type_cache() monkeypatch.setenv("TOWER_URL", "https://api.example.com") monkeypatch.setenv("TOWER_ENVIRONMENT", "production") monkeypatch.setenv("TOWER_API_KEY", "ambient-api-key") @@ -241,10 +268,8 @@ def vend(*, client, **kwargs): _storage._describe_tower_catalog_type(ctx, "analytics", "production") == _storage.TOWER_CATALOG_TYPE ) - assert ( - _storage._StorageResolver()._request_catalog_credentials("analytics", "read") - is vended - ) + resolver = _storage._StorageResolver() + assert resolver._request_catalog_credentials("analytics", "read") is vended assert captured_auth == [ ("describe", "Bearer ambient-jwt", None), ("vend", "Bearer ambient-jwt", None), @@ -261,7 +286,6 @@ def test_storage_resolver_allows_explicit_http_tower_url(monkeypatch): def test_get_tower_catalog_credentials_allows_http_and_reaches_vend(monkeypatch): - _storage._clear_credential_cache() monkeypatch.setenv("TOWER_URL", "http://localhost:9000") monkeypatch.setenv("TOWER_ENVIRONMENT", "production") monkeypatch.setenv("TOWER_API_KEY", "api-key") @@ -287,7 +311,7 @@ def vend(*, name, client, environment, body): result = _storage.get_tower_catalog_credentials("analytics") - assert result is credentials + assert result == credentials assert vend_calls == [ ( "analytics", @@ -316,11 +340,6 @@ def test_get_tower_catalog_credentials_rejects_invalid_mode_before_cache_or_vend monkeypatch, ): monkeypatch.setenv("TOWER_API_KEY", "api-key") - monkeypatch.setattr( - _storage, - "_prune_credential_cache", - lambda now: pytest.fail("cache access must not run for an invalid mode"), - ) monkeypatch.setattr( _storage.vend_catalog_credentials_api, "sync", @@ -350,7 +369,8 @@ def vend(*, client, **kwargs): monkeypatch.setattr(_storage.vend_catalog_credentials_api, "sync", vend) with pytest.raises(StorageConnectionError) as error: - _storage._StorageResolver()._request_catalog_credentials("analytics", "read") + resolver = _storage._StorageResolver() + resolver._request_catalog_credentials("analytics", "read") assert error.value.__cause__ is cause assert clients[0]._client is not None @@ -369,79 +389,203 @@ def test_storage_resolver_uses_runtime_environment_only_as_target_config(monkeyp ) -def test_get_tower_catalog_credentials_caches_vended_credentials(monkeypatch): - _storage._clear_credential_cache() +def test_access_cache_is_owned_by_resolver(monkeypatch): monkeypatch.setenv("TOWER_URL", "https://api.example.com") monkeypatch.setenv("TOWER_ENVIRONMENT", "production") monkeypatch.setenv("TOWER_API_KEY", "api-key") - expires_at = datetime.now(timezone.utc) + timedelta(hours=1) - credentials = CatalogCredentials( - catalog_uri="https://catalog.example.com", - expires_at=expires_at, - mode="read", - oauth_token="oauth-token", - warehouse="warehouse-id", + first_resolver = _storage._StorageResolver() + second_resolver = _storage._StorageResolver() + calls = script_vend( + monkeypatch, + [ + make_vend_response( + environment="production", + token=f"provider-token-{number}", + expires_at=datetime.now(timezone.utc) + timedelta(hours=1), + ) + for number in (1, 2) + ], ) - calls = [] - def vend(client, name, mode): - calls.append((name, client._target_environment, mode)) - return VendCatalogCredentialsResponse( - credentials=credentials, - environment=client._target_environment, - ) + first = first_resolver._resolve_catalog_access("analytics", "read") + first_again = first_resolver._resolve_catalog_access("analytics", "read") + second = second_resolver._resolve_catalog_access("analytics", "read") + second_again = second_resolver._resolve_catalog_access("analytics", "read") + + assert first == first_again + assert second == second_again + assert {first.oauth_token, second.oauth_token} == { + "provider-token-1", + "provider-token-2", + } + assert len(calls) == 2 - monkeypatch.setattr(_storage._StorageResolver, "_request_catalog_credentials", vend) - first = _storage.get_tower_catalog_credentials("default") - second = _storage.get_tower_catalog_credentials("default") +def test_one_shot_credential_loads_do_not_share_cache(monkeypatch): + monkeypatch.setenv("TOWER_API_KEY", "api-key") + calls = script_vend( + monkeypatch, + [ + make_vend_response( + environment="default", + token=f"provider-token-{number}", + expires_at=datetime.now(timezone.utc) + timedelta(hours=1), + ) + for number in (1, 2) + ], + ) - assert first is credentials - assert second is credentials - assert calls == [("default", "production", "read")] + first = _storage.get_tower_catalog_credentials("analytics") + second = _storage.get_tower_catalog_credentials("analytics") + assert first.oauth_token == "provider-token-1" + assert second.oauth_token == "provider-token-2" + assert len(calls) == 2 -def test_get_tower_catalog_credentials_prunes_expired_cache_entries(monkeypatch): - _storage._clear_credential_cache() - monkeypatch.setenv("TOWER_URL", "https://api.example.com") + +def test_near_expiry_access_retains_inherited_and_local_identity(monkeypatch): monkeypatch.setenv("TOWER_ENVIRONMENT", "production") monkeypatch.setenv("TOWER_API_KEY", "api-key") - expired_credentials = CatalogCredentials( - catalog_uri="https://old-catalog.example.com", - expires_at=datetime.now(timezone.utc) - timedelta(minutes=1), - mode="read", - oauth_token="old-oauth-token", - warehouse="old-warehouse-id", - ) - fresh_credentials = CatalogCredentials( - catalog_uri="https://catalog.example.com", - expires_at=datetime.now(timezone.utc) + timedelta(hours=1), - mode="read", - oauth_token="oauth-token", - warehouse="warehouse-id", - ) - storage_resolver = _storage._StorageResolver() - expired_key = _storage._cache_key(storage_resolver, "stale", "read") - _storage._credential_cache[expired_key] = _storage._CachedCredentials( - expired_credentials - ) + responses = [ + make_vend_response( + environment="default", + token="inherited-token", + expires_at=datetime.now(timezone.utc) + + _storage.CREDENTIAL_REFRESH_WINDOW + - timedelta(seconds=1), + ), + make_vend_response( + environment="production", + token="local-token", + expires_at=datetime.now(timezone.utc) + timedelta(hours=1), + ), + ] + calls = script_vend(monkeypatch, responses) + + resolver = _storage._StorageResolver() + inherited = resolver._resolve_catalog_access("analytics", "read") + local = resolver._resolve_catalog_access("analytics", "read") + resolver._resolve_catalog_access("analytics", "read") + + assert inherited.target_environment == "production" + assert inherited.catalog_environment == "default" + assert inherited.is_inherited() is True + assert inherited.oauth_token == "inherited-token" + assert local.catalog_environment == "production" + assert local.is_inherited() is False + assert local.oauth_token == "local-token" + assert len(calls) == 2 + + +def test_concurrent_access_shares_one_vend_request(monkeypatch): + monkeypatch.setenv("TOWER_API_KEY", "api-key") + timeout = 10 + vend_started = Event() + waiter_joined = Event() + release_vend = Event() + calls = [] + + class ObservableFuture(_storage.Future): + def result(self, timeout=None): + waiter_joined.set() + return super().result(timeout) def vend(client, name, mode): - return VendCatalogCredentialsResponse( - credentials=fresh_credentials, - environment=client._target_environment, + calls.append((name, mode)) + vend_started.set() + assert release_vend.wait(timeout=timeout) + return make_vend_response( + environment="default", + token="shared-token", + expires_at=datetime.now(timezone.utc) + timedelta(hours=1), ) + def resolve(resolver): + return resolver._resolve_catalog_access("analytics", "read") + + monkeypatch.setattr(_storage, "Future", ObservableFuture) monkeypatch.setattr(_storage._StorageResolver, "_request_catalog_credentials", vend) - result = _storage.get_tower_catalog_credentials("default") + resolver = _storage._StorageResolver() + with ThreadPoolExecutor(max_workers=2) as executor: + leader = executor.submit(resolve, resolver) + started = vend_started.wait(timeout=timeout) + waiter = executor.submit(resolve, resolver) + joined = waiter_joined.wait(timeout=timeout) + release_vend.set() + accesses = [ + leader.result(timeout=timeout), + waiter.result(timeout=timeout), + ] + + assert started + assert joined + assert len(calls) == 1 + assert all(access == accesses[0] for access in accesses) + + +def test_failed_vend_is_not_cached(monkeypatch): + monkeypatch.setenv("TOWER_API_KEY", "api-key") + calls = script_vend( + monkeypatch, + [ + RuntimeError("vend failed"), + make_vend_response( + environment="default", + token="retry-token", + expires_at=datetime.now(timezone.utc) + timedelta(hours=1), + ), + ], + ) + + resolver = _storage._StorageResolver() + with pytest.raises(RuntimeError, match="vend failed"): + resolver._resolve_catalog_access("analytics", "read") + access = resolver._resolve_catalog_access("analytics", "read") + + assert access.oauth_token == "retry-token" + assert len(calls) == 2 + + +@pytest.mark.parametrize( + ("environment", "mode", "message"), + [ + ("staging", "read", "unexpected environment"), + ("default", "read-write", "after read access was requested"), + ], +) +def test_inconsistent_vend_response_is_not_cached( + monkeypatch, + environment, + mode, + message, +): + monkeypatch.setenv("TOWER_API_KEY", "api-key") + responses = [ + make_vend_response( + environment=environment, + token="invalid-token", + expires_at=datetime.now(timezone.utc) + timedelta(hours=1), + mode=mode, + ), + make_vend_response( + environment="default", + token="valid-token", + expires_at=datetime.now(timezone.utc) + timedelta(hours=1), + ), + ] + calls = script_vend(monkeypatch, responses) + + resolver = _storage._StorageResolver() + with pytest.raises(StorageError, match=message): + resolver._resolve_catalog_access("analytics", "read") + access = resolver._resolve_catalog_access("analytics", "read") - assert result is fresh_credentials - assert expired_key not in _storage._credential_cache + assert access.oauth_token == "valid-token" + assert len(calls) == 2 def test_default_catalog_retries_reuse_client_auth_snapshot(monkeypatch): - _storage._clear_credential_cache() monkeypatch.setenv("TOWER_API_KEY", "operation-token") credentials = CatalogCredentials( catalog_uri="https://catalog.example.com", @@ -483,7 +627,7 @@ def legacy_default(*, client): result = _storage.get_tower_catalog_credentials("default") - assert result is credentials + assert result == credentials assert vend_tokens == ["operation-token"] * 3 assert legacy_tokens == ["operation-token"] * 2 assert len({id(client) for client in clients}) == 5 @@ -494,7 +638,7 @@ def legacy_default(*, client): def test_describe_tower_catalog_type_uses_timeout_and_recovers_after_cooldown( monkeypatch, ): - _storage._clear_credential_cache() + _storage._clear_catalog_type_cache() ctx = TowerContext( tower_url="https://api.example.com", environment="production", diff --git a/tests/tower/test_tables.py b/tests/tower/test_tables.py index 4bec5997..acd428f7 100644 --- a/tests/tower/test_tables.py +++ b/tests/tower/test_tables.py @@ -155,7 +155,7 @@ def sql_catalog(): def test_string_catalog_precedence( monkeypatch, tower_credentials, catalog_type, has_pyiceberg_config, expected_source ): - _storage._clear_credential_cache() + _storage._clear_catalog_type_cache() patch_tower_context(monkeypatch) vend_catalog = FakeCatalog("vend") configured_catalog = FakeCatalog("configured") @@ -271,7 +271,7 @@ def unexpected_call(*args, **kwargs): def test_no_tower_auth_preserves_ambient_pyiceberg_catalog(monkeypatch): - _storage._clear_credential_cache() + _storage._clear_catalog_type_cache() patch_tower_context(monkeypatch, api_key=None) monkeypatch.setenv( "PYICEBERG_CATALOG__S3_TABLES__URI", "https://s3tables.example.com" @@ -298,7 +298,7 @@ def load_catalog(name): def test_managed_catalog_vend_failure_does_not_fall_back_to_pyiceberg(monkeypatch): - _storage._clear_credential_cache() + _storage._clear_catalog_type_cache() patch_tower_context(monkeypatch) calls = [] @@ -334,7 +334,7 @@ def load_catalog(name): def test_external_catalog_write_mode_keeps_ambient_pyiceberg_catalog( monkeypatch, catalog_type ): - _storage._clear_credential_cache() + _storage._clear_catalog_type_cache() patch_tower_context(monkeypatch) catalog = FakeCatalog("configured") @@ -357,7 +357,7 @@ def unexpected_vend(*args, **kwargs): def test_string_catalog_type_describe_is_cached(monkeypatch): - _storage._clear_credential_cache() + _storage._clear_catalog_type_cache() patch_tower_context(monkeypatch) calls = [] vend_catalogs = [] From e034b6c01c7e81d6db9b97983364e98d7b08adbc Mon Sep 17 00:00:00 2001 From: Ben Lovell Date: Fri, 28 Aug 2026 10:37:20 +0200 Subject: [PATCH 3/5] fix: help text tests no longer depend on terminal width (#338) --- crates/tower-cmd/src/catalogs.rs | 10 ++++++---- crates/tower-cmd/src/lib.rs | 4 +++- 2 files changed, 9 insertions(+), 5 deletions(-) diff --git a/crates/tower-cmd/src/catalogs.rs b/crates/tower-cmd/src/catalogs.rs index f829069f..c53b30c9 100644 --- a/crates/tower-cmd/src/catalogs.rs +++ b/crates/tower-cmd/src/catalogs.rs @@ -1854,8 +1854,10 @@ mod tests { #[test] fn catalog_help_marks_storage_beta_in_short_and_long_help() { - let short_help = catalogs_cmd().render_help().to_string(); - let long_help = catalogs_cmd().render_long_help().to_string(); + // term_width(0) turns off help wrapping, which otherwise follows the + // width of the terminal the tests run in. + let short_help = catalogs_cmd().term_width(0).render_help().to_string(); + let long_help = catalogs_cmd().term_width(0).render_long_help().to_string(); for help in [short_help, long_help] { assert!(help.contains("includes Storage [beta]")); @@ -1865,7 +1867,7 @@ mod tests { #[test] fn storage_specific_command_and_flag_are_marked_beta() { - let mut command = catalogs_cmd(); + let mut command = catalogs_cmd().mut_subcommand("credentials", |c| c.term_width(0)); let credentials_help = command .find_subcommand_mut("credentials") .expect("credentials command should exist") @@ -1874,7 +1876,7 @@ mod tests { assert!(credentials_help .contains("Vend short-lived catalog credentials for external tools [beta]")); - let mut command = catalogs_cmd(); + let mut command = catalogs_cmd().mut_subcommand("list", |c| c.term_width(0)); let list_help = command .find_subcommand_mut("list") .expect("list command should exist") diff --git a/crates/tower-cmd/src/lib.rs b/crates/tower-cmd/src/lib.rs index 02cd8803..0bb92756 100644 --- a/crates/tower-cmd/src/lib.rs +++ b/crates/tower-cmd/src/lib.rs @@ -351,7 +351,9 @@ mod tests { #[test] fn root_help_scopes_beta_label_to_storage() { - let help = root_cmd().render_help().to_string(); + // term_width(0) turns off help wrapping, which otherwise follows the + // width of the terminal the tests run in. + let help = root_cmd().term_width(0).render_help().to_string(); assert!(help.contains( "Interact with the catalogs in your Tower account (includes Storage [beta])" From 4c0b7c5b0555e225d537afbc329a45b424c55627 Mon Sep 17 00:00:00 2001 From: Konstantinos St Date: Tue, 8 Sep 2026 16:31:43 +0200 Subject: [PATCH 4/5] feat: add request fields to ResourceLimits (#366) --- crates/tower-cmd/src/run.rs | 9 ++++++--- crates/tower-runtime/src/execution.rs | 18 ++++++++++++++---- crates/tower-runtime/tests/subprocess_test.rs | 9 ++++++--- 3 files changed, 26 insertions(+), 10 deletions(-) diff --git a/crates/tower-cmd/src/run.rs b/crates/tower-cmd/src/run.rs index 514df320..ad7e8d1b 100644 --- a/crates/tower-cmd/src/run.rs +++ b/crates/tower-cmd/src/run.rs @@ -322,9 +322,12 @@ fn build_cli_execution_spec( parameters: params, env_vars, resources: ResourceLimits { - cpu_millicores: None, - memory_mb: None, - storage_mb: None, + cpu_limit_millicores: None, + cpu_request_millicores: None, + memory_limit_mb: None, + memory_request_mb: None, + storage_limit_mb: None, + storage_request_mb: None, max_pids: None, gpu_count: 0, timeout_seconds: 3600, diff --git a/crates/tower-runtime/src/execution.rs b/crates/tower-runtime/src/execution.rs index f69e51f9..387b54f9 100644 --- a/crates/tower-runtime/src/execution.rs +++ b/crates/tower-runtime/src/execution.rs @@ -114,17 +114,27 @@ pub enum CacheBackend { None, } -/// ResourceLimits defines compute resource constraints +/// ResourceLimits defines compute resource requests and limits. +/// Missing values are resolved by the execution backend. #[derive(Debug, Clone)] pub struct ResourceLimits { /// CPU limit in millicores (e.g., 1000 = 1 CPU) - pub cpu_millicores: Option, + pub cpu_limit_millicores: Option, + + /// CPU request in millicores + pub cpu_request_millicores: Option, /// Memory limit in megabytes - pub memory_mb: Option, + pub memory_limit_mb: Option, + + /// Memory request in megabytes + pub memory_request_mb: Option, /// Ephemeral storage limit in megabytes - pub storage_mb: Option, + pub storage_limit_mb: Option, + + /// Ephemeral storage request in megabytes + pub storage_request_mb: Option, /// Maximum number of processes pub max_pids: Option, diff --git a/crates/tower-runtime/tests/subprocess_test.rs b/crates/tower-runtime/tests/subprocess_test.rs index 8d7448c6..7c80cbe3 100644 --- a/crates/tower-runtime/tests/subprocess_test.rs +++ b/crates/tower-runtime/tests/subprocess_test.rs @@ -65,9 +65,12 @@ async fn create_execution_spec(id: String, package: Package) -> ExecutionSpec { command: None, }, resources: ResourceLimits { - cpu_millicores: None, - memory_mb: None, - storage_mb: None, + cpu_limit_millicores: None, + cpu_request_millicores: None, + memory_limit_mb: None, + memory_request_mb: None, + storage_limit_mb: None, + storage_request_mb: None, max_pids: None, gpu_count: 0, timeout_seconds: 300, From 9ae9de1c2c26784c55e0547f27b91f8396409a93 Mon Sep 17 00:00:00 2001 From: Konstantinos Stefanidis Vozikis Date: Tue, 8 Sep 2026 16:39:01 +0200 Subject: [PATCH 5/5] bump version --- Cargo.lock | 24 ++++++++++++------------ Cargo.toml | 2 +- pyproject.toml | 2 +- uv.lock | 2 +- 4 files changed, 15 insertions(+), 15 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index d81ddbf5..fa58cd6c 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -841,7 +841,7 @@ dependencies = [ [[package]] name = "config" -version = "0.3.72" +version = "0.3.73" dependencies = [ "base64", "chrono", @@ -1033,7 +1033,7 @@ checksum = "460fbee9c2c2f33933d720630a6a0bac33ba7053db5344fac858d4b8952d77d5" [[package]] name = "crypto" -version = "0.3.72" +version = "0.3.73" dependencies = [ "aes-gcm", "base64", @@ -4456,7 +4456,7 @@ dependencies = [ [[package]] name = "testutils" -version = "0.3.72" +version = "0.3.73" dependencies = [ "pem", "rsa", @@ -4765,7 +4765,7 @@ checksum = "5d99f8c9a7727884afe522e9bd5edbfc91a3312b36a77b5fb8926e4c31a41801" [[package]] name = "tower" -version = "0.3.72" +version = "0.3.73" dependencies = [ "config", "pyo3", @@ -4793,7 +4793,7 @@ dependencies = [ [[package]] name = "tower-api" -version = "0.3.72" +version = "0.3.73" dependencies = [ "reqwest", "serde", @@ -4805,7 +4805,7 @@ dependencies = [ [[package]] name = "tower-cmd" -version = "0.3.72" +version = "0.3.73" dependencies = [ "axum", "bytes", @@ -4854,7 +4854,7 @@ dependencies = [ [[package]] name = "tower-duckdb" -version = "0.3.72" +version = "0.3.73" dependencies = [ "chrono", "duckdb", @@ -4889,7 +4889,7 @@ checksum = "121c2a6cda46980bb0fcd1647ffaf6cd3fc79a013de288782836f6df9c48780e" [[package]] name = "tower-package" -version = "0.3.72" +version = "0.3.73" dependencies = [ "async-compression", "flate2", @@ -4914,7 +4914,7 @@ dependencies = [ [[package]] name = "tower-runtime" -version = "0.3.72" +version = "0.3.73" dependencies = [ "async-trait", "chrono", @@ -4938,7 +4938,7 @@ checksum = "8df9b6e13f2d32c91b9bd719c00d1958837bc7dec474d94952798cc8e69eeec3" [[package]] name = "tower-telemetry" -version = "0.3.72" +version = "0.3.73" dependencies = [ "tracing", "tracing-appender", @@ -4947,7 +4947,7 @@ dependencies = [ [[package]] name = "tower-uv" -version = "0.3.72" +version = "0.3.73" dependencies = [ "async-compression", "async_zip", @@ -4967,7 +4967,7 @@ dependencies = [ [[package]] name = "tower-version" -version = "0.3.72" +version = "0.3.73" dependencies = [ "anyhow", "chrono", diff --git a/Cargo.toml b/Cargo.toml index 70db28b6..428fa352 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -4,7 +4,7 @@ resolver = "2" [workspace.package] edition = "2021" -version = "0.3.72" +version = "0.3.73" description = "Tower is the best way to host Python data apps in production" # Matches rust-toolchain.toml. The two had drifted: the toolchain has been 1.88 # for a while, and the dependency tree (testcontainers and its transitive deps, diff --git a/pyproject.toml b/pyproject.toml index ae73132f..0bd20ea5 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -4,7 +4,7 @@ build-backend = "maturin" [project] name = "tower" -version = "0.3.72" +version = "0.3.73" description = "Tower CLI and runtime environment for Tower." authors = [{ name = "Tower Computing GmbH", email = "brad@tower.dev" }] readme = "README.md" diff --git a/uv.lock b/uv.lock index e88fc41c..083e5533 100644 --- a/uv.lock +++ b/uv.lock @@ -2264,7 +2264,7 @@ wheels = [ [[package]] name = "tower" -version = "0.3.72" +version = "0.3.73" source = { editable = "." } dependencies = [ { name = "attrs" },