From 0ea5fa410834b7dc6f9e2828d356287fff3317d9 Mon Sep 17 00:00:00 2001 From: Konstantinos Stefanidis Vozikis Date: Wed, 19 Aug 2026 12:59:53 +0200 Subject: [PATCH 1/3] feat: Create dedicated control plane client for Storage --- src/tower/_storage.py | 231 +++++++++++++++---- src/tower/exceptions.py | 20 ++ tests/tower/test_storage.py | 427 ++++++++++++++++++++++++++++++------ 3 files changed, 565 insertions(+), 113 deletions(-) diff --git a/src/tower/_storage.py b/src/tower/_storage.py index bff78bc1..8251636d 100644 --- a/src/tower/_storage.py +++ b/src/tower/_storage.py @@ -5,11 +5,17 @@ import time from dataclasses import dataclass from datetime import datetime, timedelta, timezone -from http import HTTPStatus from typing import Any, Optional -from ._client import _env_client +import httpx + 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 +39,150 @@ # 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 _build_tower_control_plane_client( + *, + context: TowerContext, + tower_url: str, + timeout: float, + verify_tls: bool, +) -> tuple[str, str, AuthenticatedClient]: + token, auth_header_name, prefix = _auth_from_context(context) + base_url = _api_base_url(tower_url) + # hash whatever auth token was provided so we can cache catalogs tokens per account. + auth_hash = hashlib.sha256(token.encode("utf-8")).hexdigest() + client = AuthenticatedClient( + verify_ssl=verify_tls, + base_url=base_url, + token=token, + auth_header_name=auth_header_name, + prefix=prefix, + timeout=httpx.Timeout(timeout), + raise_on_unexpected_status=True, + ) + return base_url, auth_hash, client + + +class StorageClient: + """Configuration and authenticated Tower Storage-specific client. + + This remains an internal foundation until the public catalog-loading surface is + added. It intentionally does not alter the clients used by unrelated SDK calls. + """ + + def __init__( + self, + *, + tower_url: str | None = None, + environment: str | None = None, + timeout: float = DEFAULT_STORAGE_TIMEOUT_SECONDS, + verify_tls: bool = True, + ) -> None: + context = TowerContext.build() + + if tower_url is not None and not isinstance(tower_url, str): + raise TypeError("tower_url must be a string or None") + + if tower_url is not None and not tower_url.strip(): + raise ValueError("tower_url must not be blank") + + 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") + + if isinstance(timeout, bool) or not isinstance(timeout, (int, float)): + raise TypeError("timeout must be a positive number") + + if not 0 < float(timeout) < float("inf"): + raise ValueError("timeout must be a positive finite number") + + if not isinstance(verify_tls, bool): + raise TypeError("verify_tls must be a bool") + + self.tower_url = tower_url or context.tower_url + self.environment = ( + environment or context.environment or DEFAULT_ENVIRONMENT_NAME + ) + self.timeout = float(timeout) + self.verify_tls = verify_tls + self._base_url, self._auth_hash, self._tower_client = ( + _build_tower_control_plane_client( + context=context, + tower_url=self.tower_url, + timeout=self.timeout, + verify_tls=self.verify_tls, + ) + ) + + def _request_catalog_credentials( + self, + name: str, + mode: str, + ) -> ErrorModel | VendCatalogCredentialsResponse | None: + body = VendCatalogCredentialsBody(mode=_vend_mode(mode)) + + try: + return vend_catalog_credentials_api.sync( + name=name, + client=self._tower_client, + environment=self.environment, + body=body, + ) + except httpx.RequestError as error: + raise StorageConnectionError( + f"Could not connect to Tower at {self._base_url!r}." + ) 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!r}") from error + if not url.is_absolute_url or url.scheme not in ("http", "https"): + raise ValueError(f"Invalid Tower URL: {tower_url!r}") + return str(url.copy_with(path="/v1", query=None, fragment=None)).rstrip("/") + + @dataclass class _CachedCredentials: credentials: CatalogCredentials @@ -56,7 +199,7 @@ 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( @@ -76,10 +219,9 @@ def get_tower_catalog_credentials( environment: Optional[str] = None, mode: str = "read", ) -> CatalogCredentials: - ctx = TowerContext.build() - environment = environment or ctx.environment or DEFAULT_ENVIRONMENT_NAME + storage_client = StorageClient(environment=environment) mode = _normalize_mode(mode) - cache_key = _cache_key(ctx, name, environment, mode) + cache_key = _cache_key(storage_client, name, mode) now = datetime.now(timezone.utc) _prune_credential_cache(now) @@ -87,7 +229,8 @@ 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) + with storage_client._tower_client: + credentials = _vend_with_default_catalog_fallback(storage_client, name, mode) _credential_cache[cache_key] = _CachedCredentials(credentials) return credentials @@ -105,20 +248,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_client: StorageClient, + name: str, + mode: str, ) -> CatalogCredentials: - result = _vend_catalog_credentials(ctx, name, environment, mode) + environment = storage_client.environment + result = storage_client._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_client) for delay in DEFAULT_CATALOG_PROVISION_RETRY_DELAYS: time.sleep(delay) - result = _vend_catalog_credentials(ctx, name, environment, mode) + result = storage_client._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_client) return _unwrap_vend_result(result, name, environment) @@ -127,26 +273,19 @@ 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) + base_url, auth_hash, tower_client = _build_tower_control_plane_client( + context=ctx, + tower_url=ctx.tower_url, + timeout=CATALOG_TYPE_DESCRIBE_TIMEOUT_SECONDS, + verify_tls=True, + ) + 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: @@ -158,11 +297,12 @@ def _describe_tower_catalog_type( _catalog_type_cache.pop(cache_key, None) 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 +339,9 @@ def _failed_catalog_type_cache_entry() -> _CachedCatalogType: ) -def _ensure_legacy_default_catalog(ctx: TowerContext) -> None: +def _ensure_legacy_default_catalog(storage_client: StorageClient) -> None: try: - response = describe_default_catalog_api.sync_detailed(client=_env_client(ctx)) - if response.status_code not in (HTTPStatus.OK, HTTPStatus.ACCEPTED): - return + describe_default_catalog_api.sync(client=storage_client._tower_client) except Exception: # The following vend retry will surface the actionable backend/auth error. return @@ -230,19 +368,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_client: StorageClient, + 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_client._base_url, + storage_client._auth_hash, + name, + storage_client.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..7877ae24 100644 --- a/tests/tower/test_storage.py +++ b/tests/tower/test_storage.py @@ -1,7 +1,17 @@ +import hashlib 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 +21,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 +32,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 +44,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 +56,323 @@ 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_client_configuration_and_tls_defaults(monkeypatch): + monkeypatch.setenv("TOWER_API_KEY", "ambient-key") + client = _storage.StorageClient( + tower_url="https://tower.example.com/", + environment="production", + timeout=12.5, + ) + transport = client._tower_client + + assert client.tower_url == "https://tower.example.com/" + assert client.environment == "production" + assert client.timeout == 12.5 + assert client.verify_tls is True + assert client._base_url == "https://tower.example.com/v1" + assert client._auth_hash == hashlib.sha256(b"ambient-key").hexdigest() + assert transport._base_url == "https://tower.example.com/v1" + assert transport._timeout == httpx.Timeout(12.5) + assert transport._verify_ssl is True + + http_client = transport.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_client_static_auth_precedence( + monkeypatch, + ambient_auth, + expected_header, + expected_value, +): + for name, value in ambient_auth.items(): + monkeypatch.setenv(name, value) + transport = _storage.StorageClient()._tower_client + http_client = transport.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" ) - else: - raise AssertionError("expected missing-auth error") - - _storage._ensure_tower_auth( - TowerContext( - tower_url="https://api.example.com", - environment="production", - api_key="api-key", + assert other_header not in http_client.headers + finally: + http_client.close() + + +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_client_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.StorageClient() + + +@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.StorageClient()._request_catalog_credentials("analytics", "read") + + assert result is rejected + assert len(vend_calls) == 1 + + +def test_storage_client_vends_with_ambient_api_key(monkeypatch): + monkeypatch.setenv("TOWER_API_KEY", "service-account-key") + captured = {} + response = ErrorModel(status=418, detail="captured") + + def vend(*, name, client, environment, body): + http_client = client.get_httpx_client() + try: + captured.update( + name=name, + environment=environment, + api_key=http_client.headers.get("X-API-Key"), + authorization=http_client.headers.get("Authorization"), + mode=body.mode, + ) + finally: + http_client.close() + return response + + monkeypatch.setattr(_storage.vend_catalog_credentials_api, "sync", vend) + + result = _storage.StorageClient( + tower_url="https://api.example.com", + 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, + } + + +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() + try: + captured_auth.append( + ( + operation, + http_client.headers.get("Authorization"), + http_client.headers.get("X-API-Key"), + ) + ) + finally: + http_client.close() + + 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.StorageClient()._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), + ] + + +@pytest.mark.parametrize("verify_tls", [True, False], ids=("verified", "unverified")) +def test_storage_client_allows_http_independently_of_tls_verification( + monkeypatch, verify_tls +): + monkeypatch.setenv("TOWER_API_KEY", "key") + transport = _storage.StorageClient( + tower_url="http://localhost:9000", + verify_tls=verify_tls, + )._tower_client + + assert transport._base_url == "http://localhost:9000/v1" + assert transport._verify_ssl is verify_tls + + +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_client_normalizes_and_validates_tower_api_url(monkeypatch): + monkeypatch.setenv("TOWER_API_KEY", "key") + client = _storage.StorageClient( + tower_url="https://TOWER.example.com:443/old/path?debug=true#fragment", + ) + assert client._tower_client._base_url == "https://tower.example.com/v1" + + with pytest.raises(ValueError, match="Invalid Tower URL"): + _storage.StorageClient(tower_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") + + +@pytest.mark.parametrize("field", ["tower_url", "environment"]) +def test_storage_client_rejects_non_string_text_configuration(field): + with pytest.raises(TypeError, match=f"{field} must be a string or None"): + _storage.StorageClient(**{field: 123}) + + +def test_storage_client_maps_httpx_connection_errors(monkeypatch): + monkeypatch.setenv("TOWER_API_KEY", "key") + cause = httpx.ConnectError("connection refused") + + def vend(**kwargs): + raise cause + + monkeypatch.setattr(_storage.vend_catalog_credentials_api, "sync", vend) + + with pytest.raises(StorageConnectionError) as error: + _storage.StorageClient()._request_catalog_credentials("analytics", "read") + + assert error.value.__cause__ is cause + + +def test_storage_client_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.StorageClient().environment == "run-env" + assert ( + _storage.StorageClient(environment="explicit-env").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 +383,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.environment, mode)) return VendCatalogCredentialsResponse( credentials=credentials, - environment=environment, + environment=client.environment, ) - monkeypatch.setattr(_storage.TowerContext, "build", staticmethod(lambda: ctx)) - monkeypatch.setattr(_storage, "_vend_catalog_credentials", vend) + monkeypatch.setattr(_storage.StorageClient, "_request_catalog_credentials", vend) first = _storage.get_tower_catalog_credentials("default") second = _storage.get_tower_catalog_credentials("default") @@ -113,11 +402,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 +419,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_client = _storage.StorageClient() + expired_key = _storage._cache_key(storage_client, "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.environment, ) - monkeypatch.setattr(_storage.TowerContext, "build", staticmethod(lambda: ctx)) - monkeypatch.setattr(_storage, "_vend_catalog_credentials", vend) + monkeypatch.setattr(_storage.StorageClient, "_request_catalog_credentials", vend) result = _storage.get_tower_catalog_credentials("default") @@ -152,17 +439,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 +457,32 @@ def test_default_catalog_vend_retries_after_legacy_provisioning(monkeypatch): environment="default", ), ] - legacy_calls = [] - - def vend(ctx, name, environment, mode): + vend_tokens = [] + legacy_tokens = [] + transports = [] + + def vend(*, client, **kwargs): + transports.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): + transports.append(client._tower_client) + legacy_tokens.append(client._tower_client.token) - monkeypatch.setattr(_storage.TowerContext, "build", staticmethod(lambda: ctx)) - monkeypatch.setattr(_storage, "_vend_catalog_credentials", vend) + monkeypatch.setattr(_storage.vend_catalog_credentials_api, "sync", vend) monkeypatch.setattr(_storage, "_ensure_legacy_default_catalog", 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 all(transport is transports[0] for transport in transports) + assert transports[0]._client is not None + assert transports[0]._client.is_closed def test_describe_tower_catalog_type_uses_timeout_and_recovers_after_cooldown( @@ -205,8 +496,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 +525,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 +543,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 567a9564d52c110d6537988f6c5c992602bf3868 Mon Sep 17 00:00:00 2001 From: Konstantinos Stefanidis Vozikis Date: Thu, 20 Aug 2026 10:43:12 +0200 Subject: [PATCH 2/3] refactor: keep Storage resolution private and request scoped --- src/tower/_storage.py | 140 +++++++++++++-------------- tests/tower/test_storage.py | 184 +++++++++++++++++++----------------- 2 files changed, 161 insertions(+), 163 deletions(-) diff --git a/src/tower/_storage.py b/src/tower/_storage.py index 8251636d..6528b86d 100644 --- a/src/tower/_storage.py +++ b/src/tower/_storage.py @@ -77,19 +77,16 @@ def _auth_from_context(context: TowerContext) -> tuple[str, str, str]: return token, auth_header_name, prefix -def _build_tower_control_plane_client( +def _new_tower_control_plane_client( *, - context: TowerContext, - tower_url: str, + base_url: str, + token: str, + auth_header_name: str, + prefix: str, timeout: float, - verify_tls: bool, -) -> tuple[str, str, AuthenticatedClient]: - token, auth_header_name, prefix = _auth_from_context(context) - base_url = _api_base_url(tower_url) - # hash whatever auth token was provided so we can cache catalogs tokens per account. - auth_hash = hashlib.sha256(token.encode("utf-8")).hexdigest() - client = AuthenticatedClient( - verify_ssl=verify_tls, +) -> AuthenticatedClient: + return AuthenticatedClient( + verify_ssl=True, base_url=base_url, token=token, auth_header_name=auth_header_name, @@ -97,60 +94,49 @@ def _build_tower_control_plane_client( timeout=httpx.Timeout(timeout), raise_on_unexpected_status=True, ) - return base_url, auth_hash, client -class StorageClient: - """Configuration and authenticated Tower Storage-specific client. +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() - This remains an internal foundation until the public catalog-loading surface is - added. It intentionally does not alter the clients used by unrelated SDK calls. - """ + +class _StorageResolver: + """Private Tower configuration and authentication for catalog resolution.""" def __init__( self, *, - tower_url: str | None = None, environment: str | None = None, - timeout: float = DEFAULT_STORAGE_TIMEOUT_SECONDS, - verify_tls: bool = True, ) -> None: context = TowerContext.build() - if tower_url is not None and not isinstance(tower_url, str): - raise TypeError("tower_url must be a string or None") - - if tower_url is not None and not tower_url.strip(): - raise ValueError("tower_url must not be blank") - 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") - if isinstance(timeout, bool) or not isinstance(timeout, (int, float)): - raise TypeError("timeout must be a positive number") - - if not 0 < float(timeout) < float("inf"): - raise ValueError("timeout must be a positive finite number") - - if not isinstance(verify_tls, bool): - raise TypeError("verify_tls must be a bool") - - self.tower_url = tower_url or context.tower_url - self.environment = ( + self._target_environment = ( environment or context.environment or DEFAULT_ENVIRONMENT_NAME ) - self.timeout = float(timeout) - self.verify_tls = verify_tls - self._base_url, self._auth_hash, self._tower_client = ( - _build_tower_control_plane_client( - context=context, - tower_url=self.tower_url, - timeout=self.timeout, - verify_tls=self.verify_tls, - ) + 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( @@ -161,12 +147,13 @@ def _request_catalog_credentials( body = VendCatalogCredentialsBody(mode=_vend_mode(mode)) try: - return vend_catalog_credentials_api.sync( - name=name, - client=self._tower_client, - environment=self.environment, - body=body, - ) + 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!r}." @@ -219,9 +206,9 @@ def get_tower_catalog_credentials( environment: Optional[str] = None, mode: str = "read", ) -> CatalogCredentials: - storage_client = StorageClient(environment=environment) + storage_resolver = _StorageResolver(environment=environment) mode = _normalize_mode(mode) - cache_key = _cache_key(storage_client, name, mode) + cache_key = _cache_key(storage_resolver, name, mode) now = datetime.now(timezone.utc) _prune_credential_cache(now) @@ -229,8 +216,7 @@ def get_tower_catalog_credentials( if cached is not None and cached.is_usable(now): return cached.credentials - with storage_client._tower_client: - credentials = _vend_with_default_catalog_fallback(storage_client, name, mode) + credentials = _vend_with_default_catalog_fallback(storage_resolver, name, mode) _credential_cache[cache_key] = _CachedCredentials(credentials) return credentials @@ -248,23 +234,23 @@ def load_vended_catalog(name: str, credentials: CatalogCredentials) -> Any: def _vend_with_default_catalog_fallback( - storage_client: StorageClient, + storage_resolver: _StorageResolver, name: str, mode: str, ) -> CatalogCredentials: - environment = storage_client.environment - result = storage_client._request_catalog_credentials(name, 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(storage_client) + _ensure_legacy_default_catalog(storage_resolver) for delay in DEFAULT_CATALOG_PROVISION_RETRY_DELAYS: time.sleep(delay) - result = storage_client._request_catalog_credentials(name, 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(storage_client) + _ensure_legacy_default_catalog(storage_resolver) return _unwrap_vend_result(result, name, environment) @@ -279,12 +265,9 @@ def _describe_tower_catalog_type( if ctx.jwt is None and ctx.api_key is None: return None - base_url, auth_hash, tower_client = _build_tower_control_plane_client( - context=ctx, - tower_url=ctx.tower_url, - timeout=CATALOG_TYPE_DESCRIBE_TIMEOUT_SECONDS, - verify_tls=True, - ) + 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: @@ -296,6 +279,14 @@ 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: with tower_client: result = describe_catalog_api.sync( @@ -339,9 +330,10 @@ def _failed_catalog_type_cache_entry() -> _CachedCatalogType: ) -def _ensure_legacy_default_catalog(storage_client: StorageClient) -> None: +def _ensure_legacy_default_catalog(storage_resolver: _StorageResolver) -> None: try: - describe_default_catalog_api.sync(client=storage_client._tower_client) + 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 @@ -369,15 +361,15 @@ def _unwrap_vend_result( def _cache_key( - storage_client: StorageClient, + storage_resolver: _StorageResolver, name: str, mode: str, ) -> tuple[str, str, str, str, str]: return ( - storage_client._base_url, - storage_client._auth_hash, + storage_resolver._base_url, + storage_resolver._auth_hash, name, - storage_client.environment, + storage_resolver._target_environment, mode, ) diff --git a/tests/tower/test_storage.py b/tests/tower/test_storage.py index 7877ae24..7f4307fd 100644 --- a/tests/tower/test_storage.py +++ b/tests/tower/test_storage.py @@ -1,4 +1,3 @@ -import hashlib from datetime import datetime, timedelta, timezone from http import HTTPStatus @@ -56,26 +55,19 @@ def test_context_treats_blank_auth_env_as_missing(monkeypatch): assert ctx.jwt is None -def test_storage_client_configuration_and_tls_defaults(monkeypatch): +def test_storage_resolver_configuration_and_tls_defaults(monkeypatch): + monkeypatch.setenv("TOWER_URL", "https://tower.example.com/") monkeypatch.setenv("TOWER_API_KEY", "ambient-key") - client = _storage.StorageClient( - tower_url="https://tower.example.com/", - environment="production", - timeout=12.5, - ) - transport = client._tower_client + resolver = _storage._StorageResolver(environment="production") + client = resolver._new_client() - assert client.tower_url == "https://tower.example.com/" - assert client.environment == "production" - assert client.timeout == 12.5 - assert client.verify_tls is True + 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._auth_hash == hashlib.sha256(b"ambient-key").hexdigest() - assert transport._base_url == "https://tower.example.com/v1" - assert transport._timeout == httpx.Timeout(12.5) - assert transport._verify_ssl is True + assert client._timeout == httpx.Timeout(_storage.DEFAULT_STORAGE_TIMEOUT_SECONDS) + assert client._verify_ssl is True - http_client = transport.get_httpx_client() + http_client = client.get_httpx_client() try: assert http_client.headers["X-API-Key"] == "ambient-key" assert "Authorization" not in http_client.headers @@ -95,7 +87,7 @@ def test_storage_client_configuration_and_tls_defaults(monkeypatch): ], ids=("jwt-over-api-key", "api-key-fallback"), ) -def test_storage_client_static_auth_precedence( +def test_storage_resolver_static_auth_precedence( monkeypatch, ambient_auth, expected_header, @@ -104,8 +96,8 @@ def test_storage_client_static_auth_precedence( for name, value in ambient_auth.items(): monkeypatch.setenv(name, value) - transport = _storage.StorageClient()._tower_client - http_client = transport.get_httpx_client() + client = _storage._StorageResolver()._new_client() + http_client = client.get_httpx_client() try: assert http_client.headers[expected_header] == expected_value other_header = ( @@ -116,6 +108,12 @@ def test_storage_client_static_auth_precedence( 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() @@ -134,12 +132,12 @@ def test_missing_auth_fails_before_cache_or_vend(monkeypatch): _storage.get_tower_catalog_credentials("analytics") -def test_storage_client_rejects_redacted_jwt_without_falling_back(monkeypatch): +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.StorageClient() + _storage._StorageResolver() @pytest.mark.parametrize("status", [HTTPStatus.UNAUTHORIZED, HTTPStatus.FORBIDDEN]) @@ -154,36 +152,37 @@ def vend(**kwargs): monkeypatch.setattr(_storage.vend_catalog_credentials_api, "sync", vend) - result = _storage.StorageClient()._request_catalog_credentials("analytics", "read") + result = _storage._StorageResolver()._request_catalog_credentials( + "analytics", "read" + ) assert result is rejected assert len(vend_calls) == 1 -def test_storage_client_vends_with_ambient_api_key(monkeypatch): +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() - try: - captured.update( - name=name, - environment=environment, - api_key=http_client.headers.get("X-API-Key"), - authorization=http_client.headers.get("Authorization"), - mode=body.mode, - ) - finally: - http_client.close() + 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.StorageClient( - tower_url="https://api.example.com", - environment="production", + result = _storage._StorageResolver( + environment="production" )._request_catalog_credentials("analytics", "read") assert result is response @@ -194,6 +193,8 @@ def vend(*, name, client, environment, body): "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): @@ -206,16 +207,13 @@ def test_describe_and_vend_prefer_jwt_when_both_auth_vars_are_set(monkeypatch): def capture_auth(operation, client): http_client = client.get_httpx_client() - try: - captured_auth.append( - ( - operation, - http_client.headers.get("Authorization"), - http_client.headers.get("X-API-Key"), - ) + captured_auth.append( + ( + operation, + http_client.headers.get("Authorization"), + http_client.headers.get("X-API-Key"), ) - finally: - http_client.close() + ) def describe(*, name, client, environment): capture_auth("describe", client) @@ -244,7 +242,7 @@ def vend(*, client, **kwargs): == _storage.TOWER_CATALOG_TYPE ) assert ( - _storage.StorageClient()._request_catalog_credentials("analytics", "read") + _storage._StorageResolver()._request_catalog_credentials("analytics", "read") is vended ) assert captured_auth == [ @@ -253,18 +251,13 @@ def vend(*, client, **kwargs): ] -@pytest.mark.parametrize("verify_tls", [True, False], ids=("verified", "unverified")) -def test_storage_client_allows_http_independently_of_tls_verification( - monkeypatch, verify_tls -): +def test_storage_resolver_allows_explicit_http_tower_url(monkeypatch): + monkeypatch.setenv("TOWER_URL", "http://localhost:9000") monkeypatch.setenv("TOWER_API_KEY", "key") - transport = _storage.StorageClient( - tower_url="http://localhost:9000", - verify_tls=verify_tls, - )._tower_client + client = _storage._StorageResolver()._new_client() - assert transport._base_url == "http://localhost:9000/v1" - assert transport._verify_ssl is verify_tls + 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): @@ -306,15 +299,17 @@ def vend(*, name, client, environment, body): ] -def test_storage_client_normalizes_and_validates_tower_api_url(monkeypatch): - monkeypatch.setenv("TOWER_API_KEY", "key") - client = _storage.StorageClient( - tower_url="https://TOWER.example.com:443/old/path?debug=true#fragment", +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", ) - assert client._tower_client._base_url == "https://tower.example.com/v1" + 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.StorageClient(tower_url="not-a-url") + _storage._api_base_url("not-a-url") def test_get_tower_catalog_credentials_rejects_invalid_mode_before_cache_or_vend( @@ -336,35 +331,41 @@ def test_get_tower_catalog_credentials_rejects_invalid_mode_before_cache_or_vend _storage.get_tower_catalog_credentials("analytics", mode="write") -@pytest.mark.parametrize("field", ["tower_url", "environment"]) -def test_storage_client_rejects_non_string_text_configuration(field): - with pytest.raises(TypeError, match=f"{field} must be a string or None"): - _storage.StorageClient(**{field: 123}) +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_client_maps_httpx_connection_errors(monkeypatch): +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(**kwargs): + 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.StorageClient()._request_catalog_credentials("analytics", "read") + _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_client_uses_runtime_environment_only_as_target_config(monkeypatch): +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.StorageClient().environment == "run-env" + assert _storage._StorageResolver()._target_environment == "run-env" assert ( - _storage.StorageClient(environment="explicit-env").environment == "explicit-env" + _storage._StorageResolver(environment="explicit-env")._target_environment + == "explicit-env" ) @@ -384,13 +385,13 @@ def test_get_tower_catalog_credentials_caches_vended_credentials(monkeypatch): calls = [] def vend(client, name, mode): - calls.append((name, client.environment, mode)) + calls.append((name, client._target_environment, mode)) return VendCatalogCredentialsResponse( credentials=credentials, - environment=client.environment, + environment=client._target_environment, ) - monkeypatch.setattr(_storage.StorageClient, "_request_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") @@ -419,8 +420,8 @@ def test_get_tower_catalog_credentials_prunes_expired_cache_entries(monkeypatch) oauth_token="oauth-token", warehouse="warehouse-id", ) - storage_client = _storage.StorageClient() - expired_key = _storage._cache_key(storage_client, "stale", "read") + storage_resolver = _storage._StorageResolver() + expired_key = _storage._cache_key(storage_resolver, "stale", "read") _storage._credential_cache[expired_key] = _storage._CachedCredentials( expired_credentials ) @@ -428,10 +429,10 @@ def test_get_tower_catalog_credentials_prunes_expired_cache_entries(monkeypatch) def vend(client, name, mode): return VendCatalogCredentialsResponse( credentials=fresh_credentials, - environment=client.environment, + environment=client._target_environment, ) - monkeypatch.setattr(_storage.StorageClient, "_request_catalog_credentials", vend) + monkeypatch.setattr(_storage._StorageResolver, "_request_catalog_credentials", vend) result = _storage.get_tower_catalog_credentials("default") @@ -459,20 +460,25 @@ def test_default_catalog_retries_reuse_client_auth_snapshot(monkeypatch): ] vend_tokens = [] legacy_tokens = [] - transports = [] + clients = [] def vend(*, client, **kwargs): - transports.append(client) + clients.append(client) vend_tokens.append(client.token) monkeypatch.setenv("TOWER_API_KEY", "changed-after-client-construction") return responses.pop(0) - def legacy_default(client): - transports.append(client._tower_client) - legacy_tokens.append(client._tower_client.token) + def legacy_default(*, client): + clients.append(client) + legacy_tokens.append(client.token) + return ErrorModel(status=404, detail="not provisioned") monkeypatch.setattr(_storage.vend_catalog_credentials_api, "sync", vend) - monkeypatch.setattr(_storage, "_ensure_legacy_default_catalog", legacy_default) + 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") @@ -480,9 +486,9 @@ def legacy_default(client): assert result is credentials assert vend_tokens == ["operation-token"] * 3 assert legacy_tokens == ["operation-token"] * 2 - assert all(transport is transports[0] for transport in transports) - assert transports[0]._client is not None - assert transports[0]._client.is_closed + 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( From ecfda23cce3afddf7093e5d0888647c212b23fe4 Mon Sep 17 00:00:00 2001 From: Konstantinos Stefanidis Vozikis Date: Thu, 20 Aug 2026 12:16:19 +0200 Subject: [PATCH 3/3] Address Storage URL review feedback --- src/tower/_storage.py | 21 ++++++++++++--------- 1 file changed, 12 insertions(+), 9 deletions(-) diff --git a/src/tower/_storage.py b/src/tower/_storage.py index 6528b86d..9975a171 100644 --- a/src/tower/_storage.py +++ b/src/tower/_storage.py @@ -5,10 +5,13 @@ import time from dataclasses import dataclass from datetime import datetime, timedelta, timezone -from typing import Any, Optional +from typing import TYPE_CHECKING import httpx +if TYPE_CHECKING: + from pyiceberg.catalog import Catalog + from ._context import TowerContext from .exceptions import ( StorageConnectionError, @@ -156,7 +159,7 @@ def _request_catalog_credentials( ) except httpx.RequestError as error: raise StorageConnectionError( - f"Could not connect to Tower at {self._base_url!r}." + f"Could not connect to Tower at {self._base_url}." ) from error @@ -164,10 +167,10 @@ 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!r}") from 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!r}") - return str(url.copy_with(path="/v1", query=None, fragment=None)).rstrip("/") + raise ValueError(f"Invalid Tower URL: {tower_url}") + return str(url.copy_with(path="/v1", query=None, fragment=None)) @dataclass @@ -191,9 +194,9 @@ class _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. """ @@ -203,7 +206,7 @@ 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: storage_resolver = _StorageResolver(environment=environment) @@ -221,7 +224,7 @@ def get_tower_catalog_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(