Skip to content

Commit b8e5a4b

Browse files
authored
Add snapshot-loading-mode option to RESTCatalog (#1998)
<!-- Thanks for opening a pull request! --> <!-- In the case this PR will resolve an issue, please replace ${GITHUB_ISSUE_ID} below with the actual Github issue id. --> <!-- Closes #${GITHUB_ISSUE_ID} --> # Rationale for this change This allows users only to fetch the snapshots that are referenced by a tag or a branch. # Are these changes tested? # Are there any user-facing changes? <!-- In the case of user-facing changes, please add the changelog label. -->
1 parent f0e49c6 commit b8e5a4b

File tree

3 files changed

+36
-2
lines changed

3 files changed

+36
-2
lines changed

mkdocs/docs/configuration.md

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -332,6 +332,7 @@ catalog:
332332
| rest.signing-region | us-east-1 | The region to use when SigV4 signing a request |
333333
| rest.signing-name | execute-api | The service signing name to use when SigV4 signing a request |
334334
| oauth2-server-uri | <https://auth-service/cc> | Authentication URL to use for client credentials authentication (default: uri + 'v1/oauth/tokens') |
335+
| snapshot-loading-mode | refs | The snapshots to return in the body of the metadata. Setting the value to `all` would return the full set of snapshots currently valid for the table. Setting the value to `refs` would load all snapshots referenced by branches or tags. |
335336

336337
<!-- markdown-link-check-enable-->
337338

pyiceberg/catalog/rest/__init__.py

Lines changed: 11 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -138,6 +138,7 @@ class IdentifierKind(Enum):
138138
SIGV4_REGION = "rest.signing-region"
139139
SIGV4_SERVICE = "rest.signing-name"
140140
OAUTH2_SERVER_URI = "oauth2-server-uri"
141+
SNAPSHOT_LOADING_MODE = "snapshot-loading-mode"
141142

142143
NAMESPACE_SEPARATOR = b"\x1f".decode(UTF8)
143144

@@ -678,7 +679,16 @@ def list_tables(self, namespace: Union[str, Identifier]) -> List[Identifier]:
678679

679680
@retry(**_RETRY_ARGS)
680681
def load_table(self, identifier: Union[str, Identifier]) -> Table:
681-
response = self._session.get(self.url(Endpoints.load_table, prefixed=True, **self._split_identifier_for_path(identifier)))
682+
params = {}
683+
if mode := self.properties.get(SNAPSHOT_LOADING_MODE):
684+
if mode in {"all", "refs"}:
685+
params["snapshots"] = mode
686+
else:
687+
raise ValueError("Invalid snapshot-loading-mode: {}")
688+
689+
response = self._session.get(
690+
self.url(Endpoints.load_table, prefixed=True, **self._split_identifier_for_path(identifier)), params=params
691+
)
682692
try:
683693
response.raise_for_status()
684694
except HTTPError as exc:

tests/catalog/test_rest.py

Lines changed: 24 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -24,7 +24,7 @@
2424

2525
import pyiceberg
2626
from pyiceberg.catalog import PropertiesUpdateSummary, load_catalog
27-
from pyiceberg.catalog.rest import OAUTH2_SERVER_URI, RestCatalog
27+
from pyiceberg.catalog.rest import OAUTH2_SERVER_URI, SNAPSHOT_LOADING_MODE, RestCatalog
2828
from pyiceberg.exceptions import (
2929
AuthorizationExpiredError,
3030
NamespaceAlreadyExistsError,
@@ -853,6 +853,29 @@ def test_load_table_200(rest_mock: Mocker, example_table_metadata_with_snapshot_
853853
assert actual == expected
854854

855855

856+
def test_load_table_200_loading_mode(
857+
rest_mock: Mocker, example_table_metadata_with_snapshot_v1_rest_json: Dict[str, Any]
858+
) -> None:
859+
rest_mock.get(
860+
f"{TEST_URI}v1/namespaces/fokko/tables/table?snapshots=refs",
861+
json=example_table_metadata_with_snapshot_v1_rest_json,
862+
status_code=200,
863+
request_headers=TEST_HEADERS,
864+
)
865+
catalog = RestCatalog("rest", uri=TEST_URI, token=TEST_TOKEN, **{SNAPSHOT_LOADING_MODE: "refs"})
866+
actual = catalog.load_table(("fokko", "table"))
867+
expected = Table(
868+
identifier=("fokko", "table"),
869+
metadata_location=example_table_metadata_with_snapshot_v1_rest_json["metadata-location"],
870+
metadata=TableMetadataV1(**example_table_metadata_with_snapshot_v1_rest_json["metadata"]),
871+
io=load_file_io(),
872+
catalog=catalog,
873+
)
874+
# First compare the dicts
875+
assert actual.metadata.model_dump() == expected.metadata.model_dump()
876+
assert actual == expected
877+
878+
856879
def test_load_table_honor_access_delegation(
857880
rest_mock: Mocker, example_table_metadata_with_snapshot_v1_rest_json: Dict[str, Any]
858881
) -> None:

0 commit comments

Comments
 (0)