Skip to content
Open
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
19 changes: 15 additions & 4 deletions pyiceberg/catalog/rest/__init__.py
Original file line number Diff line number Diff line change
Expand Up @@ -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()
Expand Down
7 changes: 7 additions & 0 deletions pyiceberg/io/fsspec.py
Original file line number Diff line number Diff line change
Expand Up @@ -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():
Expand Down
30 changes: 29 additions & 1 deletion tests/catalog/test_rest.py
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand Down Expand Up @@ -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",
Expand Down
27 changes: 25 additions & 2 deletions tests/io/test_fsspec.py
Original file line number Diff line number Diff line change
Expand Up @@ -17,6 +17,7 @@

import os
import pickle
import sys
import tempfile
import threading
import uuid
Expand All @@ -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
Expand Down Expand Up @@ -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(
Expand Down
Loading