feat: send Idempotency-Key on REST mutations when the server advertises it

mkdocs/docs/configuration.md

@@ -449,6 +449,9 @@ as a single request timeout (floored to whole seconds). Retries are applied to i
 only (`GET`, `HEAD`, `OPTIONS`, `PUT`, `DELETE`) and to the transient HTTP status codes `429`, `500`,
 `502`, `503`, `504`. Other failures are not retried.
 
+When the server advertises `idempotency-key-lifetime` in its config response, mutating requests carry an
+`Idempotency-Key` header and `POST` is retried as well, reusing the same key for each retry.
+
 #### Headers in REST Catalog
 
 To configure custom headers in REST Catalog, include them in the catalog properties with `header.<Header-Name>`. This

pyiceberg/catalog/rest/__init__.py

@@ -17,7 +17,10 @@
 from __future__ import annotations
 
 import logging
+import os
+import sys
 import time
+import uuid
 from collections import deque
 from collections.abc import Mapping
 from enum import Enum, StrEnum
@@ -298,6 +301,7 @@ class ScanPlanningMode(Enum):
 # (e.g. setting raise_on_status=False would swallow 4xx errors silently).
 _CONNECTION_RETRY_STATUS_FORCELIST = (429, 500, 502, 503, 504)
 _CONNECTION_RETRY_ALLOWED_METHODS = frozenset({"GET", "HEAD", "OPTIONS", "PUT", "DELETE"})
+IDEMPOTENCY_KEY_HEADER = "Idempotency-Key"
 EMPTY_BODY_SHA256: str = "e3b0c44298fc1c149afbf4c8996fb92427ae41e4649b934ca495991b7852b855"
 OAUTH2_SERVER_URI = "oauth2-server-uri"
 SNAPSHOT_LOADING_MODE = "snapshot-loading-mode"
@@ -339,6 +343,17 @@ def _parse_scan_planning_mode(properties: Properties, *, strict: bool = True) ->
         return None
 
 
+if sys.version_info >= (3, 14):
+    from uuid import uuid7 as _uuid7
+else:
+
+    def _uuid7() -> uuid.UUID:
+        value = (time.time_ns() // 1_000_000) << 80 | int.from_bytes(os.urandom(10))
+        value = (value & ~(0xF << 76)) | (0x7 << 76)
+        value = (value & ~(0x3 << 62)) | (0x2 << 62)
+        return uuid.UUID(int=value)
+
+
 def _retry_hook(retry_state: RetryCallState) -> None:
     rest_catalog: RestCatalog = retry_state.args[0]
     rest_catalog._refresh_token()  # pylint: disable=protected-access
@@ -421,6 +436,7 @@ class ConfigResponse(IcebergBaseModel):
     defaults: Properties | None = Field(default_factory=dict)
     overrides: Properties | None = Field(default_factory=dict)
     endpoints: set[Endpoint] | None = Field(default=None)
+    idempotency_key_lifetime: str | None = Field(alias="idempotency-key-lifetime", default=None)
 
     @field_validator("endpoints", mode="before")
     @classmethod
@@ -502,11 +518,12 @@ def send(
         return super().send(request, stream=stream, timeout=timeout, verify=verify, cert=cert, proxies=proxies)
 
 
-def _create_connection_adapter(properties: Properties) -> _RetryTimeoutHTTPAdapter | None:
+def _create_connection_adapter(properties: Properties, retry_post: bool = False) -> _RetryTimeoutHTTPAdapter | None:
     """Build a connection adapter from the optional `rest.client.*` properties.
 
     Returns None when no connection properties are supplied, leaving the default
-    Session behavior unchanged. Raises ValueError on invalid input.
+    Session behavior unchanged. POST requests are retried only when `retry_post` is set.
+    Raises ValueError on invalid input.
     """
     connection_timeout_ms = property_as_int(properties, REST_CLIENT_CONNECTION_TIMEOUT_MS)
     if connection_timeout_ms is not None and connection_timeout_ms <= 0:
@@ -541,7 +558,7 @@ def _create_connection_adapter(properties: Properties) -> _RetryTimeoutHTTPAdapt
             total=retries if retries is not None else DEFAULT_RETRIES,
             backoff_factor=backoff_factor if backoff_factor is not None else 0.0,
             status_forcelist=list(_CONNECTION_RETRY_STATUS_FORCELIST),
-            allowed_methods=_CONNECTION_RETRY_ALLOWED_METHODS,
+            allowed_methods=_CONNECTION_RETRY_ALLOWED_METHODS | {"POST"} if retry_post else _CONNECTION_RETRY_ALLOWED_METHODS,
             # Return the final response on retry exhaustion (instead of raising MaxRetryError)
             # so `_handle_non_200_response` can map the 5xx status to a typed exception
             # (ServiceUnavailableError, etc.). 4xx codes are not in status_forcelist and are
@@ -557,6 +574,7 @@ class RestCatalog(Catalog):
     _auth_manager: AuthManager | None
     _supported_endpoints: set[Endpoint]
     _namespace_separator: str
+    _idempotency_enabled: bool = False
 
     def __init__(self, name: str, **properties: str):
         """Rest Catalog.
@@ -579,7 +597,7 @@ def _create_session(self) -> Session:
 
         # Mount the retry/timeout adapter when `connection.*` properties are set.
         # SigV4's adapter mounted below at `self.uri` is a longer prefix and still wins for that host.
-        if (connection_adapter := _create_connection_adapter(self.properties)) is not None:
+        if (connection_adapter := _create_connection_adapter(self.properties, self._idempotency_enabled)) is not None:
             session.mount("http://", connection_adapter)
             session.mount("https://", connection_adapter)
 
@@ -1052,6 +1070,8 @@ def _fetch_config(self) -> None:
         # Update URI based on overrides
         self.uri = config[URI]
 
+        self._idempotency_enabled = config_response.idempotency_key_lifetime is not None
+
         # Determine supported endpoints
         endpoints = config_response.endpoints
         if endpoints:
@@ -1068,6 +1088,9 @@ def _fetch_config(self) -> None:
             raise ValueError("Namespace separator cannot be an empty string")
         self._namespace_separator = unquote(separator_from_properties)
 
+    def _mutation_headers(self) -> dict[str, str]:
+        return {IDEMPOTENCY_KEY_HEADER: str(_uuid7())} if self._idempotency_enabled else {}
+
     def _identifier_to_validated_tuple(self, identifier: str | Identifier) -> Identifier:
         identifier_tuple = self.identifier_to_tuple(identifier)
         if len(identifier_tuple) <= 1:
@@ -1244,6 +1267,7 @@ def _create_table(
         response = self._session.post(
             self.url(Endpoints.create_table, namespace=namespace_and_table["namespace"]),
             data=serialized_json,
+            headers=self._mutation_headers(),
         )
         try:
             response.raise_for_status()
@@ -1326,6 +1350,7 @@ def create_view(
         response = self._session.post(
             self.url(Endpoints.create_view, namespace=namespace_and_view["namespace"]),
             data=serialized_json,
+            headers=self._mutation_headers(),
         )
 
         try:
@@ -1363,6 +1388,7 @@ def register_table(self, identifier: str | Identifier, metadata_location: str, o
         response = self._session.post(
             self.url(Endpoints.register_table, namespace=namespace_and_table["namespace"]),
             data=serialized_json,
+            headers=self._mutation_headers(),
         )
         try:
             response.raise_for_status()
@@ -1389,6 +1415,7 @@ def unregister_table(self, identifier: str | Identifier) -> tuple[str, TableMeta
         namespace_and_table = self._split_identifier_for_path(identifier)
         response = self._session.post(
             self.url(Endpoints.unregister_table, prefixed=True, **namespace_and_table),
+            headers=self._mutation_headers(),
         )
         try:
             response.raise_for_status()
@@ -1489,6 +1516,7 @@ def drop_table(self, identifier: str | Identifier, purge_requested: bool = False
         response = self._session.delete(
             self.url(Endpoints.drop_table, prefixed=True, **self._split_identifier_for_path(identifier)),
             params={"purgeRequested": purge_requested},
+            headers=self._mutation_headers(),
         )
         try:
             response.raise_for_status()
@@ -1518,7 +1546,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}")
 
-        response = self._session.post(self.url(Endpoints.rename_table), json=payload)
+        response = self._session.post(self.url(Endpoints.rename_table), json=payload, headers=self._mutation_headers())
         try:
             response.raise_for_status()
         except HTTPError as exc:
@@ -1622,7 +1650,7 @@ def commit_table(
         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),
-            headers=headers,
+            headers={**headers, **self._mutation_headers()},
         )
         try:
             response.raise_for_status()
@@ -1644,7 +1672,7 @@ def create_namespace(self, namespace: str | Identifier, properties: Properties =
         self._check_endpoint(Capability.V1_CREATE_NAMESPACE)
         namespace_tuple = self._check_valid_namespace_identifier(namespace)
         payload = {"namespace": namespace_tuple, "properties": properties}
-        response = self._session.post(self.url(Endpoints.create_namespace), json=payload)
+        response = self._session.post(self.url(Endpoints.create_namespace), json=payload, headers=self._mutation_headers())
         try:
             response.raise_for_status()
         except HTTPError as exc:
@@ -1656,7 +1684,7 @@ def drop_namespace(self, namespace: str | Identifier) -> None:
         self._check_endpoint(Capability.V1_DELETE_NAMESPACE)
         namespace_tuple = self._check_valid_namespace_identifier(namespace)
         namespace = self._encode_namespace_path(namespace_tuple)
-        response = self._session.delete(self.url(Endpoints.drop_namespace, namespace=namespace))
+        response = self._session.delete(self.url(Endpoints.drop_namespace, namespace=namespace), headers=self._mutation_headers())
         try:
             response.raise_for_status()
         except HTTPError as exc:
@@ -1723,7 +1751,9 @@ def update_namespace_properties(
         namespace_tuple = self._check_valid_namespace_identifier(namespace)
         namespace = self._encode_namespace_path(namespace_tuple)
         payload = {"removals": list(removals or []), "updates": updates}
-        response = self._session.post(self.url(Endpoints.update_namespace_properties, namespace=namespace), json=payload)
+        response = self._session.post(
+            self.url(Endpoints.update_namespace_properties, namespace=namespace), json=payload, headers=self._mutation_headers()
+        )
         try:
             response.raise_for_status()
         except HTTPError as exc:
@@ -1838,6 +1868,7 @@ def register_view(self, identifier: str | Identifier, metadata_location: str) ->
         response = self._session.post(
             self.url(Endpoints.register_view, namespace=namespace_and_view["namespace"]),
             data=serialized_json,
+            headers=self._mutation_headers(),
         )
         try:
             response.raise_for_status()
@@ -1853,6 +1884,7 @@ def drop_view(self, identifier: str | Identifier) -> None:
         self._check_endpoint(Capability.V1_DELETE_VIEW)
         response = self._session.delete(
             self.url(Endpoints.drop_view, prefixed=True, **self._split_identifier_for_path(identifier, IdentifierKind.VIEW)),
+            headers=self._mutation_headers(),
         )
         try:
             response.raise_for_status()
@@ -1876,7 +1908,7 @@ def rename_view(self, from_identifier: str | Identifier, to_identifier: str | Id
             raise NoSuchNamespaceError(f"Destination namespace does not exist: {destination_namespace}")
 
         payload = {"source": source, "destination": destination}
-        response = self._session.post(self.url(Endpoints.rename_view), json=payload)
+        response = self._session.post(self.url(Endpoints.rename_view), json=payload, headers=self._mutation_headers())
         try:
             response.raise_for_status()
         except HTTPError as exc:

tests/catalog/test_rest.py

@@ -19,6 +19,7 @@
 
 import base64
 import os
+import uuid
 from collections.abc import Callable, Iterator
 from contextlib import contextmanager
 from typing import Any, cast
@@ -37,6 +38,7 @@
 from pyiceberg.catalog.rest import (
     DEFAULT_ENDPOINTS,
     EMPTY_BODY_SHA256,
+    IDEMPOTENCY_KEY_HEADER,
     OAUTH2_SERVER_URI,
     PAGE_SIZE,
     REST_CLIENT_CONNECTION_TIMEOUT_MS,
@@ -2460,28 +2462,38 @@ def test_session_with_socket_timeout_only(rest_mock: Mocker) -> None:
 
 
 @contextmanager
-def _local_rest_server_503_then_200(num_failures: int) -> Iterator[dict[str, Any]]:
+def _local_rest_server_503_then_200(num_failures: int, idempotency_key_lifetime: str | None = None) -> Iterator[dict[str, Any]]:
     """Stand up a loopback HTTP server that returns `num_failures` 503s for `/v1/namespaces` then a 200.
 
     Used in place of `requests_mock`, which replaces the HTTPAdapter and would bypass the retry logic.
 
-    Yields a dict with `port` and `namespace_calls` keys (the latter is updated in-place as requests arrive).
+    Yields a dict with `port`, `namespace_calls` and `idempotency_keys` keys, updated in-place as requests arrive.
     """
     import json
     import threading
     from http.server import BaseHTTPRequestHandler, HTTPServer
 
-    state: dict[str, Any] = {"namespace_calls": 0}
-    config_body = json.dumps(
-        {"defaults": {}, "overrides": {}, "endpoints": [str(endpoint) for endpoint in TEST_SUPPORTED_ENDPOINTS]}
-    ).encode()
+    state: dict[str, Any] = {"namespace_calls": 0, "idempotency_keys": []}
+    config: dict[str, Any] = {
+        "defaults": {},
+        "overrides": {},
+        "endpoints": [str(endpoint) for endpoint in TEST_SUPPORTED_ENDPOINTS],
+    }
+    if idempotency_key_lifetime is not None:
+        config["idempotency-key-lifetime"] = idempotency_key_lifetime
+    config_body = json.dumps(config).encode()
 
     class _Handler(BaseHTTPRequestHandler):
+        def do_POST(self) -> None:
+            self.rfile.read(int(self.headers.get("Content-Length", 0)))
+            self.do_GET()
+
         def do_GET(self) -> None:
             if self.path.endswith("/v1/config"):
                 self._respond(200, config_body)
             elif self.path.endswith("/v1/namespaces"):
                 state["namespace_calls"] += 1
+                state["idempotency_keys"].append(self.headers.get(IDEMPOTENCY_KEY_HEADER))
                 if state["namespace_calls"] <= num_failures:
                     self._respond(503, b"")
                 else:
@@ -2548,6 +2560,48 @@ def test_session_exhausted_retries_surfaces_typed_exception() -> None:
         assert server["namespace_calls"] == 3
 
 
+def test_session_retries_keyed_post_with_the_same_key() -> None:
+    with _local_rest_server_503_then_200(num_failures=1, idempotency_key_lifetime="PT30M") as server:
+        catalog = RestCatalog(
+            "rest",
+            **{
+                "uri": f"http://127.0.0.1:{server['port']}/",
+                "token": TEST_TOKEN,
+                REST_CLIENT_MAX_RETRIES: "1",
+                REST_CLIENT_RETRY_BACKOFF_FACTOR: "0",
+            },
+        )
+        catalog.create_namespace("foo")
+
+        key = server["idempotency_keys"][0]
+        assert key is not None
+        assert server["idempotency_keys"] == [key, key]
+
+
+def test_mutations_send_idempotency_key_when_advertised(requests_mock: Mocker) -> None:
+    requests_mock.get(f"{TEST_URI}v1/config", json={"defaults": {}, "overrides": {}, "idempotency-key-lifetime": "PT30M"})
+    requests_mock.post(f"{TEST_URI}v1/namespaces", json={"namespace": ["foo"], "properties": {}})
+    requests_mock.get(f"{TEST_URI}v1/namespaces", json={"namespaces": []})
+    catalog = RestCatalog("rest", uri=TEST_URI, token=TEST_TOKEN)
+
+    catalog.create_namespace("foo")
+    catalog.create_namespace("foo")
+    catalog.list_namespaces()
+
+    first, second, listed = (request.headers.get(IDEMPOTENCY_KEY_HEADER) for request in requests_mock.request_history[1:])
+    assert uuid.UUID(first).version == 7
+    assert first != second
+    assert listed is None
+
+
+def test_mutations_omit_idempotency_key_when_not_advertised(rest_mock: Mocker) -> None:
+    rest_mock.post(f"{TEST_URI}v1/namespaces", json={"namespace": ["foo"], "properties": {}})
+
+    RestCatalog("rest", uri=TEST_URI, token=TEST_TOKEN).create_namespace("foo")
+
+    assert IDEMPOTENCY_KEY_HEADER not in rest_mock.last_request.headers
+
+
 def test_session_with_invalid_connection_timeout_raises(rest_mock: Mocker) -> None:
     catalog_properties = {
         "uri": TEST_URI,