feat: support freshness-aware table loading in RestCatalog

mkdocs/docs/configuration.md

@@ -406,6 +406,8 @@ catalog:
 | uri                 | <https://rest-catalog/ws>        | URI identifying the REST Server                                                                    |
 | warehouse           | myWarehouse                      | Warehouse location or identifier to request from the catalog service. May be used to determine server-side overrides, such as the warehouse location. |
 | 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. |
+| rest-table-cache.max-entries | 100                        | Maximum number of loaded tables kept to answer `304 Not Modified` on reload, keyed by the table's `ETag`. Loads that vend short-lived storage credentials are never cached. Set to `0` to disable (default: 100). |
+| rest-table-cache.expire-after-write-ms | 300000           | How long a loaded table stays cached for reload. Must be positive (default: 5 minutes). |
 | `header.X-Iceberg-Access-Delegation` | `vended-credentials` | Signal to the server that the client supports delegated access via a comma-separated list of access mechanisms. The server may choose to supply access via any or none of the requested mechanisms. When using `vended-credentials`, the server provides temporary credentials to the client. When using `remote-signing`, the server signs requests on behalf of the client. (default: `vended-credentials`) |
 | view-endpoints-supported | false                           | For backwards compatibility with older REST servers. Set to `true` if the server supports view endpoints but doesn't send the `endpoints` field in the ConfigResponse. |
 | scan-planning-mode | client | When set to `server`, and the catalog advertises the plan-table-scan endpoint, `table.scan()` uses REST server-side scan planning. May be set by the client, returned by the catalog via `GET /v1/config` defaults/overrides, or returned per table in the `config` of the `loadTable` response. The `loadTable` value takes precedence over the catalog-level value, which lets a server request server-side planning only for specific tables. Async plans (`status=submitted`) are polled via `GET .../plan/{plan-id}` until terminal state (completed / failed / cancelled). |

pyiceberg/catalog/rest/__init__.py

@@ -17,6 +17,7 @@
 from __future__ import annotations
 
 import logging
+import threading
 import time
 from collections import deque
 from collections.abc import Mapping
@@ -27,6 +28,7 @@
 )
 from urllib.parse import quote, unquote
 
+from cachetools import TTLCache
 from pydantic import ConfigDict, Field, TypeAdapter, field_validator
 from requests import HTTPError, PreparedRequest, Response, Session
 from requests.adapters import DEFAULT_RETRIES, HTTPAdapter
@@ -63,16 +65,20 @@
     NoSuchTableError,
     NoSuchViewError,
     RemotePlanTimeoutError,
+    RESTError,
     TableAlreadyExistsError,
     UnauthorizedError,
     ViewAlreadyExistsError,
 )
 from pyiceberg.io import (
+    ADLS_SAS_TOKEN,
     AWS_ACCESS_KEY_ID,
     AWS_PROFILE_NAME,
     AWS_REGION,
     AWS_SECRET_ACCESS_KEY,
     AWS_SESSION_TOKEN,
+    GCS_TOKEN,
+    S3_SESSION_TOKEN,
     FileIO,
     load_file_io,
 )
@@ -298,6 +304,10 @@ class ScanPlanningMode(Enum):
 EMPTY_BODY_SHA256: str = "e3b0c44298fc1c149afbf4c8996fb92427ae41e4649b934ca495991b7852b855"
 OAUTH2_SERVER_URI = "oauth2-server-uri"
 SNAPSHOT_LOADING_MODE = "snapshot-loading-mode"
+REST_TABLE_CACHE_MAX_ENTRIES = "rest-table-cache.max-entries"
+REST_TABLE_CACHE_MAX_ENTRIES_DEFAULT = 100
+REST_TABLE_CACHE_EXPIRE_AFTER_WRITE_MS = "rest-table-cache.expire-after-write-ms"
+REST_TABLE_CACHE_EXPIRE_AFTER_WRITE_MS_DEFAULT = 5 * 60 * 1000  # 5 minutes, matches Java
 AUTH = "auth"
 CUSTOM = "custom"
 SCAN_PLANNING_MODE = "scan-planning-mode"
@@ -336,6 +346,12 @@ def _parse_scan_planning_mode(properties: Properties, *, strict: bool = True) ->
         return None
 
 
+def _is_short_lived_credential(key: str) -> bool:
+    """Whether a loadTable config key carries a short-lived storage credential."""
+    key = key.lower()
+    return key in (S3_SESSION_TOKEN, GCS_TOKEN) or key.startswith(ADLS_SAS_TOKEN)
+
+
 def _retry_hook(retry_state: RetryCallState) -> None:
     rest_catalog: RestCatalog = retry_state.args[0]
     rest_catalog._refresh_token()  # pylint: disable=protected-access
@@ -554,6 +570,8 @@ class RestCatalog(Catalog):
     _auth_manager: AuthManager | None
     _supported_endpoints: set[Endpoint]
     _namespace_separator: str
+    # Loads that vend short-lived storage credentials are not cached, since a 304 cannot refresh them.
+    _table_cache: TTLCache[Identifier, tuple[str, TableResponse]] | None
 
     def __init__(self, name: str, **properties: str):
         """Rest Catalog.
@@ -569,6 +587,17 @@ def __init__(self, name: str, **properties: str):
         self.uri = properties[URI]
         self._fetch_config()
         self._session = self._create_session()
+        max_entries = int(self.properties.get(REST_TABLE_CACHE_MAX_ENTRIES, REST_TABLE_CACHE_MAX_ENTRIES_DEFAULT))
+        ttl_ms = int(self.properties.get(REST_TABLE_CACHE_EXPIRE_AFTER_WRITE_MS, REST_TABLE_CACHE_EXPIRE_AFTER_WRITE_MS_DEFAULT))
+        if ttl_ms <= 0:
+            raise ValueError(f"Invalid value for {REST_TABLE_CACHE_EXPIRE_AFTER_WRITE_MS}: {ttl_ms} (must be positive)")
+        self._table_cache = TTLCache(maxsize=max_entries, ttl=ttl_ms / 1000) if max_entries > 0 else None
+        self._table_cache_lock = threading.Lock()
+
+    def _invalidate_cached_table(self, identifier: str | Identifier) -> None:
+        if self._table_cache is not None:
+            with self._table_cache_lock:
+                self._table_cache.pop(self.identifier_to_tuple(identifier), None)
 
     def _create_session(self) -> Session:
         """Create a request session with provided catalog configuration."""
@@ -1356,6 +1385,7 @@ def register_table(self, identifier: str | Identifier, metadata_location: str, o
             overwrite=overwrite,
         )
         serialized_json = request.model_dump_json().encode(UTF8)
+        self._invalidate_cached_table(identifier)
         response = self._session.post(
             self.url(Endpoints.register_table, namespace=namespace_and_table["namespace"]),
             data=serialized_json,
@@ -1441,16 +1471,37 @@ def load_table(self, identifier: str | Identifier) -> Table:
             else:
                 raise ValueError("Invalid snapshot-loading-mode: {}")
 
+        identifier_tuple = self.identifier_to_tuple(identifier)
+        with self._table_cache_lock:
+            cached = self._table_cache.get(identifier_tuple) if self._table_cache is not None else None
         response = self._session.get(
-            self.url(Endpoints.load_table, prefixed=True, **self._split_identifier_for_path(identifier)), params=params
+            self.url(Endpoints.load_table, prefixed=True, **self._split_identifier_for_path(identifier)),
+            params=params,
+            headers={"If-None-Match": cached[0]} if cached else None,
         )
+        if response.status_code == 304:
+            if cached is None:
+                raise RESTError("Received 304 Not Modified for a table load without If-None-Match")
+            return self._response_to_table(identifier_tuple, cached[1])
         try:
             response.raise_for_status()
         except HTTPError as exc:
+            if response.status_code == 404:
+                self._invalidate_cached_table(identifier)
             _handle_non_200_response(exc, {404: NoSuchTableError})
 
         table_response = TableResponse.model_validate_json(response.text)
-        return self._response_to_table(self.identifier_to_tuple(identifier), table_response)
+        if self._table_cache is not None:
+            etag = response.headers.get("ETag")
+            vends_credentials = bool(table_response.storage_credentials) or any(
+                map(_is_short_lived_credential, table_response.config)
+            )
+            with self._table_cache_lock:
+                if etag and not vends_credentials:
+                    self._table_cache[identifier_tuple] = (etag, table_response)
+                else:
+                    self._table_cache.pop(identifier_tuple, None)
+        return self._response_to_table(identifier_tuple, table_response)
 
     @retry(**_RETRY_ARGS)
     def _load_credentials(
@@ -1482,6 +1533,7 @@ def load_credentials(
     @override
     def drop_table(self, identifier: str | Identifier, purge_requested: bool = False) -> None:
         self._check_endpoint(Capability.V1_DELETE_TABLE)
+        self._invalidate_cached_table(identifier)
         response = self._session.delete(
             self.url(Endpoints.drop_table, prefixed=True, **self._split_identifier_for_path(identifier)),
             params={"purgeRequested": purge_requested},
@@ -1514,6 +1566,7 @@ def rename_table(self, from_identifier: str | Identifier, to_identifier: str | I
         if not self.namespace_exists(destination_namespace):
             raise NoSuchNamespaceError(f"Destination namespace does not exist: {destination_namespace}")
 
+        self._invalidate_cached_table(from_identifier)
         response = self._session.post(self.url(Endpoints.rename_table), json=payload)
         try:
             response.raise_for_status()
@@ -1615,6 +1668,7 @@ def commit_table(
         if table_token := table.config.get(TOKEN):
             headers[AUTHORIZATION_HEADER] = f"{BEARER_PREFIX} {table_token}"
 
+        self._invalidate_cached_table(identifier)
         response = self._session.post(
             self.url(Endpoints.update_table, prefixed=True, **self._split_identifier_for_path(table_request.identifier)),
             data=table_request.model_dump_json().encode(UTF8),

tests/catalog/test_rest.py

@@ -27,7 +27,7 @@
 import pytest
 from requests import Request
 from requests.adapters import HTTPAdapter
-from requests.exceptions import HTTPError
+from requests.exceptions import ConnectionError, HTTPError
 from requests_mock import Mocker
 
 import pyiceberg
@@ -41,6 +41,8 @@
     REST_CLIENT_MAX_RETRIES,
     REST_CLIENT_RETRY_BACKOFF_FACTOR,
     REST_CLIENT_SOCKET_TIMEOUT_MS,
+    REST_TABLE_CACHE_EXPIRE_AFTER_WRITE_MS,
+    REST_TABLE_CACHE_MAX_ENTRIES,
     SIGV4_MAX_RETRIES,
     SIGV4_MAX_RETRIES_DEFAULT,
     SNAPSHOT_LOADING_MODE,
@@ -49,6 +51,7 @@
     HttpMethod,
     RestCatalog,
     ScanPlanningMode,
+    _is_short_lived_credential,
     _RetryTimeoutHTTPAdapter,
 )
 from pyiceberg.exceptions import (
@@ -60,6 +63,7 @@
     NoSuchTableError,
     NoSuchViewError,
     OAuthError,
+    RESTError,
     ServerError,
     ServiceUnavailableError,
     TableAlreadyExistsError,
@@ -1449,6 +1453,152 @@ def test_load_table_200_loading_mode(
     assert actual == expected
 
 
+def test_load_table_reuses_cached_table_on_not_modified(
+    rest_mock: Mocker, example_table_metadata_with_snapshot_v1_rest_json: dict[str, Any]
+) -> None:
+    rest_mock.get(
+        f"{TEST_URI}v1/namespaces/fokko/tables/table",
+        [
+            {"json": example_table_metadata_with_snapshot_v1_rest_json, "status_code": 200, "headers": {"ETag": '"v1"'}},
+            {"status_code": 304},
+        ],
+    )
+    catalog = RestCatalog("rest", uri=TEST_URI, token=TEST_TOKEN)
+
+    loaded = catalog.load_table(("fokko", "table"))
+    reloaded = catalog.load_table(("fokko", "table"))
+
+    assert rest_mock.last_request.headers["If-None-Match"] == '"v1"'
+    assert reloaded == loaded
+    assert reloaded is not loaded
+
+
+def test_load_table_not_modified_without_cached_table(rest_mock: Mocker) -> None:
+    rest_mock.get(f"{TEST_URI}v1/namespaces/fokko/tables/table", status_code=304)
+    catalog = RestCatalog("rest", uri=TEST_URI, token=TEST_TOKEN)
+
+    with pytest.raises(RESTError, match="304 Not Modified"):
+        catalog.load_table(("fokko", "table"))
+
+
+def test_load_table_cache_disabled_with_zero_max_entries(
+    rest_mock: Mocker, example_table_metadata_with_snapshot_v1_rest_json: dict[str, Any]
+) -> None:
+    rest_mock.get(
+        f"{TEST_URI}v1/namespaces/fokko/tables/table",
+        json=example_table_metadata_with_snapshot_v1_rest_json,
+        status_code=200,
+        headers={"ETag": '"v1"'},
+    )
+    catalog = RestCatalog("rest", uri=TEST_URI, token=TEST_TOKEN, **{REST_TABLE_CACHE_MAX_ENTRIES: "0"})
+
+    catalog.load_table(("fokko", "table"))
+    catalog.load_table(("fokko", "table"))
+
+    assert "If-None-Match" not in rest_mock.last_request.headers
+
+
+@pytest.mark.parametrize("ttl_ms", ["0", "-1"])
+def test_table_cache_rejects_non_positive_expiry(rest_mock: Mocker, ttl_ms: str) -> None:
+    with pytest.raises(ValueError, match=REST_TABLE_CACHE_EXPIRE_AFTER_WRITE_MS):
+        RestCatalog("rest", uri=TEST_URI, token=TEST_TOKEN, **{REST_TABLE_CACHE_EXPIRE_AFTER_WRITE_MS: ttl_ms})
+
+
+@pytest.mark.parametrize("change", ["drop", "rename", "register", "commit"])
+def test_table_changes_invalidate_cached_table(
+    rest_mock: Mocker, example_table_metadata_with_snapshot_v1_rest_json: dict[str, Any], change: str
+) -> None:
+    table_url = f"{TEST_URI}v1/namespaces/fokko/tables/table"
+    body = example_table_metadata_with_snapshot_v1_rest_json
+    rest_mock.get(table_url, json=body, headers={"ETag": '"v1"'})
+    rest_mock.get(f"{TEST_URI}v1/namespaces/fokko/tables/renamed", json=body)
+    rest_mock.head(f"{TEST_URI}v1/namespaces/fokko", status_code=204)
+    rest_mock.delete(table_url, status_code=204)
+    rest_mock.post(f"{TEST_URI}v1/tables/rename", status_code=204)
+    rest_mock.post(f"{TEST_URI}v1/namespaces/fokko/register", json=body)
+    rest_mock.post(table_url, json=body)
+    catalog = RestCatalog("rest", uri=TEST_URI, token=TEST_TOKEN)
+    identifier = ("fokko", "table")
+    table = catalog.load_table(identifier)
+
+    {
+        "drop": lambda: catalog.drop_table(identifier),
+        "rename": lambda: catalog.rename_table(identifier, ("fokko", "renamed")),
+        "register": lambda: catalog.register_table(identifier, body["metadata-location"]),
+        "commit": lambda: catalog.commit_table(table, (), ()),
+    }[change]()
+    catalog.load_table(identifier)
+
+    assert "If-None-Match" not in rest_mock.last_request.headers
+
+
+def test_failed_drop_still_invalidates_cached_table(
+    rest_mock: Mocker, example_table_metadata_with_snapshot_v1_rest_json: dict[str, Any]
+) -> None:
+    table_url = f"{TEST_URI}v1/namespaces/fokko/tables/table"
+    rest_mock.get(table_url, json=example_table_metadata_with_snapshot_v1_rest_json, headers={"ETag": '"v1"'})
+    rest_mock.delete(table_url, exc=ConnectionError)
+    catalog = RestCatalog("rest", uri=TEST_URI, token=TEST_TOKEN)
+
+    catalog.load_table(("fokko", "table"))
+    with pytest.raises(ConnectionError):
+        catalog.drop_table(("fokko", "table"))
+    catalog.load_table(("fokko", "table"))
+
+    assert "If-None-Match" not in rest_mock.last_request.headers
+
+
+def test_load_table_404_invalidates_cached_table(
+    rest_mock: Mocker, example_table_metadata_with_snapshot_v1_rest_json: dict[str, Any]
+) -> None:
+    found = {"json": example_table_metadata_with_snapshot_v1_rest_json, "headers": {"ETag": '"v1"'}}
+    rest_mock.get(f"{TEST_URI}v1/namespaces/fokko/tables/table", [found, {"status_code": 404, "json": {}}, found])
+    catalog = RestCatalog("rest", uri=TEST_URI, token=TEST_TOKEN)
+
+    catalog.load_table(("fokko", "table"))
+    with pytest.raises(NoSuchTableError):
+        catalog.load_table(("fokko", "table"))
+    catalog.load_table(("fokko", "table"))
+
+    assert "If-None-Match" not in rest_mock.last_request.headers
+
+
+@pytest.mark.parametrize(
+    "vended",
+    [
+        {"config": {"s3.session-token": "token"}},
+        {"storage-credentials": [{"prefix": "s3://", "config": {"s3.session-token": "token"}}]},
+    ],
+)
+def test_load_table_vending_credentials_evicts_cached_table(
+    rest_mock: Mocker, example_table_metadata_with_snapshot_v1_rest_json: dict[str, Any], vended: dict[str, Any]
+) -> None:
+    plain = {"json": example_table_metadata_with_snapshot_v1_rest_json, "headers": {"ETag": '"v1"'}}
+    vending = {"json": {**example_table_metadata_with_snapshot_v1_rest_json, **vended}, "headers": {"ETag": '"v2"'}}
+    rest_mock.get(f"{TEST_URI}v1/namespaces/fokko/tables/table", [plain, vending, plain])
+    catalog = RestCatalog("rest", uri=TEST_URI, token=TEST_TOKEN)
+
+    for _ in range(3):
+        catalog.load_table(("fokko", "table"))
+
+    assert "If-None-Match" not in rest_mock.last_request.headers
+
+
+@pytest.mark.parametrize(
+    "key, expected",
+    [
+        ("S3.Session-Token", True),
+        ("gcs.oauth2.token", True),
+        ("adls.sas-token.account", True),
+        ("s3.sse.key", False),
+        ("s3.access-key-id", False),
+        ("encryption.key-id", False),
+    ],
+)
+def test_is_short_lived_credential(key: str, expected: bool) -> None:
+    assert _is_short_lived_credential(key) is expected
+
+
 def test_load_table_honor_access_delegation(
     rest_mock: Mocker, example_table_metadata_with_snapshot_v1_rest_json: dict[str, Any]
 ) -> None:

tests/integration/test_rest_catalog.py

@@ -17,6 +17,7 @@
 # pylint:disable=redefined-outer-name
 
 import time
+from unittest import mock
 
 import pytest
 from pytest_lazy_fixtures import lf
@@ -69,6 +70,28 @@ def test_create_namespace_if_already_existing(catalog: RestCatalog) -> None:
     assert catalog.namespace_exists(TEST_NAMESPACE_IDENTIFIER)
 
 
+@pytest.mark.integration
+@pytest.mark.parametrize("catalog", [lf("session_catalog")])
+def test_load_table_reuses_cached_table_on_not_modified(
+    catalog: RestCatalog, table_schema_nested: Schema, database_name: str, table_name: str
+) -> None:
+    identifier = (database_name, table_name)
+    catalog.create_namespace_if_not_exists(database_name)
+    catalog.create_table(identifier, table_schema_nested)
+    statuses: list[int] = []
+
+    # The fixture predates apache/iceberg#17598, so its unquoted ETags get a `--gzip` suffix that Jetty never strips.
+    with (
+        mock.patch.dict(catalog._session.headers, {"Accept-Encoding": "identity"}),
+        mock.patch.object(catalog._session, "hooks", {"response": [lambda r, **_: statuses.append(r.status_code)]}),
+    ):
+        loaded = catalog.load_table(identifier)
+        reloaded = catalog.load_table(identifier)
+
+    assert statuses == [200, 304]
+    assert reloaded.metadata_location == loaded.metadata_location
+
+
 @pytest.mark.integration
 @pytest.mark.parametrize("catalog", [lf("session_catalog")])
 def test_load_view(catalog: RestCatalog, table_schema_nested: Schema, database_name: str, view_name: str) -> None: