diff --git a/pyiceberg/catalog/rest/__init__.py b/pyiceberg/catalog/rest/__init__.py index d8f58773e6..ce7c35842a 100644 --- a/pyiceberg/catalog/rest/__init__.py +++ b/pyiceberg/catalog/rest/__init__.py @@ -1603,14 +1603,25 @@ def commit_table( table_identifier = TableIdentifier(namespace=identifier[:-1], name=identifier[-1]) table_request = CommitTableRequest(identifier=table_identifier, requirements=requirements, updates=updates) - headers = self._session.headers + post_kwargs: dict[str, Any] = { + "data": table_request.model_dump_json().encode(UTF8), + } if table_token := table.config.get(TOKEN): - headers[AUTHORIZATION_HEADER] = f"{BEARER_PREFIX} {table_token}" + table_auth = f"{BEARER_PREFIX} {table_token}" + post_kwargs["headers"] = {AUTHORIZATION_HEADER: table_auth} + + # Session.auth would otherwise overwrite Authorization with the + # catalog token. Apply the table-scoped token on this request only + # and leave the session headers untouched. + def _apply_table_token(request: PreparedRequest) -> PreparedRequest: + request.headers[AUTHORIZATION_HEADER] = table_auth + return request + + post_kwargs["auth"] = _apply_table_token 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, + **post_kwargs, ) try: response.raise_for_status() diff --git a/pyiceberg/io/fsspec.py b/pyiceberg/io/fsspec.py index 09bbe6f1d6..8e8b346945 100644 --- a/pyiceberg/io/fsspec.py +++ b/pyiceberg/io/fsspec.py @@ -231,6 +231,13 @@ def _s3(properties: Properties) -> AbstractFileSystem: if s3_additional_kwargs: s3_fs_kwargs["s3_additional_kwargs"] = s3_additional_kwargs + # Signers are registered after construction, so they are not part of the + # constructor kwargs fsspec uses as a cache key. Disable instance caching + # whenever a custom signer is in play so two catalogs do not share one FS + # and overwrite each other's before-sign handler. + if register_events: + s3_fs_kwargs["skip_instance_cache"] = True + fs = S3FileSystem(**s3_fs_kwargs) for event_name, event_function in register_events.items(): diff --git a/tests/catalog/test_rest.py b/tests/catalog/test_rest.py index a918829c24..8401d7fc77 100644 --- a/tests/catalog/test_rest.py +++ b/tests/catalog/test_rest.py @@ -31,7 +31,7 @@ from requests_mock import Mocker import pyiceberg -from pyiceberg.catalog import PropertiesUpdateSummary, load_catalog +from pyiceberg.catalog import TOKEN, PropertiesUpdateSummary, load_catalog from pyiceberg.catalog.rest import ( DEFAULT_ENDPOINTS, EMPTY_BODY_SHA256, @@ -2747,6 +2747,34 @@ def test_table_identifier_in_commit_table_request( ) +def test_commit_table_does_not_mutate_session_headers(rest_mock: Mocker, example_table_metadata_v2: dict[str, Any]) -> None: + metadata_location = "s3://some_bucket/metadata.json" + table_token = "table-scoped-token" + rest_mock.post( + url=f"{TEST_URI}v1/namespaces/namespace/tables/table_name", + json={ + "metadata": example_table_metadata_v2, + "metadata-location": metadata_location, + }, + status_code=200, + ) + catalog = RestCatalog("catalog_name", uri=TEST_URI, token=TEST_TOKEN) + session_auth = catalog._session.headers.get("Authorization") + table = Table( + identifier=("namespace", "table_name"), + metadata=None, # type: ignore + metadata_location=metadata_location, + io=None, # type: ignore + catalog=catalog, + config={TOKEN: table_token}, + ) + + catalog.commit_table(table, (), ()) + + assert catalog._session.headers.get("Authorization") == session_auth + assert rest_mock.last_request.headers["Authorization"] == f"Bearer {table_token}" + + def test_register_view_200(rest_mock: Mocker, example_view_metadata_rest_json: dict[str, Any]) -> None: rest_mock.head( f"{TEST_URI}v1/namespaces/default/tables/registered_view", diff --git a/tests/io/test_fsspec.py b/tests/io/test_fsspec.py index 45835a08eb..b0fb20d9f4 100644 --- a/tests/io/test_fsspec.py +++ b/tests/io/test_fsspec.py @@ -17,6 +17,7 @@ import os import pickle +import sys import tempfile import threading import uuid @@ -30,8 +31,8 @@ from pyiceberg.catalog.rest.auth import AUTH_MANAGER from pyiceberg.exceptions import SignError -from pyiceberg.io import fsspec -from pyiceberg.io.fsspec import FsspecFileIO, S3V4RestSigner +from pyiceberg.io import S3_SIGNER, fsspec +from pyiceberg.io.fsspec import FsspecFileIO, S3V4RestSigner, _s3 from pyiceberg.io.pyarrow import PyArrowFileIO from pyiceberg.typedef import Properties from tests.conftest import UNIFIED_AWS_SESSION_PROPERTIES @@ -924,6 +925,28 @@ def _test_fsspec_pickle_round_trip(fsspec_fileio: FsspecFileIO, location: str) - TEST_URI = "https://iceberg-test-signer" +def test_s3_custom_signer_skips_fsspec_instance_cache() -> None: + captured: list[dict[str, object]] = [] + + class FakeS3: + def __init__(self, **kwargs: object) -> None: + captured.append(kwargs) + self.s3 = mock.Mock() + self.s3.meta.events.unregister = mock.Mock() + self.s3.meta.events.register_last = mock.Mock() + + fake_s3fs = mock.MagicMock() + fake_s3fs.S3FileSystem = FakeS3 + with mock.patch.dict(sys.modules, {"s3fs": fake_s3fs}): + first = _s3({S3_SIGNER: "S3V4RestSigner", "token": "one", "uri": TEST_URI}) + second = _s3({S3_SIGNER: "S3V4RestSigner", "token": "two", "uri": TEST_URI}) + + assert first is not second + assert len(captured) == 2 + assert captured[0]["skip_instance_cache"] is True + assert captured[1]["skip_instance_cache"] is True + + def test_s3v4_rest_signer(requests_mock: Mocker) -> None: new_uri = "https://other-bucket/metadata/snap-8048355899640248710-1-a5c8ea2d-aa1f-48e8-89f4-1fa69db8c742.avro" requests_mock.post(