diff --git a/INSTALL-AND-REFERENCE.md b/INSTALL-AND-REFERENCE.md index 0cccf67e..06a4e5ee 100644 --- a/INSTALL-AND-REFERENCE.md +++ b/INSTALL-AND-REFERENCE.md @@ -128,7 +128,54 @@ pip install "tower[ai]" pip install "tower[iceberg]" ``` -- `tower.tables(...)`: load, create, update, and delete Iceberg table data +- `tower.load_catalog(...)`: use a Tower-managed catalog through native PyIceberg APIs +- `tower.tables(...)`: load, create, update, and delete Iceberg table data through Tower's convenience API + +#### Native PyIceberg access + +Set `TOWER_JWT` or `TOWER_API_KEY` in the Python process, then load a catalog for +the target environment: + +```python +import tower + +catalog = tower.load_catalog( + "analytics", + environment="production", +) + +catalog.list_namespaces() +events = catalog.load_table(("analytics", "events")) +``` + +`tower.load_catalog()` returns an unwrapped `pyiceberg.catalog.Catalog`. Tower +first looks for the named catalog in the target environment, then uses the +same-named catalog from the shared `default` environment when it is not locally +defined. Only Tower-managed catalogs are supported by this API. + +Access is read-only by default. Request write credentials explicitly for +namespace or table mutations: + +```python +catalog = tower.load_catalog( + "analytics", + environment="production", + mode="read-write", +) +``` + +When both authentication variables are set, `TOWER_JWT` takes precedence over +`TOWER_API_KEY`. `tower login` authenticates the CLI and MCP server only; it does +not authenticate Python SDK calls. + +The returned catalog has one fixed temporary provider token and does not renew +itself. If the token expires, the original PyIceberg or provider error is +preserved. Call `tower.load_catalog()` again to obtain a fresh handle. Missing, +forbidden, unsupported, and connection failures are available as distinct +exceptions in `tower.exceptions`. A missing `default` catalog is not created +automatically. + +#### Tower table convenience API Delete filters are SQL-like strings or native PyIceberg boolean expressions. The `Table.column()` builder creates composable PyIceberg predicates: diff --git a/README.md b/README.md index 438c5414..6b02f441 100644 --- a/README.md +++ b/README.md @@ -122,7 +122,7 @@ If you find Tower useful, consider giving the repo a ⭐. - **Consistent execution environment everywhere** - `tower run` on Tower serverless or your own compute. - **Deploy in under 30 seconds** - `tower deploy` packages and ships. - **Secrets** - CLI-managed; injected as env vars in runner only (E2E encrypted). -- **Optional AI inference, Iceberg or dbt** - `tower[ai]`, `tower[iceberg]` or `tower[dbt]`; [details](INSTALL-AND-REFERENCE.md#optional-features). +- **Optional AI inference, native PyIceberg catalogs, or dbt** - `tower[ai]`, `tower[iceberg]` or `tower[dbt]`; [details](INSTALL-AND-REFERENCE.md#optional-features). - **MCP server** - Deploy and launch runs from AI coding assistants; [details](https://docs.tower.dev/docs/reference/mcp-server). --- diff --git a/src/tower/__init__.py b/src/tower/__init__.py index 27c1bc1d..6f2e7d94 100644 --- a/src/tower/__init__.py +++ b/src/tower/__init__.py @@ -30,6 +30,7 @@ from ._features import override_get_attr, get_available_features, is_feature_enabled if TYPE_CHECKING: + from ._storage import load_catalog as load_catalog from ._tables import tables as tables from ._llms import llms as llms from ._dbt import dbt as dbt diff --git a/src/tower/_features.py b/src/tower/_features.py index 752fd057..0e031323 100644 --- a/src/tower/_features.py +++ b/src/tower/_features.py @@ -42,7 +42,7 @@ def get_package_version(pkg_name: str) -> Optional[str]: # export_type can be "function" (return specific function) or "module" (return entire module) _feature_modules: Dict[str, tuple[str, List[str], str]] = { "ai": ("_llms", ["llms"], "function"), - "iceberg": ("_tables", ["tables"], "function"), + "iceberg": ("_iceberg", ["load_catalog", "tables"], "function"), "dbt": ("_dbt", ["dbt"], "function"), } diff --git a/src/tower/_iceberg.py b/src/tower/_iceberg.py new file mode 100644 index 00000000..6709553f --- /dev/null +++ b/src/tower/_iceberg.py @@ -0,0 +1,6 @@ +"""Public exports for Tower's optional Apache Iceberg feature.""" + +from ._storage import load_catalog +from ._tables import tables + +__all__ = ["load_catalog", "tables"] diff --git a/src/tower/_storage.py b/src/tower/_storage.py index c0db1286..5867a427 100644 --- a/src/tower/_storage.py +++ b/src/tower/_storage.py @@ -7,25 +7,24 @@ from dataclasses import dataclass, field from datetime import datetime, timedelta, timezone from threading import Lock -from typing import TYPE_CHECKING +from typing import Literal import httpx - -if TYPE_CHECKING: - from pyiceberg.catalog import Catalog +from pyiceberg.catalog import Catalog from ._context import TowerContext from .exceptions import ( + StorageAuthenticationError, + StorageCatalogNotFoundError, StorageConnectionError, StorageError, StorageInvalidCredentialError, StorageMissingAuthenticationError, + StoragePermissionError, + StorageUnsupportedCatalogError, ) 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, -) from .tower_api_client.api.default import ( vend_catalog_credentials as vend_catalog_credentials_api, ) @@ -46,7 +45,6 @@ # 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" @@ -208,10 +206,10 @@ def _resolve_catalog_access( return flight.result() try: - response = _vend_with_default_catalog_fallback( - self, + response = _unwrap_vend_result( + self._request_catalog_credentials(name, mode), name, - mode, + target_environment, ) credentials = response.credentials if credentials.mode != mode: @@ -301,16 +299,58 @@ class _CachedCatalogType: _catalog_type_cache: dict[_CatalogTypeCacheKey, _CachedCatalogType] = {} -def get_tower_catalog( +def load_catalog( name: str = DEFAULT_CATALOG_NAME, + *, environment: str | None = None, - mode: str = "read", + mode: Literal["read", "read-write"] = "read", ) -> Catalog: + """Load a Tower-managed catalog as a native PyIceberg catalog. + + The SDK authenticates from the process environment. ``TOWER_JWT`` takes + precedence over ``TOWER_API_KEY`` when both are set. A Tower CLI login does + not authenticate this Python API. + + The returned catalog uses one fixed set of temporary provider credentials. + It does not retain a Tower resolver or renew itself. If those credentials + expire, the provider's original error is raised; call ``load_catalog`` again + to obtain a fresh catalog handle. + + Args: + name: Name of the Tower-managed catalog. Defaults to ``"default"``. + environment: Target Tower environment. Tower first resolves the catalog + in this environment, then in the shared ``default`` environment. If + omitted, the configured Tower environment is used. + mode: ``"read"`` for read-only credentials, or ``"read-write"`` for + namespace and table mutations. Defaults to ``"read"``. + + Returns: + A caller-owned :class:`pyiceberg.catalog.Catalog` configured with the + resolved REST endpoint, warehouse, and temporary provider token. + + Raises: + StorageMissingAuthenticationError: If neither supported Tower + credential is configured. + StorageAuthenticationError: If Tower rejects the configured credential. + StorageCatalogNotFoundError: If the environment or catalog is not found. + StoragePermissionError: If the credential lacks the requested access. + StorageUnsupportedCatalogError: If the catalog is not Tower-managed or + otherwise cannot vend native PyIceberg access. + StorageConnectionError: If Tower's control-plane API cannot be reached. + TypeError: If ``environment`` is not a string or ``None``. + ValueError: If ``environment`` is blank or ``mode`` is invalid. + + Examples: + >>> import tower + >>> catalog = tower.load_catalog( + ... "analytics", + ... environment="production", + ... ) + >>> catalog.list_namespaces() """ - Load a PyIceberg REST catalog using short-lived credentials vended by Tower. - """ - credentials = get_tower_catalog_credentials(name, environment, mode) - return load_vended_catalog(name, credentials) + storage_resolver = _StorageResolver(environment=environment) + access = storage_resolver._resolve_catalog_access(name, mode) + return load_vended_catalog(name, access.to_credentials()) def get_tower_catalog_credentials( @@ -335,32 +375,6 @@ def load_vended_catalog(name: str, credentials: CatalogCredentials) -> Catalog: ) -def _vend_with_default_catalog_fallback( - storage_resolver: _StorageResolver, - name: str, - mode: str, -) -> VendCatalogCredentialsResponse: - 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_resolver) - for delay in DEFAULT_CATALOG_PROVISION_RETRY_DELAYS: - time.sleep(delay) - 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_resolver) - - return _unwrap_vend_result(result, name, environment) - - raise RuntimeError( - f"Tower catalog {name} does not exist in environment {environment}." - ) - - def _describe_tower_catalog_type( ctx: TowerContext, name: str, environment: str ) -> str | None: @@ -437,15 +451,6 @@ def _failed_catalog_type_cache_entry() -> _CachedCatalogType: ) -def _ensure_legacy_default_catalog(storage_resolver: _StorageResolver) -> None: - try: - 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 - - def _unwrap_vend_result( result: ErrorModel | VendCatalogCredentialsResponse | None, name: str, @@ -456,14 +461,25 @@ def _unwrap_vend_result( if isinstance(result, ErrorModel): detail = _error_text(result) - raise RuntimeError( - f"Failed to vend credentials for Tower catalog {name} " - f"in environment {environment}: {detail}" + message = ( + f"Could not load Tower catalog {name!r} " + f"in environment {environment!r}: {detail}" ) - - raise RuntimeError( - f"Failed to vend credentials for Tower catalog {name} " - f"in environment {environment}." + error_type: type[StorageError] + if result.status == 401: + error_type = StorageAuthenticationError + elif result.status == 403: + error_type = StoragePermissionError + elif result.status == 404: + error_type = StorageCatalogNotFoundError + elif result.status == 422: + error_type = StorageUnsupportedCatalogError + else: + error_type = StorageError + raise error_type(message) + + raise StorageError( + f"Could not load Tower catalog {name!r} in environment {environment!r}." ) @@ -481,10 +497,6 @@ def _vend_mode(mode: str) -> VendCatalogCredentialsBodyMode: ) -def _is_not_found(result: ErrorModel | VendCatalogCredentialsResponse | None) -> bool: - return isinstance(result, ErrorModel) and result.status == 404 - - def _error_text(error: ErrorModel) -> str: for value in (error.detail, error.title): if not isinstance(value, Unset) and value: diff --git a/src/tower/exceptions.py b/src/tower/exceptions.py index 0f0bfc9b..9f2e7b6c 100644 --- a/src/tower/exceptions.py +++ b/src/tower/exceptions.py @@ -63,3 +63,15 @@ class StorageInvalidCredentialError(StorageAuthenticationError): class StorageConnectionError(StorageError): """Tower's control-plane API could not be reached.""" + + +class StorageCatalogNotFoundError(StorageError): + """The target environment or requested catalog could not be found.""" + + +class StoragePermissionError(StorageError): + """The Tower credential lacks permission for the requested catalog access.""" + + +class StorageUnsupportedCatalogError(StorageError): + """Tower cannot vend native PyIceberg access for the requested catalog.""" diff --git a/tests/tower/test_storage.py b/tests/tower/test_storage.py index 2a0a136f..e444a372 100644 --- a/tests/tower/test_storage.py +++ b/tests/tower/test_storage.py @@ -2,17 +2,27 @@ from datetime import datetime, timedelta, timezone from http import HTTPStatus from threading import Event +from typing import get_type_hints import httpx +import pyiceberg.catalog import pytest +from pyiceberg.catalog import Catalog as PyIcebergCatalog +from pyiceberg.catalog.memory import InMemoryCatalog +import tower +from tower import _features from tower import _storage from tower._context import TowerContext from tower.exceptions import ( + StorageAuthenticationError, + StorageCatalogNotFoundError, StorageConnectionError, StorageError, StorageInvalidCredentialError, StorageMissingAuthenticationError, + StoragePermissionError, + StorageUnsupportedCatalogError, ) from tower.tower_api_client.models import ( Catalog, @@ -158,7 +168,7 @@ def test_missing_auth_fails_before_cache_or_vend(monkeypatch): ) with pytest.raises(StorageMissingAuthenticationError): - _storage.get_tower_catalog_credentials("analytics") + tower.load_catalog("analytics") def test_storage_resolver_rejects_redacted_jwt_without_falling_back(monkeypatch): @@ -336,7 +346,7 @@ def test_storage_resolver_normalizes_and_validates_tower_api_url(monkeypatch): _storage._api_base_url("not-a-url") -def test_get_tower_catalog_credentials_rejects_invalid_mode_before_cache_or_vend( +def test_load_catalog_rejects_invalid_mode_before_cache_or_vend( monkeypatch, ): monkeypatch.setenv("TOWER_API_KEY", "api-key") @@ -347,7 +357,7 @@ def test_get_tower_catalog_credentials_rejects_invalid_mode_before_cache_or_vend ) with pytest.raises(ValueError, match="mode must be 'read' or 'read-write'"): - _storage.get_tower_catalog_credentials("analytics", mode="write") + tower.load_catalog("analytics", mode="write") def test_storage_resolver_rejects_invalid_environment(): @@ -369,8 +379,7 @@ def vend(*, client, **kwargs): monkeypatch.setattr(_storage.vend_catalog_credentials_api, "sync", vend) with pytest.raises(StorageConnectionError) as error: - resolver = _storage._StorageResolver() - resolver._request_catalog_credentials("analytics", "read") + tower.load_catalog("analytics") assert error.value.__cause__ is cause assert clients[0]._client is not None @@ -421,7 +430,7 @@ def test_access_cache_is_owned_by_resolver(monkeypatch): assert len(calls) == 2 -def test_one_shot_credential_loads_do_not_share_cache(monkeypatch): +def test_load_catalog_calls_do_not_share_resolvers(monkeypatch): monkeypatch.setenv("TOWER_API_KEY", "api-key") calls = script_vend( monkeypatch, @@ -434,15 +443,172 @@ def test_one_shot_credential_loads_do_not_share_cache(monkeypatch): for number in (1, 2) ], ) + loaded_credentials = [] - first = _storage.get_tower_catalog_credentials("analytics") - second = _storage.get_tower_catalog_credentials("analytics") + def load_vended_catalog(name, credentials): + loaded_credentials.append(credentials) + return InMemoryCatalog(name) - assert first.oauth_token == "provider-token-1" - assert second.oauth_token == "provider-token-2" + monkeypatch.setattr(_storage, "load_vended_catalog", load_vended_catalog) + + first = tower.load_catalog("analytics") + second = tower.load_catalog("analytics") + + assert isinstance(first, PyIcebergCatalog) + assert isinstance(second, PyIcebergCatalog) + assert first is not second + assert [credentials.oauth_token for credentials in loaded_credentials] == [ + "provider-token-1", + "provider-token-2", + ] assert len(calls) == 2 +def test_load_catalog_returns_plain_native_catalog(monkeypatch): + monkeypatch.setenv("TOWER_API_KEY", "api-key") + expires_at = datetime.now(timezone.utc) + timedelta(hours=1) + calls = script_vend( + monkeypatch, + [ + make_vend_response( + environment="production", + token="provider-token", + expires_at=expires_at, + ) + ], + ) + native_catalog = InMemoryCatalog("analytics") + loaded = {} + + def load_pyiceberg_catalog(name, **properties): + loaded["name"] = name + loaded["properties"] = properties + return native_catalog + + monkeypatch.setattr(pyiceberg.catalog, "load_catalog", load_pyiceberg_catalog) + + catalog = tower.load_catalog("analytics", environment="production") + + assert catalog is native_catalog + assert catalog.list_namespaces() == [] + assert calls == [("analytics", "read")] + assert loaded["name"] == "analytics" + assert loaded["properties"] == { + "type": "rest", + "uri": "https://catalog.example.com", + "warehouse": "warehouse-id", + "token": "provider-token", + } + + +def test_load_catalog_preserves_provider_errors_and_reloads_explicitly(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) + ], + ) + expired_catalog = InMemoryCatalog("analytics") + fresh_catalog = InMemoryCatalog("analytics") + catalogs = iter((expired_catalog, fresh_catalog)) + provider_error = RuntimeError("provider token expired") + + def load_pyiceberg_catalog(name, **properties): + return next(catalogs) + + def raise_provider_error(*args, **kwargs): + raise provider_error + + monkeypatch.setattr(pyiceberg.catalog, "load_catalog", load_pyiceberg_catalog) + monkeypatch.setattr(expired_catalog, "list_namespaces", raise_provider_error) + + first = tower.load_catalog("analytics") + with pytest.raises(RuntimeError) as error: + first.list_namespaces() + + assert error.value is provider_error + assert len(calls) == 1 + + second = tower.load_catalog("analytics") + + assert second is fresh_catalog + assert len(calls) == 2 + + +def test_load_catalog_supports_explicit_read_write_access(monkeypatch): + monkeypatch.setenv("TOWER_API_KEY", "api-key") + calls = script_vend( + monkeypatch, + [ + make_vend_response( + environment="default", + token="write-token", + expires_at=datetime.now(timezone.utc) + timedelta(hours=1), + mode="read-write", + ) + ], + ) + loaded = [] + + def load_vended_catalog(name, credentials): + loaded.append((name, credentials)) + return InMemoryCatalog(name) + + monkeypatch.setattr(_storage, "load_vended_catalog", load_vended_catalog) + + tower.load_catalog("analytics", mode="read-write") + + assert calls == [("analytics", "read-write")] + assert loaded[0][1].mode == "read-write" + + +def test_load_catalog_rejects_unsupported_catalog_without_fallback(monkeypatch): + monkeypatch.setenv("TOWER_API_KEY", "api-key") + monkeypatch.setenv( + "PYICEBERG_CATALOG__ANALYTICS__URI", + "https://configured.example.com", + ) + calls = script_vend( + monkeypatch, + [ErrorModel(status=422, detail="only tower-managed catalogs are supported")], + ) + + def unexpected_call(*args, **kwargs): + raise AssertionError("unsupported catalogs must not use another catalog path") + + monkeypatch.setattr(_storage, "_describe_tower_catalog_type", unexpected_call) + monkeypatch.setattr(pyiceberg.catalog, "load_catalog", unexpected_call) + + with pytest.raises(StorageUnsupportedCatalogError, match="tower-managed"): + tower.load_catalog("analytics") + + assert calls == [("analytics", "read")] + + +def test_load_catalog_is_an_optional_iceberg_export(monkeypatch): + assert tower.load_catalog is _storage.load_catalog + assert get_type_hints(tower.load_catalog)["return"] is PyIcebergCatalog + assert tower.get_available_features()["iceberg"]["exports"] == [ + "load_catalog", + "tables", + ] + assert not hasattr(tower, "StorageResolver") + + monkeypatch.setattr( + _features, + "is_installed", + lambda dependency: dependency != "pyiceberg", + ) + with pytest.raises(ImportError, match=r"tower\[iceberg\]"): + _features.override_get_attr("load_catalog") + + def test_near_expiry_access_retains_inherited_and_local_identity(monkeypatch): monkeypatch.setenv("TOWER_ENVIRONMENT", "production") monkeypatch.setenv("TOWER_API_KEY", "api-key") @@ -585,54 +751,32 @@ def test_inconsistent_vend_response_is_not_cached( assert len(calls) == 2 -def test_default_catalog_retries_reuse_client_auth_snapshot(monkeypatch): - monkeypatch.setenv("TOWER_API_KEY", "operation-token") - 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", - ) - responses = [ - ErrorModel(status=404, detail="not found"), - ErrorModel(status=404, detail="still provisioning"), - VendCatalogCredentialsResponse( - credentials=credentials, - environment="default", - ), - ] - 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(*, 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.describe_default_catalog_api, - "sync", - legacy_default, +@pytest.mark.parametrize( + ("status", "error_type"), + [ + (HTTPStatus.UNAUTHORIZED, StorageAuthenticationError), + (HTTPStatus.FORBIDDEN, StoragePermissionError), + (HTTPStatus.NOT_FOUND, StorageCatalogNotFoundError), + (HTTPStatus.UNPROCESSABLE_ENTITY, StorageUnsupportedCatalogError), + (HTTPStatus.INTERNAL_SERVER_ERROR, StorageError), + ], +) +def test_load_catalog_maps_tower_errors_without_retry( + monkeypatch, + status, + error_type, +): + monkeypatch.setenv("TOWER_API_KEY", "api-key") + calls = script_vend( + monkeypatch, + [ErrorModel(status=int(status), detail="request rejected")], ) - monkeypatch.setattr(_storage.time, "sleep", lambda delay: None) - result = _storage.get_tower_catalog_credentials("default") + with pytest.raises(error_type, match="request rejected") as error: + tower.load_catalog("default") - assert result == credentials - 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) + assert type(error.value) is error_type + assert calls == [("default", "read")] def test_describe_tower_catalog_type_uses_timeout_and_recovers_after_cooldown(