Check storage-credentials before config in table response

Author: Alexandra Popova

pyiceberg/catalog/rest/__init__.py

@@ -20,7 +20,6 @@
     Any,
     Dict,
     List,
-    Optional,
     Set,
     Tuple,
     Union,
@@ -153,18 +152,18 @@ def _retry_hook(retry_state: RetryCallState) -> None:
 
 
 class TableResponse(IcebergBaseModel):
-    metadata_location: Optional[str] = Field(alias="metadata-location", default=None)
+    metadata_location: str | None = Field(alias="metadata-location", default=None)
     metadata: TableMetadata
     config: Properties = Field(default_factory=dict)
-    storage_credentials: Optional[Properties] = Field(alias="storage-credentials", default=None)
+    storage_credentials: Properties | None = Field(alias="storage-credentials", default=None)
 
 
 class CreateTableRequest(IcebergBaseModel):
     name: str = Field()
-    location: Optional[str] = Field()
+    location: str | None = Field()
     table_schema: Schema = Field(alias="schema")
-    partition_spec: Optional[PartitionSpec] = Field(alias="partition-spec")
-    write_order: Optional[SortOrder] = Field(alias="write-order")
+    partition_spec: PartitionSpec | None = Field(alias="partition-spec")
+    write_order: SortOrder | None = Field(alias="write-order")
     stage_create: bool = Field(alias="stage-create", default=False)
     properties: Dict[str, str] = Field(default_factory=dict)
 
@@ -180,8 +179,8 @@ class RegisterTableRequest(IcebergBaseModel):
 
 
 class ConfigResponse(IcebergBaseModel):
-    defaults: Optional[Properties] = Field(default_factory=dict)
-    overrides: Optional[Properties] = Field(default_factory=dict)
+    defaults: Properties | None = Field(default_factory=dict)
+    overrides: Properties | None = Field(default_factory=dict)
 
 
 class ListNamespaceResponse(IcebergBaseModel):
@@ -295,7 +294,7 @@ def _create_legacy_oauth2_auth_manager(self, session: Session) -> AuthManager:
 
         return AuthManagerFactory.create("legacyoauth2", auth_config)
 
-    def _check_valid_namespace_identifier(self, identifier: Union[str, Identifier]) -> Identifier:
+    def _check_valid_namespace_identifier(self, identifier: str | Identifier) -> Identifier:
         """Check if the identifier has at least one element."""
         identifier_tuple = Catalog.identifier_to_tuple(identifier)
         if len(identifier_tuple) < 1:
@@ -378,22 +377,22 @@ def _fetch_config(self) -> None:
         # Update URI based on overrides
         self.uri = config[URI]
 
-    def _identifier_to_validated_tuple(self, identifier: Union[str, Identifier]) -> Identifier:
+    def _identifier_to_validated_tuple(self, identifier: str | Identifier) -> Identifier:
         identifier_tuple = self.identifier_to_tuple(identifier)
         if len(identifier_tuple) <= 1:
             raise NoSuchIdentifierError(f"Missing namespace or invalid identifier: {'.'.join(identifier_tuple)}")
         return identifier_tuple
 
     def _split_identifier_for_path(
-        self, identifier: Union[str, Identifier, TableIdentifier], kind: IdentifierKind = IdentifierKind.TABLE
+        self, identifier: str | Identifier | TableIdentifier, kind: IdentifierKind = IdentifierKind.TABLE
     ) -> Properties:
         if isinstance(identifier, TableIdentifier):
             return {"namespace": NAMESPACE_SEPARATOR.join(identifier.namespace.root), kind.value: identifier.name}
         identifier_tuple = self._identifier_to_validated_tuple(identifier)
 
         return {"namespace": NAMESPACE_SEPARATOR.join(identifier_tuple[:-1]), kind.value: identifier_tuple[-1]}
 
-    def _split_identifier_for_json(self, identifier: Union[str, Identifier]) -> Dict[str, Union[Identifier, str]]:
+    def _split_identifier_for_json(self, identifier: str | Identifier) -> Dict[str, Identifier | str]:
         identifier_tuple = self._identifier_to_validated_tuple(identifier)
         return {"namespace": identifier_tuple[:-1], "name": identifier_tuple[-1]}
 
@@ -497,9 +496,9 @@ def _config_headers(self, session: Session) -> None:
 
     def _create_table(
         self,
-        identifier: Union[str, Identifier],
+        identifier: str | Identifier,
         schema: Union[Schema, "pa.Schema"],
-        location: Optional[str] = None,
+        location: str | None = None,
         partition_spec: PartitionSpec = UNPARTITIONED_PARTITION_SPEC,
         sort_order: SortOrder = UNSORTED_SORT_ORDER,
         properties: Properties = EMPTY_DICT,
@@ -539,9 +538,9 @@ def _create_table(
     @retry(**_RETRY_ARGS)
     def create_table(
         self,
-        identifier: Union[str, Identifier],
+        identifier: str | Identifier,
         schema: Union[Schema, "pa.Schema"],
-        location: Optional[str] = None,
+        location: str | None = None,
         partition_spec: PartitionSpec = UNPARTITIONED_PARTITION_SPEC,
         sort_order: SortOrder = UNSORTED_SORT_ORDER,
         properties: Properties = EMPTY_DICT,
@@ -560,9 +559,9 @@ def create_table(
     @retry(**_RETRY_ARGS)
     def create_table_transaction(
         self,
-        identifier: Union[str, Identifier],
+        identifier: str | Identifier,
         schema: Union[Schema, "pa.Schema"],
-        location: Optional[str] = None,
+        location: str | None = None,
         partition_spec: PartitionSpec = UNPARTITIONED_PARTITION_SPEC,
         sort_order: SortOrder = UNSORTED_SORT_ORDER,
         properties: Properties = EMPTY_DICT,
@@ -580,7 +579,7 @@ def create_table_transaction(
         return CreateTableTransaction(staged_table)
 
     @retry(**_RETRY_ARGS)
-    def register_table(self, identifier: Union[str, Identifier], metadata_location: str) -> Table:
+    def register_table(self, identifier: str | Identifier, metadata_location: str) -> Table:
         """Register a new table using existing metadata.
 
         Args:
@@ -612,7 +611,7 @@ def register_table(self, identifier: Union[str, Identifier], metadata_location:
         return self._response_to_table(self.identifier_to_tuple(identifier), table_response)
 
     @retry(**_RETRY_ARGS)
-    def list_tables(self, namespace: Union[str, Identifier]) -> List[Identifier]:
+    def list_tables(self, namespace: str | Identifier) -> List[Identifier]:
         namespace_tuple = self._check_valid_namespace_identifier(namespace)
         namespace_concat = NAMESPACE_SEPARATOR.join(namespace_tuple)
         response = self._session.get(self.url(Endpoints.list_tables, namespace=namespace_concat))
@@ -623,7 +622,7 @@ def list_tables(self, namespace: Union[str, Identifier]) -> List[Identifier]:
         return [(*table.namespace, table.name) for table in ListTablesResponse.model_validate_json(response.text).identifiers]
 
     @retry(**_RETRY_ARGS)
-    def load_table(self, identifier: Union[str, Identifier]) -> Table:
+    def load_table(self, identifier: str | Identifier) -> Table:
         params = {}
         if mode := self.properties.get(SNAPSHOT_LOADING_MODE):
             if mode in {"all", "refs"}:
@@ -643,7 +642,7 @@ def load_table(self, identifier: Union[str, Identifier]) -> Table:
         return self._response_to_table(self.identifier_to_tuple(identifier), table_response)
 
     @retry(**_RETRY_ARGS)
-    def drop_table(self, identifier: Union[str, Identifier], purge_requested: bool = False) -> None:
+    def drop_table(self, identifier: str | Identifier, purge_requested: bool = False) -> None:
         response = self._session.delete(
             self.url(Endpoints.drop_table, prefixed=True, **self._split_identifier_for_path(identifier)),
             params={"purgeRequested": purge_requested},
@@ -654,11 +653,11 @@ def drop_table(self, identifier: Union[str, Identifier], purge_requested: bool =
             _handle_non_200_response(exc, {404: NoSuchTableError})
 
     @retry(**_RETRY_ARGS)
-    def purge_table(self, identifier: Union[str, Identifier]) -> None:
+    def purge_table(self, identifier: str | Identifier) -> None:
         self.drop_table(identifier=identifier, purge_requested=True)
 
     @retry(**_RETRY_ARGS)
-    def rename_table(self, from_identifier: Union[str, Identifier], to_identifier: Union[str, Identifier]) -> Table:
+    def rename_table(self, from_identifier: str | Identifier, to_identifier: str | Identifier) -> Table:
         payload = {
             "source": self._split_identifier_for_json(from_identifier),
             "destination": self._split_identifier_for_json(to_identifier),
@@ -683,7 +682,7 @@ def _remove_catalog_name_from_table_request_identifier(self, table_request: Comm
         return table_request
 
     @retry(**_RETRY_ARGS)
-    def list_views(self, namespace: Union[str, Identifier]) -> List[Identifier]:
+    def list_views(self, namespace: str | Identifier) -> List[Identifier]:
         namespace_tuple = self._check_valid_namespace_identifier(namespace)
         namespace_concat = NAMESPACE_SEPARATOR.join(namespace_tuple)
         response = self._session.get(self.url(Endpoints.list_views, namespace=namespace_concat))
@@ -740,7 +739,7 @@ def commit_table(
         return CommitTableResponse.model_validate_json(response.text)
 
     @retry(**_RETRY_ARGS)
-    def create_namespace(self, namespace: Union[str, Identifier], properties: Properties = EMPTY_DICT) -> None:
+    def create_namespace(self, namespace: str | Identifier, properties: Properties = EMPTY_DICT) -> None:
         namespace_tuple = self._check_valid_namespace_identifier(namespace)
         payload = {"namespace": namespace_tuple, "properties": properties}
         response = self._session.post(self.url(Endpoints.create_namespace), json=payload)
@@ -750,7 +749,7 @@ def create_namespace(self, namespace: Union[str, Identifier], properties: Proper
             _handle_non_200_response(exc, {409: NamespaceAlreadyExistsError})
 
     @retry(**_RETRY_ARGS)
-    def drop_namespace(self, namespace: Union[str, Identifier]) -> None:
+    def drop_namespace(self, namespace: str | Identifier) -> None:
         namespace_tuple = self._check_valid_namespace_identifier(namespace)
         namespace = NAMESPACE_SEPARATOR.join(namespace_tuple)
         response = self._session.delete(self.url(Endpoints.drop_namespace, namespace=namespace))
@@ -760,7 +759,7 @@ def drop_namespace(self, namespace: Union[str, Identifier]) -> None:
             _handle_non_200_response(exc, {404: NoSuchNamespaceError, 409: NamespaceNotEmptyError})
 
     @retry(**_RETRY_ARGS)
-    def list_namespaces(self, namespace: Union[str, Identifier] = ()) -> List[Identifier]:
+    def list_namespaces(self, namespace: str | Identifier = ()) -> List[Identifier]:
         namespace_tuple = self.identifier_to_tuple(namespace)
         response = self._session.get(
             self.url(
@@ -777,7 +776,7 @@ def list_namespaces(self, namespace: Union[str, Identifier] = ()) -> List[Identi
         return ListNamespaceResponse.model_validate_json(response.text).namespaces
 
     @retry(**_RETRY_ARGS)
-    def load_namespace_properties(self, namespace: Union[str, Identifier]) -> Properties:
+    def load_namespace_properties(self, namespace: str | Identifier) -> Properties:
         namespace_tuple = self._check_valid_namespace_identifier(namespace)
         namespace = NAMESPACE_SEPARATOR.join(namespace_tuple)
         response = self._session.get(self.url(Endpoints.load_namespace_metadata, namespace=namespace))
@@ -790,7 +789,7 @@ def load_namespace_properties(self, namespace: Union[str, Identifier]) -> Proper
 
     @retry(**_RETRY_ARGS)
     def update_namespace_properties(
-        self, namespace: Union[str, Identifier], removals: Optional[Set[str]] = None, updates: Properties = EMPTY_DICT
+        self, namespace: str | Identifier, removals: Set[str] | None = None, updates: Properties = EMPTY_DICT
     ) -> PropertiesUpdateSummary:
         namespace_tuple = self._check_valid_namespace_identifier(namespace)
         namespace = NAMESPACE_SEPARATOR.join(namespace_tuple)
@@ -808,7 +807,7 @@ def update_namespace_properties(
         )
 
     @retry(**_RETRY_ARGS)
-    def namespace_exists(self, namespace: Union[str, Identifier]) -> bool:
+    def namespace_exists(self, namespace: str | Identifier) -> bool:
         namespace_tuple = self._check_valid_namespace_identifier(namespace)
         namespace = NAMESPACE_SEPARATOR.join(namespace_tuple)
         response = self._session.head(self.url(Endpoints.namespace_exists, namespace=namespace))
@@ -826,7 +825,7 @@ def namespace_exists(self, namespace: Union[str, Identifier]) -> bool:
         return False
 
     @retry(**_RETRY_ARGS)
-    def table_exists(self, identifier: Union[str, Identifier]) -> bool:
+    def table_exists(self, identifier: str | Identifier) -> bool:
         """Check if a table exists.
 
         Args:
@@ -852,7 +851,7 @@ def table_exists(self, identifier: Union[str, Identifier]) -> bool:
         return False
 
     @retry(**_RETRY_ARGS)
-    def view_exists(self, identifier: Union[str, Identifier]) -> bool:
+    def view_exists(self, identifier: str | Identifier) -> bool:
         """Check if a view exists.
 
         Args:
@@ -877,7 +876,7 @@ def view_exists(self, identifier: Union[str, Identifier]) -> bool:
         return False
 
     @retry(**_RETRY_ARGS)
-    def drop_view(self, identifier: Union[str]) -> None:
+    def drop_view(self, identifier: str) -> None:
         response = self._session.delete(
             self.url(Endpoints.drop_view, prefixed=True, **self._split_identifier_for_path(identifier, IdentifierKind.VIEW)),
         )

tests/catalog/test_rest.py

@@ -18,7 +18,6 @@
 import base64
 import os
 from typing import Any, Callable, Dict, cast
-from copy import deepcopy
 from unittest import mock
 
 import pytest
@@ -40,7 +39,7 @@
     ServerError,
     TableAlreadyExistsError,
 )
-from pyiceberg.io import load_file_io
+from pyiceberg.io import AWS_ACCESS_KEY_ID, load_file_io
 from pyiceberg.partitioning import PartitionField, PartitionSpec
 from pyiceberg.schema import Schema
 from pyiceberg.table import Table
@@ -859,41 +858,51 @@ def test_load_table_200(rest_mock: Mocker, example_table_metadata_with_snapshot_
     assert actual == expected
 
 
-def test_load_table_prefers_storage_credentials_over_config(
-    rest_mock: Mocker, example_table_metadata_with_snapshot_v1_rest_json: Dict[str, Any], monkeypatch: pytest.MonkeyPatch
-) -> None:
-    table_resp = deepcopy(example_table_metadata_with_snapshot_v1_rest_json)
-    table_resp["config"] = {"some.key": "from-config", "only.config": "only-config"}
-    table_resp["storage-credentials"] = {"some.key": "from-cred", "only.creds": "only-creds"}
+def test_storage_credentials_over_config(rest_mock: Mocker, example_table_metadata_with_snapshot_v1_rest_json: Dict[str, Any]) -> None:
+    response = {
+        "metadata-location": example_table_metadata_with_snapshot_v1_rest_json["metadata-location"],
+        "metadata": example_table_metadata_with_snapshot_v1_rest_json["metadata"],
+        "config": {
+            AWS_ACCESS_KEY_ID: "from_config",
+        },
+        "storage-credentials": {
+            AWS_ACCESS_KEY_ID: "from_storage_credentials",
+        },
+    }
 
-    table_resp["metadata"].setdefault("properties", {})["meta.key"] = "from-metadata"
+    rest_mock.get(
+        f"{TEST_URI}v1/namespaces/fokko/tables/table",
+        json=response,
+        status_code=200,
+        request_headers=TEST_HEADERS,
+    )
+
+    catalog = RestCatalog("rest", uri=TEST_URI, token=TEST_TOKEN)
+    table = catalog.load_table(("fokko", "table"))
 
-    captured: Dict[str, Any] = {}
+    assert table.io.properties[AWS_ACCESS_KEY_ID] == "from_storage_credentials"
 
-    def _capture_io(_self: Any, properties: Dict[str, Any], location: Any) -> Any:  # type: ignore
-        captured["properties"] = dict(properties)
-        captured["location"] = location
-        class _DummyIO:
-            pass
-        return _DummyIO()
 
-    monkeypatch.setattr(RestCatalog, "_load_file_io", _capture_io, raising=True)
+def test_config_when_no_storage_credentials(rest_mock: Mocker, example_table_metadata_with_snapshot_v1_rest_json: Dict[str, Any]) -> None:
+    response = {
+        "metadata-location": example_table_metadata_with_snapshot_v1_rest_json["metadata-location"],
+        "metadata": example_table_metadata_with_snapshot_v1_rest_json["metadata"],
+        "config": {
+            AWS_ACCESS_KEY_ID: "from_config",
+        },
+    }
 
     rest_mock.get(
         f"{TEST_URI}v1/namespaces/fokko/tables/table",
-        json=table_resp,
+        json=response,
         status_code=200,
         request_headers=TEST_HEADERS,
     )
 
     catalog = RestCatalog("rest", uri=TEST_URI, token=TEST_TOKEN)
-    _ = catalog.load_table(("fokko", "table"))
-
-    props = cast(Dict[str, Any], captured["properties"])
-    assert props["meta.key"] == "from-metadata"
-    assert props["some.key"] == "from-cred"
-    assert props["only.creds"] == "only-creds"
+    table = catalog.load_table(("fokko", "table"))
 
+    assert table.io.properties[AWS_ACCESS_KEY_ID] == "from_config"
 
 def test_load_table_200_loading_mode(
     rest_mock: Mocker, example_table_metadata_with_snapshot_v1_rest_json: Dict[str, Any]