Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Add base headers in properties to signer_headers #1610

Merged
merged 4 commits into from
Feb 10, 2025
Merged
Show file tree
Hide file tree
Changes from 1 commit
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
8 changes: 2 additions & 6 deletions pyiceberg/catalog/rest.py
Original file line number Diff line number Diff line change
Expand Up @@ -79,7 +79,7 @@
from pyiceberg.typedef import EMPTY_DICT, UTF8, IcebergBaseModel, Identifier, Properties
from pyiceberg.types import transform_dict_value_to_str
from pyiceberg.utils.deprecated import deprecation_message
from pyiceberg.utils.properties import get_first_property_value, property_as_bool
from pyiceberg.utils.properties import get_first_property_value, get_header_properties, property_as_bool

if TYPE_CHECKING:
import pyarrow as pa
Expand Down Expand Up @@ -139,7 +139,6 @@ class IdentifierKind(Enum):
SIGV4_SERVICE = "rest.signing-name"
AUTH_URL = "rest.authorization-url"
OAUTH2_SERVER_URI = "oauth2-server-uri"
HEADER_PREFIX = "header."

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

Expand Down Expand Up @@ -554,16 +553,13 @@ def _refresh_token(self, session: Optional[Session] = None, initial_token: Optio
session.headers[AUTHORIZATION_HEADER] = f"{BEARER_PREFIX} {token}"

def _config_headers(self, session: Session) -> None:
header_properties = self._extract_headers_from_properties()
header_properties = get_header_properties(self.properties)
session.headers.update(header_properties)
session.headers["Content-type"] = "application/json"
session.headers["X-Client-Version"] = ICEBERG_REST_SPEC_VERSION
session.headers["User-Agent"] = f"PyIceberg/{__version__}"
session.headers.setdefault("X-Iceberg-Access-Delegation", ACCESS_DELEGATION_DEFAULT)

def _extract_headers_from_properties(self) -> Dict[str, str]:
return {key[len(HEADER_PREFIX) :]: value for key, value in self.properties.items() if key.startswith(HEADER_PREFIX)}

def _create_table(
self,
identifier: Union[str, Identifier],
Expand Down
3 changes: 2 additions & 1 deletion pyiceberg/io/fsspec.py
Original file line number Diff line number Diff line change
Expand Up @@ -79,7 +79,7 @@
OutputStream,
)
from pyiceberg.typedef import Properties
from pyiceberg.utils.properties import get_first_property_value, property_as_bool
from pyiceberg.utils.properties import get_first_property_value, get_header_properties, property_as_bool

logger = logging.getLogger(__name__)

Expand All @@ -91,6 +91,7 @@ def s3v4_rest_signer(properties: Properties, request: AWSRequest, **_: Any) -> A
signer_headers = {}
if token := properties.get(TOKEN):
signer_headers = {"Authorization": f"Bearer {token}"}
signer_headers.update(get_header_properties(properties))

signer_body = {
"method": request.method,
Expand Down
13 changes: 13 additions & 0 deletions pyiceberg/utils/properties.py
Original file line number Diff line number Diff line change
Expand Up @@ -25,6 +25,9 @@
from pyiceberg.types import strtobool


HEADER_PREFIX = "header."


def property_as_int(
properties: Dict[str, str],
property_name: str,
Expand Down Expand Up @@ -74,3 +77,13 @@ def get_first_property_value(
if property_value := properties.get(property_name):
return property_value
return None


def get_header_properties(
properties: Properties,
) -> Properties:
return {
key[len(HEADER_PREFIX):]: value
tom-s-powell marked this conversation as resolved.
Show resolved Hide resolved
for key, value in properties.items()
if key.startswith(HEADER_PREFIX)
}
3 changes: 2 additions & 1 deletion tests/io/test_fsspec.py
Original file line number Diff line number Diff line change
Expand Up @@ -692,6 +692,7 @@ def test_s3v4_rest_signer(requests_mock: Mocker) -> None:
"X-Amz-Security-Token": [
"YQoJb3JpZ2luX2VjEDoaCXVzLXdlc3QtMiJGMEQCID/fFxZP5oaEgQmcwP6XhZa0xSq9lmLSx8ffaWbySfUPAiAesa7sjd/WV4uwRTO0S03y/MWVtgpH+/NyZQ4bZgLVriqrAggTEAEaDDAzMzQwNzIyMjE1OSIMOeFOWhZIurMmAqjsKogCxMCqxX8ZjK0gacAkcDqBCyA7qTSLhdfKQIH/w7WpLBU1km+cRUWWCudan6gZsAq867DBaKEP7qI05DAWr9MChAkgUgyI8/G3Z23ET0gAedf3GsJbakB0F1kklx8jPmj4BPCht9RcTiXiJ5DxTS/cRCcalIQXmPFbaJSqpBusVG2EkWnm1v7VQrNPE2Os2b2P293vpbhwkyCEQiGRVva4Sw9D1sKvqSsK10QCRG+os6dFEOu1kARaXi6pStvR4OVmj7OYeAYjzaFchn7nz2CSae0M4IluiYQ01eQAywbfRo9DpKSmDM/DnPZWJnD/woLhaaaCrCxSSEaFsvGOHFhLd3Rknw1v0jADMILUtJoGOp4BpqKqyMz0CY3kpKL0jfR3ykTf/ge9wWVE0Alr7wRIkGCIURkhslGHqSyFRGoTqIXaxU+oPbwlw/0w/nYO7qQ6bTANOWye/wgw4h/NmJ6vU7wnZTXwREf1r6MF72++bE/fMk19LfVb8jN/qrUqAUXTc8gBAUxL5pgy8+oT/JnI2BkVrrLS4ilxEXP9Ahm+6GDUYXV4fBpqpZwdkzQ/5Gw="
],
"X-Custom-Header": ["value"],
},
"extensions": {},
},
Expand All @@ -714,7 +715,7 @@ def test_s3v4_rest_signer(requests_mock: Mocker) -> None:
"retries": {"attempt": 1, "invocation-id": "75d143fb-0219-439b-872c-18213d1c8d54"},
}

signed_request = s3v4_rest_signer({"token": "abc", "uri": TEST_URI}, request)
signed_request = s3v4_rest_signer({"token": "abc", "uri": TEST_URI, "header.X-Custom-Header": "value"}, request)

assert signed_request.url == new_uri
assert dict(signed_request.headers) == {
Expand Down
Loading