Skip to content

Commit a38bbe3

Browse files
authored
feat: Add REST loadCredentials support (#3499)
## Summary - add REST `loadCredentials` endpoint support - parse `LoadCredentialsResponse` - expose `RestCatalog.load_credentials(...)` with longest-prefix resolution for a target location Contributes to #3495. ## Notes The pinned `apache/iceberg-rest-fixture:1.10.1` used by integration tests does not serve the `/credentials` route, so this is covered with REST catalog unit tests rather than fixture-backed integration tests. ## Testing - `PYTHONPATH=. pytest tests/catalog/test_rest.py -k 'load_credentials or storage_credentials'` - `ruff check pyiceberg/catalog/rest/__init__.py tests/catalog/test_rest.py` - `ruff format --check pyiceberg/catalog/rest/__init__.py tests/catalog/test_rest.py`
1 parent 3a3bc62 commit a38bbe3

2 files changed

Lines changed: 65 additions & 0 deletions

File tree

pyiceberg/catalog/rest/__init__.py

Lines changed: 32 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -147,6 +147,7 @@ class Endpoints:
147147
create_table: str = "namespaces/{namespace}/tables"
148148
register_table: str = "namespaces/{namespace}/register"
149149
load_table: str = "namespaces/{namespace}/tables/{table}"
150+
load_credentials: str = "namespaces/{namespace}/tables/{table}/credentials"
150151
update_table: str = "namespaces/{namespace}/tables/{table}"
151152
drop_table: str = "namespaces/{namespace}/tables/{table}"
152153
table_exists: str = "namespaces/{namespace}/tables/{table}"
@@ -181,6 +182,7 @@ class Capability:
181182
V1_DELETE_TABLE = Endpoint(http_method=HttpMethod.DELETE, path=f"{API_PREFIX}/{Endpoints.drop_table}")
182183
V1_RENAME_TABLE = Endpoint(http_method=HttpMethod.POST, path=f"{API_PREFIX}/{Endpoints.rename_table}")
183184
V1_REGISTER_TABLE = Endpoint(http_method=HttpMethod.POST, path=f"{API_PREFIX}/{Endpoints.register_table}")
185+
V1_LOAD_CREDENTIALS = Endpoint(http_method=HttpMethod.GET, path=f"{API_PREFIX}/{Endpoints.load_credentials}")
184186

185187
V1_LIST_VIEWS = Endpoint(http_method=HttpMethod.GET, path=f"{API_PREFIX}/{Endpoints.list_views}")
186188
V1_LOAD_VIEW = Endpoint(http_method=HttpMethod.GET, path=f"{API_PREFIX}/{Endpoints.load_view}")
@@ -293,6 +295,10 @@ class TableResponse(IcebergBaseModel):
293295
storage_credentials: list[StorageCredential] = Field(alias="storage-credentials", default_factory=list)
294296

295297

298+
class LoadCredentialsResponse(IcebergBaseModel):
299+
storage_credentials: list[StorageCredential] = Field(alias="storage-credentials")
300+
301+
296302
class ViewResponse(IcebergBaseModel):
297303
metadata_location: str | None = Field(alias="metadata-location", default=None)
298304
metadata: ViewMetadata
@@ -1094,6 +1100,32 @@ def load_table(self, identifier: str | Identifier) -> Table:
10941100
table_response = TableResponse.model_validate_json(response.text)
10951101
return self._response_to_table(self.identifier_to_tuple(identifier), table_response)
10961102

1103+
@retry(**_RETRY_ARGS)
1104+
def _load_credentials(
1105+
self,
1106+
identifier: str | Identifier,
1107+
) -> LoadCredentialsResponse:
1108+
"""Load raw vended storage credentials for a table."""
1109+
self._check_endpoint(Capability.V1_LOAD_CREDENTIALS)
1110+
response = self._session.get(
1111+
self.url(Endpoints.load_credentials, prefixed=True, **self._split_identifier_for_path(identifier)),
1112+
)
1113+
try:
1114+
response.raise_for_status()
1115+
except HTTPError as exc:
1116+
_handle_non_200_response(exc, {404: NoSuchTableError})
1117+
1118+
return LoadCredentialsResponse.model_validate_json(response.text)
1119+
1120+
def load_credentials(
1121+
self,
1122+
identifier: str | Identifier,
1123+
location: str,
1124+
) -> Properties:
1125+
"""Load vended storage credentials and return the best match for a location."""
1126+
credentials_response = self._load_credentials(identifier)
1127+
return self._resolve_storage_credentials(credentials_response.storage_credentials, location)
1128+
10971129
@retry(**_RETRY_ARGS)
10981130
@override
10991131
def drop_table(self, identifier: str | Identifier, purge_requested: bool = False) -> None:

tests/catalog/test_rest.py

Lines changed: 33 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -104,6 +104,7 @@
104104
Capability.V1_DELETE_TABLE,
105105
Capability.V1_RENAME_TABLE,
106106
Capability.V1_REGISTER_TABLE,
107+
Capability.V1_LOAD_CREDENTIALS,
107108
Capability.V1_LIST_VIEWS,
108109
Capability.V1_LOAD_VIEW,
109110
Capability.V1_VIEW_EXISTS,
@@ -3143,6 +3144,38 @@ def test_load_table_with_storage_credentials(rest_mock: Mocker, example_table_me
31433144
assert table.io.properties["s3.session-token"] == "vended-token"
31443145

31453146

3147+
def test_load_credentials_with_longest_prefix(rest_mock: Mocker) -> None:
3148+
rest_mock.get(
3149+
f"{TEST_URI}v1/namespaces/fokko/tables/table/credentials",
3150+
json={
3151+
"storage-credentials": [
3152+
{
3153+
"prefix": "s3://warehouse/database/",
3154+
"config": {"s3.access-key-id": "short-prefix-key"},
3155+
},
3156+
{
3157+
"prefix": "s3://warehouse/database/table",
3158+
"config": {
3159+
"s3.access-key-id": "long-prefix-key",
3160+
"s3.secret-access-key": "long-prefix-secret",
3161+
},
3162+
},
3163+
],
3164+
},
3165+
status_code=200,
3166+
request_headers=TEST_HEADERS,
3167+
)
3168+
catalog = RestCatalog("rest", uri=TEST_URI, token=TEST_TOKEN)
3169+
3170+
credentials = catalog.load_credentials(
3171+
("fokko", "table"),
3172+
"s3://warehouse/database/table/data/file.parquet",
3173+
)
3174+
3175+
assert credentials == {"s3.access-key-id": "long-prefix-key", "s3.secret-access-key": "long-prefix-secret"}
3176+
assert rest_mock.last_request.url == f"{TEST_URI}v1/namespaces/fokko/tables/table/credentials"
3177+
3178+
31463179
def test_load_table_without_storage_credentials(
31473180
rest_mock: Mocker, example_table_metadata_with_snapshot_v1_rest_json: dict[str, Any]
31483181
) -> None:

0 commit comments

Comments
 (0)