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)