Skip to content

Commit a26a8c3

Browse files
antonlin1claude
andcommitted
feat: support storage-credentials in loadTable response
Per the Iceberg REST spec, the loadTable response can include a storage-credentials field with prefix-scoped temporary credentials. Clients must check storage-credentials before falling back to the config dict. This adds longest-prefix matching for credential selection. Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
1 parent 7d4a8ef commit a26a8c3

2 files changed

Lines changed: 111 additions & 2 deletions

File tree

pyiceberg/catalog/rest/__init__.py

Lines changed: 40 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -40,6 +40,7 @@
4040
PlanSubmitted,
4141
PlanTableScanRequest,
4242
ScanTasks,
43+
StorageCredential,
4344
)
4445
from pyiceberg.exceptions import (
4546
AuthorizationExpiredError,
@@ -256,6 +257,7 @@ class TableResponse(IcebergBaseModel):
256257
metadata_location: str | None = Field(alias="metadata-location", default=None)
257258
metadata: TableMetadata
258259
config: Properties = Field(default_factory=dict)
260+
storage_credentials: list[StorageCredential] | None = Field(alias="storage-credentials", default=None)
259261

260262

261263
class CreateTableRequest(IcebergBaseModel):
@@ -728,13 +730,40 @@ def add_headers(self, request: PreparedRequest, **kwargs: Any) -> None: # pylin
728730

729731
session.mount(self.uri, SigV4Adapter(**self.properties))
730732

733+
@staticmethod
734+
def _get_credentials(
735+
storage_credentials: list[StorageCredential] | None,
736+
config: Properties,
737+
metadata_location: str | None,
738+
table_location: str | None,
739+
) -> Properties:
740+
if not storage_credentials:
741+
return config
742+
target = metadata_location or table_location
743+
if not target:
744+
return config
745+
matching = [sc for sc in storage_credentials if target.startswith(sc.prefix)]
746+
if not matching:
747+
return config
748+
selected = max(matching, key=lambda sc: len(sc.prefix))
749+
return selected.config
750+
731751
def _response_to_table(self, identifier_tuple: tuple[str, ...], table_response: TableResponse) -> Table:
732752
return Table(
733753
identifier=identifier_tuple,
734754
metadata_location=table_response.metadata_location, # type: ignore
735755
metadata=table_response.metadata,
736756
io=self._load_file_io(
737-
{**table_response.metadata.properties, **table_response.config}, table_response.metadata_location
757+
{
758+
**table_response.metadata.properties,
759+
**self._get_credentials(
760+
table_response.storage_credentials,
761+
table_response.config,
762+
table_response.metadata_location,
763+
getattr(table_response.metadata, "location", None),
764+
),
765+
},
766+
table_response.metadata_location,
738767
),
739768
catalog=self,
740769
config=table_response.config,
@@ -746,7 +775,16 @@ def _response_to_staged_table(self, identifier_tuple: tuple[str, ...], table_res
746775
metadata_location=table_response.metadata_location, # type: ignore
747776
metadata=table_response.metadata,
748777
io=self._load_file_io(
749-
{**table_response.metadata.properties, **table_response.config}, table_response.metadata_location
778+
{
779+
**table_response.metadata.properties,
780+
**self._get_credentials(
781+
table_response.storage_credentials,
782+
table_response.config,
783+
table_response.metadata_location,
784+
getattr(table_response.metadata, "location", None),
785+
),
786+
},
787+
table_response.metadata_location,
750788
),
751789
catalog=self,
752790
)

tests/catalog/test_rest.py

Lines changed: 71 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -2351,3 +2351,74 @@ def test_table_uuid_check_on_refresh(rest_mock: Mocker, example_table_metadata_v
23512351
assert "Table UUID does not match" in str(exc_info.value)
23522352
assert f"current={original_uuid}" in str(exc_info.value)
23532353
assert f"refreshed={different_uuid}" in str(exc_info.value)
2354+
2355+
2356+
def test_storage_credentials_over_config(
2357+
rest_mock: Mocker, example_table_metadata_with_snapshot_v1_rest_json: dict[str, Any]
2358+
) -> None:
2359+
response_with_storage_creds = {
2360+
**example_table_metadata_with_snapshot_v1_rest_json,
2361+
"storage-credentials": [
2362+
{
2363+
"prefix": "s3://warehouse/",
2364+
"config": {
2365+
"s3.access-key-id": "storage-cred-key",
2366+
"s3.secret-access-key": "storage-cred-secret",
2367+
},
2368+
}
2369+
],
2370+
}
2371+
rest_mock.get(
2372+
f"{TEST_URI}v1/namespaces/fokko/tables/table",
2373+
json=response_with_storage_creds,
2374+
status_code=200,
2375+
request_headers=TEST_HEADERS,
2376+
)
2377+
catalog = RestCatalog("rest", uri=TEST_URI, token=TEST_TOKEN)
2378+
table = catalog.load_table(("fokko", "table"))
2379+
assert table.io.properties["s3.access-key-id"] == "storage-cred-key"
2380+
assert table.io.properties["s3.secret-access-key"] == "storage-cred-secret"
2381+
2382+
2383+
def test_config_when_no_storage_credentials(
2384+
rest_mock: Mocker, example_table_metadata_with_snapshot_v1_rest_json: dict[str, Any]
2385+
) -> None:
2386+
rest_mock.get(
2387+
f"{TEST_URI}v1/namespaces/fokko/tables/table",
2388+
json=example_table_metadata_with_snapshot_v1_rest_json,
2389+
status_code=200,
2390+
request_headers=TEST_HEADERS,
2391+
)
2392+
catalog = RestCatalog("rest", uri=TEST_URI, token=TEST_TOKEN)
2393+
table = catalog.load_table(("fokko", "table"))
2394+
# config from the fixture should be used since there are no storage-credentials
2395+
assert table.io.properties["region"] == "us-west-2"
2396+
2397+
2398+
def test_storage_credentials_no_prefix_match() -> None:
2399+
from pyiceberg.catalog.rest.scan_planning import StorageCredential
2400+
2401+
creds = [StorageCredential(prefix="s3://other-bucket/", config={"key": "val"})]
2402+
result = RestCatalog._get_credentials(
2403+
storage_credentials=creds,
2404+
config={"fallback-key": "fallback-val"},
2405+
metadata_location="s3://warehouse/database/table/metadata/file.json",
2406+
table_location="s3://warehouse/database/table",
2407+
)
2408+
assert result == {"fallback-key": "fallback-val"}
2409+
2410+
2411+
def test_storage_credentials_longest_prefix_wins() -> None:
2412+
from pyiceberg.catalog.rest.scan_planning import StorageCredential
2413+
2414+
creds = [
2415+
StorageCredential(prefix="s3://warehouse/", config={"key": "short-prefix"}),
2416+
StorageCredential(prefix="s3://warehouse/database/table/", config={"key": "long-prefix"}),
2417+
]
2418+
result = RestCatalog._get_credentials(
2419+
storage_credentials=creds,
2420+
config={"key": "fallback"},
2421+
metadata_location="s3://warehouse/database/table/metadata/file.json",
2422+
table_location="s3://warehouse/database/table",
2423+
)
2424+
assert result == {"key": "long-prefix"}

0 commit comments

Comments
 (0)