Skip to content

Commit df984c9

Browse files
add test case
1 parent d595859 commit df984c9

File tree

2 files changed

+64
-1
lines changed

2 files changed

+64
-1
lines changed

pyiceberg/catalog/rest/__init__.py

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -216,6 +216,7 @@ class IdentifierKind(Enum):
216216
SIGV4 = "rest.sigv4-enabled"
217217
SIGV4_REGION = "rest.signing-region"
218218
SIGV4_SERVICE = "rest.signing-name"
219+
EMPTY_BODY_SHA256: str = "e3b0c44298fc1c149afbf4c8996fb92427ae41e4649b934ca495991b7852b855"
219220
OAUTH2_SERVER_URI = "oauth2-server-uri"
220221
SNAPSHOT_LOADING_MODE = "snapshot-loading-mode"
221222
AUTH = "auth"
@@ -565,7 +566,7 @@ def add_headers(self, request: PreparedRequest, **kwargs: Any) -> None: # pylin
565566
# For empty bodies, explicitly set the content hash header to the SHA256 of an empty string
566567
body = request.body
567568
if body in (None, b"", ""):
568-
request.headers["x-amz-content-sha256"] = "e3b0c44298fc1c149afbf4c8996fb92427ae41e4649b934ca495991b7852b855"
569+
request.headers["x-amz-content-sha256"] = EMPTY_BODY_SHA256
569570

570571
aws_request = AWSRequest(
571572
method=request.method, url=url, params=params, data=request.body, headers=dict(request.headers)

tests/catalog/test_rest.py

Lines changed: 62 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -23,6 +23,7 @@
2323

2424
import pytest
2525
from requests.exceptions import HTTPError
26+
from requests.adapters import HTTPAdapter
2627
from requests_mock import Mocker
2728

2829
import pyiceberg
@@ -49,6 +50,7 @@
4950
from pyiceberg.transforms import IdentityTransform, TruncateTransform
5051
from pyiceberg.typedef import RecursiveDict
5152
from pyiceberg.utils.config import Config
53+
from requests import Request
5254

5355
TEST_URI = "https://iceberg-test-catalog/"
5456
TEST_CREDENTIALS = "client:secret_with:colon"
@@ -451,6 +453,66 @@ def test_list_tables_200_sigv4(rest_mock: Mocker) -> None:
451453
assert rest_mock.called
452454

453455

456+
def test_sigv4_sign_request_without_body(rest_mock: Mocker) -> None:
457+
existing_token = "existing_token"
458+
from pyiceberg.catalog.rest import EMPTY_BODY_SHA256
459+
from requests.adapters import HTTPAdapter
460+
461+
catalog = RestCatalog(
462+
"rest",
463+
**{
464+
"uri": TEST_URI,
465+
"token": existing_token,
466+
"rest.sigv4-enabled": "true",
467+
"rest.signing-region": "us-west-2",
468+
"client.access-key-id": "id",
469+
"client.secret-access-key": "secret",
470+
},
471+
)
472+
473+
prepared = catalog._session.prepare_request(Request("GET", f"{TEST_URI}v1/config"))
474+
adapter = catalog._session.adapters[catalog.uri]
475+
assert isinstance(adapter, HTTPAdapter)
476+
adapter.add_headers(prepared)
477+
478+
assert prepared.headers["Authorization"].startswith("AWS4-HMAC-SHA256")
479+
assert prepared.headers["Original-Authorization"] == f"Bearer {existing_token}"
480+
assert prepared.headers["x-amz-content-sha256"] == EMPTY_BODY_SHA256
481+
482+
483+
def test_sigv4_sign_request_with_body(rest_mock: Mocker) -> None:
484+
existing_token = "existing_token"
485+
from pyiceberg.catalog.rest import EMPTY_BODY_SHA256
486+
from requests.adapters import HTTPAdapter
487+
488+
catalog = RestCatalog(
489+
"rest",
490+
**{
491+
"uri": TEST_URI,
492+
"token": existing_token,
493+
"rest.sigv4-enabled": "true",
494+
"rest.signing-region": "us-west-2",
495+
"client.access-key-id": "id",
496+
"client.secret-access-key": "secret",
497+
},
498+
)
499+
500+
prepared = catalog._session.prepare_request(
501+
Request(
502+
"POST",
503+
f"{TEST_URI}v1/namespaces",
504+
data={"namespace": "asdfasd"},
505+
)
506+
)
507+
adapter = catalog._session.adapters[catalog.uri]
508+
assert isinstance(adapter, HTTPAdapter)
509+
adapter.add_headers(prepared)
510+
511+
assert prepared.headers["Authorization"].startswith("AWS4-HMAC-SHA256")
512+
assert prepared.headers["Original-Authorization"] == f"Bearer {existing_token}"
513+
assert prepared.headers.get("x-amz-content-sha256") != EMPTY_BODY_SHA256
514+
515+
454516
def test_list_tables_404(rest_mock: Mocker) -> None:
455517
namespace = "examples"
456518
rest_mock.get(

0 commit comments

Comments
 (0)