Add REST load credentials support

Author: kevinjqliu

pyiceberg/catalog/rest/__init__.py

@@ -17,6 +17,7 @@
 from __future__ import annotations
 
 from collections import deque
+from collections.abc import Sequence
 from enum import Enum
 from typing import (
     TYPE_CHECKING,
@@ -147,6 +148,7 @@ class Endpoints:
     create_table: str = "namespaces/{namespace}/tables"
     register_table: str = "namespaces/{namespace}/register"
     load_table: str = "namespaces/{namespace}/tables/{table}"
+    load_credentials: str = "namespaces/{namespace}/tables/{table}/credentials"
     update_table: str = "namespaces/{namespace}/tables/{table}"
     drop_table: str = "namespaces/{namespace}/tables/{table}"
     table_exists: str = "namespaces/{namespace}/tables/{table}"
@@ -181,6 +183,7 @@ class Capability:
     V1_DELETE_TABLE = Endpoint(http_method=HttpMethod.DELETE, path=f"{API_PREFIX}/{Endpoints.drop_table}")
     V1_RENAME_TABLE = Endpoint(http_method=HttpMethod.POST, path=f"{API_PREFIX}/{Endpoints.rename_table}")
     V1_REGISTER_TABLE = Endpoint(http_method=HttpMethod.POST, path=f"{API_PREFIX}/{Endpoints.register_table}")
+    V1_LOAD_CREDENTIALS = Endpoint(http_method=HttpMethod.GET, path=f"{API_PREFIX}/{Endpoints.load_credentials}")
 
     V1_LIST_VIEWS = Endpoint(http_method=HttpMethod.GET, path=f"{API_PREFIX}/{Endpoints.list_views}")
     V1_LOAD_VIEW = Endpoint(http_method=HttpMethod.GET, path=f"{API_PREFIX}/{Endpoints.load_view}")
@@ -293,6 +296,10 @@ class TableResponse(IcebergBaseModel):
     storage_credentials: list[StorageCredential] = Field(alias="storage-credentials", default_factory=list)
 
 
+class LoadCredentialsResponse(IcebergBaseModel):
+    storage_credentials: list[StorageCredential] = Field(alias="storage-credentials")
+
+
 class ViewResponse(IcebergBaseModel):
     metadata_location: str | None = Field(alias="metadata-location", default=None)
     metadata: ViewMetadata
@@ -480,6 +487,12 @@ def _resolve_storage_credentials(storage_credentials: list[StorageCredential], l
 
         return best_match.config if best_match else {}
 
+    @staticmethod
+    def _format_referenced_by(referenced_by: str | Sequence[str]) -> str:
+        if isinstance(referenced_by, str):
+            return referenced_by
+        return ",".join(referenced_by)
+
     def _load_file_io(self, properties: Properties = EMPTY_DICT, location: str | None = None) -> FileIO:
         merged_properties = {**self.properties, **properties}
         if self._auth_manager:
@@ -545,6 +558,43 @@ def _fetch_scan_tasks(self, identifier: str | Identifier, plan_task: str) -> Sca
 
         return ScanTasks.model_validate_json(response.text)
 
+    @retry(**_RETRY_ARGS)
+    def _load_credentials(
+        self,
+        identifier: str | Identifier,
+        plan_id: str | None = None,
+        referenced_by: str | Sequence[str] | None = None,
+    ) -> LoadCredentialsResponse:
+        """Load raw vended storage credentials for a table."""
+        self._check_endpoint(Capability.V1_LOAD_CREDENTIALS)
+        params: dict[str, str] = {}
+        if plan_id is not None:
+            params["planId"] = plan_id
+        if referenced_by is not None:
+            params["referenced-by"] = self._format_referenced_by(referenced_by)
+
+        response = self._session.get(
+            self.url(Endpoints.load_credentials, prefixed=True, **self._split_identifier_for_path(identifier)),
+            params=params,
+        )
+        try:
+            response.raise_for_status()
+        except HTTPError as exc:
+            _handle_non_200_response(exc, {404: NoSuchTableError})
+
+        return LoadCredentialsResponse.model_validate_json(response.text)
+
+    def load_credentials(
+        self,
+        identifier: str | Identifier,
+        location: str,
+        plan_id: str | None = None,
+        referenced_by: str | Sequence[str] | None = None,
+    ) -> Properties:
+        """Load vended storage credentials and return the best match for a location."""
+        credentials_response = self._load_credentials(identifier, plan_id=plan_id, referenced_by=referenced_by)
+        return self._resolve_storage_credentials(credentials_response.storage_credentials, location)
+
     def plan_scan(self, identifier: str | Identifier, request: PlanTableScanRequest) -> list[FileScanTask]:
         """Plan a table scan and return FileScanTasks.
 

tests/catalog/test_rest.py

@@ -104,6 +104,7 @@
     Capability.V1_DELETE_TABLE,
     Capability.V1_RENAME_TABLE,
     Capability.V1_REGISTER_TABLE,
+    Capability.V1_LOAD_CREDENTIALS,
     Capability.V1_LIST_VIEWS,
     Capability.V1_LOAD_VIEW,
     Capability.V1_VIEW_EXISTS,
@@ -3147,6 +3148,43 @@ def test_load_table_with_storage_credentials(rest_mock: Mocker, example_table_me
     assert table.io.properties["s3.session-token"] == "vended-token"
 
 
+def test_load_credentials_with_longest_prefix_and_query_params(rest_mock: Mocker) -> None:
+    rest_mock.get(
+        f"{TEST_URI}v1/namespaces/fokko/tables/table/credentials",
+        json={
+            "storage-credentials": [
+                {
+                    "prefix": "s3://warehouse/database/",
+                    "config": {"s3.access-key-id": "short-prefix-key"},
+                },
+                {
+                    "prefix": "s3://warehouse/database/table",
+                    "config": {
+                        "s3.access-key-id": "long-prefix-key",
+                        "s3.secret-access-key": "long-prefix-secret",
+                    },
+                },
+            ],
+        },
+        status_code=200,
+        request_headers=TEST_HEADERS,
+    )
+    catalog = RestCatalog("rest", uri=TEST_URI, token=TEST_TOKEN)
+
+    credentials = catalog.load_credentials(
+        ("fokko", "table"),
+        "s3://warehouse/database/table/data/file.parquet",
+        plan_id="plan-123",
+        referenced_by=["prod.view", "prod.inner_view"],
+    )
+
+    assert credentials == {"s3.access-key-id": "long-prefix-key", "s3.secret-access-key": "long-prefix-secret"}
+    assert (
+        rest_mock.last_request.url == f"{TEST_URI}v1/namespaces/fokko/tables/table/credentials?planId=plan-123"
+        "&referenced-by=prod.view%2Cprod.inner_view"
+    )
+
+
 def test_load_table_without_storage_credentials(
     rest_mock: Mocker, example_table_metadata_with_snapshot_v1_rest_json: dict[str, Any]
 ) -> None: