diff --git a/src/tower/_storage.py b/src/tower/_storage.py index 01477b04..ebe5ca1a 100644 --- a/src/tower/_storage.py +++ b/src/tower/_storage.py @@ -1,6 +1,7 @@ from __future__ import annotations import hashlib +import logging import time from dataclasses import dataclass from datetime import datetime, timedelta, timezone @@ -12,11 +13,13 @@ 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, ) from .tower_api_client.models import ( CatalogCredentials, + DescribeCatalogResponse, ErrorModel, VendCatalogCredentialsBody, VendCatalogCredentialsBodyMode, @@ -28,6 +31,8 @@ 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" +logger = logging.getLogger("tower.storage") @dataclass @@ -40,6 +45,7 @@ def is_usable(self, now: datetime) -> bool: _credential_cache: dict[tuple[str, str, str, str, str], _CachedCredentials] = {} +_catalog_type_cache: dict[tuple[str, str, str], str | None] = {} def get_tower_catalog( @@ -123,6 +129,50 @@ def _vend_catalog_credentials( ) +def _describe_tower_catalog_type( + ctx: TowerContext, name: str, environment: str +) -> str | None: + if not (ctx.api_key or ctx.jwt): + return None + + cache_key = (ctx.tower_url, name, environment) + if cache_key in _catalog_type_cache: + return _catalog_type_cache[cache_key] + + try: + result = describe_catalog_api.sync( + name=name, + client=_env_client(ctx), + environment=environment, + ) + except Exception: + logger.debug( + "Failed to describe Tower catalog %r in environment %r; " + "falling back to PyIceberg catalog configuration detection.", + name, + environment, + exc_info=True, + ) + return None + + if isinstance(result, DescribeCatalogResponse): + catalog_type = result.catalog.type_ + _catalog_type_cache[cache_key] = catalog_type + return catalog_type + + if isinstance(result, ErrorModel): + logger.debug( + "Tower catalog describe for %r in environment %r returned %s; " + "falling back to PyIceberg catalog configuration detection.", + name, + environment, + _error_text(result), + ) + + _catalog_type_cache[cache_key] = None + return None + + def _ensure_legacy_default_catalog(ctx: TowerContext) -> None: try: response = describe_default_catalog_api.sync_detailed(client=_env_client(ctx)) @@ -210,3 +260,4 @@ def _ensure_aware(value: datetime) -> datetime: def _clear_credential_cache() -> None: _credential_cache.clear() + _catalog_type_cache.clear() diff --git a/src/tower/_tables.py b/src/tower/_tables.py index 0f56540b..2b0d8401 100644 --- a/src/tower/_tables.py +++ b/src/tower/_tables.py @@ -1,5 +1,6 @@ from __future__ import annotations +import os from dataclasses import dataclass from typing import List, Optional, TypeVar, Union @@ -20,7 +21,12 @@ from pyiceberg.table import Table as IcebergTable from ._context import TowerContext -from ._storage import get_tower_catalog_credentials, load_vended_catalog +from ._storage import ( + TOWER_CATALOG_TYPE, + _describe_tower_catalog_type, + get_tower_catalog_credentials, + load_vended_catalog, +) from .tower_api_client.models import CatalogCredentials from .utils.pyarrow import ( convert_pyarrow_expressions, @@ -60,6 +66,45 @@ def _load_tower_catalog( return load_vended_catalog(name, credentials), _vended_catalog_identity(credentials) +def _pyiceberg_catalog_env_prefix(name: str) -> str: + catalog_name = name.replace("-", "_").replace(".", "_").replace(":", "_").upper() + return f"PYICEBERG_CATALOG__{catalog_name}__" + + +def _has_pyiceberg_catalog_config(name: str) -> bool: + try: + from pyiceberg.catalog import _ENV_CONFIG + + if _ENV_CONFIG.get_catalog_config(name) is not None: + return True + except Exception: + pass + + prefix = _pyiceberg_catalog_env_prefix(name) + return any(key.upper().startswith(prefix) for key in os.environ) + + +def _should_vend_tower_credentials( + ctx: TowerContext, + name: str, + environment: str, + tower_credentials: Optional[bool], +) -> bool: + """Choose Tower vending only for managed catalogs unless explicitly overridden. + + BYO catalogs such as S3 Tables already receive PyIceberg config from the runner, + so the default path must preserve that instead of forcing Tower vending. + """ + if tower_credentials is not None: + return tower_credentials + + catalog_type = _describe_tower_catalog_type(ctx, name, environment) + if catalog_type is not None: + return catalog_type == TOWER_CATALOG_TYPE + + return not _has_pyiceberg_catalog_config(name) + + class Table: """ `Table` is a wrapper around an Iceberg table. It provides methods to read and @@ -743,10 +788,11 @@ def tables( namespace (Optional[str], optional): The namespace in which the table exists or should be created. If not provided, a default namespace will be used. tower_credentials (Optional[bool], optional): Credential resolution for string - catalogs. By default (None) and when True, credentials are vended from - Tower. Set False to fall back to existing PyIceberg configuration (the - legacy ``PYICEBERG_CATALOG__*`` env vars) — a temporary rollback hatch. - Ignored when a Catalog instance is passed. + catalogs. By default (None), Tower-managed catalogs vend credentials and + other configured catalogs use existing PyIceberg configuration (including + runner-injected ``PYICEBERG_CATALOG__*`` env vars for S3 Tables). Set + True to force Tower credential vending or False to force PyIceberg + configuration. Ignored when a Catalog instance is passed. Returns: TableReference: A reference object that can be used to: @@ -788,19 +834,20 @@ def tables( vended_catalog_identity = None if isinstance(catalog, str): - # Tower-managed catalogs always resolve through credential vending. - # `tower_credentials=False` is a rollback hatch to the legacy - # PYICEBERG_CATALOG__* env-var config (still injected by the runner); - # remove it once that injection is gone. - if tower_credentials is False: - catalog = load_catalog(catalog) - else: + if _should_vend_tower_credentials( + ctx, + catalog, + ctx.environment, + tower_credentials, + ): catalog, vended_catalog_identity = _load_tower_catalog( catalog, environment=ctx.environment, mode="read", ) tower_vended = True + else: + catalog = load_catalog(catalog) return TableReference( ctx, diff --git a/tests/tower/test_tables.py b/tests/tower/test_tables.py index 38429a7d..946e1b11 100644 --- a/tests/tower/test_tables.py +++ b/tests/tower/test_tables.py @@ -19,8 +19,13 @@ # Imports the library under test import tower import tower._tables as tables_module +from tower import _storage from tower._context import TowerContext -from tower.tower_api_client.models import CatalogCredentials +from tower.tower_api_client.models import ( + Catalog, + CatalogCredentials, + DescribeCatalogResponse, +) class FakeLoadedTable: @@ -62,6 +67,18 @@ def patch_tower_context( return ctx +def make_describe_catalog_response(name: str, catalog_type: str): + return DescribeCatalogResponse( + catalog=Catalog( + created_at=datetime.datetime.now(datetime.timezone.utc), + environment="production", + name=name, + properties=[], + type_=catalog_type, + ) + ) + + def make_catalog_credentials(mode: str, token: str | None = None): return CatalogCredentials( catalog_uri="https://catalog.example.com", @@ -123,16 +140,20 @@ def sql_catalog(): @pytest.mark.parametrize( - ("tower_credentials", "expected_source"), + ("tower_credentials", "catalog_type", "has_pyiceberg_config", "expected_source"), [ - (None, "vend"), - (True, "vend"), - (False, "load_catalog"), + (None, "tower-catalog", True, "vend"), + (None, "s3-tables", True, "load_catalog"), + (None, None, True, "load_catalog"), + (None, None, False, "vend"), + (True, "s3-tables", True, "vend"), + (False, "tower-catalog", False, "load_catalog"), ], ) -def test_string_catalog_precedence(monkeypatch, tower_credentials, expected_source): - # Tower-managed catalogs vend by default and when forced; `tower_credentials=False` - # is the only path that falls back to existing PyIceberg configuration. +def test_string_catalog_precedence( + monkeypatch, tower_credentials, catalog_type, has_pyiceberg_config, expected_source +): + _storage._clear_credential_cache() patch_tower_context(monkeypatch) vend_catalog = FakeCatalog("vend") configured_catalog = FakeCatalog("configured") @@ -150,11 +171,27 @@ def load_catalog(name): calls.append(("load_catalog", name)) return configured_catalog + def describe_catalog_api_sync(name, client, environment): + calls.append(("describe_catalog", name, environment)) + if catalog_type is None: + return None + return make_describe_catalog_response(name, catalog_type) + + def has_pyiceberg_catalog_config(name): + calls.append(("has_pyiceberg_config", name)) + return has_pyiceberg_config + monkeypatch.setattr( tables_module, "get_tower_catalog_credentials", get_tower_catalog_credentials ) monkeypatch.setattr(tables_module, "load_vended_catalog", load_vended_catalog) monkeypatch.setattr(tables_module, "load_catalog", load_catalog) + monkeypatch.setattr( + _storage.describe_catalog_api, "sync", describe_catalog_api_sync + ) + monkeypatch.setattr( + tables_module, "_has_pyiceberg_catalog_config", has_pyiceberg_catalog_config + ) ref = tables_module.tables( "events", @@ -165,14 +202,64 @@ def load_catalog(name): if expected_source == "vend": assert ref._catalog is vend_catalog assert ref._tower_vended is True - assert calls == [ - ("vend", "default", "production", "read"), - ("load_vended_catalog", "default", "read"), - ] + assert ("vend", "default", "production", "read") in calls + assert ("load_vended_catalog", "default", "read") in calls else: assert ref._catalog is configured_catalog assert ref._tower_vended is False - assert calls == [("load_catalog", "default")] + assert ("load_catalog", "default") in calls + assert ("vend", "default", "production", "read") not in calls + + if tower_credentials is None: + assert ("describe_catalog", "default", "production") in calls + if catalog_type is None: + assert ("has_pyiceberg_config", "default") in calls + else: + assert ("describe_catalog", "default", "production") not in calls + assert ("has_pyiceberg_config", "default") not in calls + + +def test_pyiceberg_catalog_config_detects_runner_env(monkeypatch): + monkeypatch.setenv("PYICEBERG_CATALOG__S3_TABLES__URI", "https://example.com") + + assert tables_module._has_pyiceberg_catalog_config("s3-tables") is True + assert tables_module._has_pyiceberg_catalog_config("other") is False + + +def test_string_catalog_type_describe_is_cached(monkeypatch): + _storage._clear_credential_cache() + patch_tower_context(monkeypatch) + calls = [] + vend_catalogs = [] + + def describe_catalog_api_sync(name, client, environment): + calls.append(("describe_catalog", name, environment)) + return make_describe_catalog_response(name, "tower-catalog") + + def get_tower_catalog_credentials(name, environment=None, mode="read"): + calls.append(("vend", name, environment, mode)) + return make_catalog_credentials(mode) + + def load_vended_catalog(name, credentials): + catalog = FakeCatalog(credentials.mode) + vend_catalogs.append(catalog) + return catalog + + monkeypatch.setattr( + _storage.describe_catalog_api, "sync", describe_catalog_api_sync + ) + monkeypatch.setattr( + tables_module, "get_tower_catalog_credentials", get_tower_catalog_credentials + ) + monkeypatch.setattr(tables_module, "load_vended_catalog", load_vended_catalog) + + first = tables_module.tables("events", catalog="default") + second = tables_module.tables("users", catalog="default") + + assert first._catalog is vend_catalogs[0] + assert second._catalog is vend_catalogs[1] + assert calls.count(("describe_catalog", "default", "production")) == 1 + assert calls.count(("vend", "default", "production", "read")) == 2 def test_vended_table_write_lazily_escalates_and_reuses_catalog(monkeypatch):