diff --git a/src/tower/_storage.py b/src/tower/_storage.py index ebe5ca1a..bff78bc1 100644 --- a/src/tower/_storage.py +++ b/src/tower/_storage.py @@ -10,10 +10,10 @@ from ._client import _env_client from ._context import TowerContext +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, ) -from .tower_api_client.api.default import describe_catalog as describe_catalog_api from .tower_api_client.api.default import ( vend_catalog_credentials as vend_catalog_credentials_api, ) @@ -28,6 +28,11 @@ from .tower_api_client.types import UNSET, Unset CREDENTIAL_REFRESH_WINDOW = timedelta(minutes=5) +# how long to wait for a catalog type describe request to complete +CATALOG_TYPE_DESCRIBE_TIMEOUT_SECONDS = 2.0 +# only cache failed catalog type describe requests for this long +# retry only after this period +CATALOG_TYPE_FAILURE_CACHE_TTL_SECONDS = 30.0 DEFAULT_CATALOG_PROVISION_RETRY_DELAYS = (0.25, 0.5, 1.0, 2.0) DEFAULT_CATALOG_NAME = "default" DEFAULT_ENVIRONMENT_NAME = "default" @@ -44,8 +49,14 @@ def is_usable(self, now: datetime) -> bool: return now < expires_at - CREDENTIAL_REFRESH_WINDOW +@dataclass +class _CachedCatalogType: + catalog_type: str | None + retry_at: float | None = None + + _credential_cache: dict[tuple[str, str, str, str, str], _CachedCredentials] = {} -_catalog_type_cache: dict[tuple[str, str, str], str | None] = {} +_catalog_type_cache: dict[tuple[str, str, str], _CachedCatalogType] = {} def get_tower_catalog( @@ -136,13 +147,20 @@ def _describe_tower_catalog_type( return None cache_key = (ctx.tower_url, name, environment) - if cache_key in _catalog_type_cache: - return _catalog_type_cache[cache_key] + cached = _catalog_type_cache.get(cache_key) + if cached is not None: + if cached.retry_at is None: + return cached.catalog_type + + if time.monotonic() < cached.retry_at: + return None + + _catalog_type_cache.pop(cache_key, None) try: result = describe_catalog_api.sync( name=name, - client=_env_client(ctx), + client=_env_client(ctx, timeout=CATALOG_TYPE_DESCRIBE_TIMEOUT_SECONDS), environment=environment, ) except Exception: @@ -153,11 +171,12 @@ def _describe_tower_catalog_type( environment, exc_info=True, ) + _catalog_type_cache[cache_key] = _failed_catalog_type_cache_entry() return None if isinstance(result, DescribeCatalogResponse): catalog_type = result.catalog.type_ - _catalog_type_cache[cache_key] = catalog_type + _catalog_type_cache[cache_key] = _CachedCatalogType(catalog_type=catalog_type) return catalog_type if isinstance(result, ErrorModel): @@ -169,10 +188,17 @@ def _describe_tower_catalog_type( _error_text(result), ) - _catalog_type_cache[cache_key] = None + _catalog_type_cache[cache_key] = _failed_catalog_type_cache_entry() return None +def _failed_catalog_type_cache_entry() -> _CachedCatalogType: + return _CachedCatalogType( + catalog_type=None, + retry_at=time.monotonic() + CATALOG_TYPE_FAILURE_CACHE_TTL_SECONDS, + ) + + def _ensure_legacy_default_catalog(ctx: TowerContext) -> None: try: response = describe_default_catalog_api.sync_detailed(client=_env_client(ctx)) diff --git a/tests/tower/test_storage.py b/tests/tower/test_storage.py index 2187b57f..43c70842 100644 --- a/tests/tower/test_storage.py +++ b/tests/tower/test_storage.py @@ -3,7 +3,9 @@ from tower._context import TowerContext from tower import _storage from tower.tower_api_client.models import ( + Catalog, CatalogCredentials, + DescribeCatalogResponse, ErrorModel, VendCatalogCredentialsResponse, ) @@ -181,3 +183,69 @@ def legacy_default(ctx): assert result is credentials assert len(legacy_calls) == 2 + + +def test_describe_tower_catalog_type_uses_timeout_and_recovers_after_cooldown( + monkeypatch, +): + _storage._clear_credential_cache() + ctx = TowerContext( + tower_url="https://api.example.com", + environment="production", + api_key="api-key", + ) + now = {"value": 100.0} + calls = [] + + def describe_catalog_api_sync(name, client, environment): + calls.append((name, environment, client._timeout)) + if len(calls) == 1: + raise TimeoutError("describe timed out") + + return DescribeCatalogResponse( + catalog=Catalog( + created_at=datetime.now(timezone.utc), + environment=environment, + name=name, + properties=[], + type_="s3-tables", + ) + ) + + monkeypatch.setattr(_storage.time, "monotonic", lambda: now["value"]) + monkeypatch.setattr( + _storage.describe_catalog_api, "sync", describe_catalog_api_sync + ) + + assert _storage._describe_tower_catalog_type(ctx, "s3-tables", "production") is None + assert _storage._describe_tower_catalog_type(ctx, "s3-tables", "production") is None + assert calls == [ + ( + "s3-tables", + "production", + _storage.CATALOG_TYPE_DESCRIBE_TIMEOUT_SECONDS, + ) + ] + + now["value"] += _storage.CATALOG_TYPE_FAILURE_CACHE_TTL_SECONDS + 0.1 + + assert ( + _storage._describe_tower_catalog_type(ctx, "s3-tables", "production") + == "s3-tables" + ) + assert ( + _storage._describe_tower_catalog_type(ctx, "s3-tables", "production") + == "s3-tables" + ) + assert calls == [ + ( + "s3-tables", + "production", + _storage.CATALOG_TYPE_DESCRIBE_TIMEOUT_SECONDS, + ), + ( + "s3-tables", + "production", + _storage.CATALOG_TYPE_DESCRIBE_TIMEOUT_SECONDS, + ), + ]