|
22 | 22 | from unittest import mock |
23 | 23 |
|
24 | 24 | import pytest |
| 25 | +from requests import Request |
| 26 | +from requests.adapters import HTTPAdapter |
25 | 27 | from requests.exceptions import HTTPError |
26 | 28 | from requests_mock import Mocker |
27 | 29 |
|
28 | 30 | import pyiceberg |
29 | 31 | from pyiceberg.catalog import PropertiesUpdateSummary, load_catalog |
30 | | -from pyiceberg.catalog.rest import DEFAULT_ENDPOINTS, OAUTH2_SERVER_URI, SNAPSHOT_LOADING_MODE, Capability, RestCatalog |
| 32 | +from pyiceberg.catalog.rest import DEFAULT_ENDPOINTS, EMPTY_BODY_SHA256, OAUTH2_SERVER_URI, SNAPSHOT_LOADING_MODE, Capability, RestCatalog |
31 | 33 | from pyiceberg.exceptions import ( |
32 | 34 | AuthorizationExpiredError, |
33 | 35 | NamespaceAlreadyExistsError, |
@@ -451,6 +453,62 @@ def test_list_tables_200_sigv4(rest_mock: Mocker) -> None: |
451 | 453 | assert rest_mock.called |
452 | 454 |
|
453 | 455 |
|
| 456 | +def test_sigv4_sign_request_without_body(rest_mock: Mocker) -> None: |
| 457 | + existing_token = "existing_token" |
| 458 | + |
| 459 | + catalog = RestCatalog( |
| 460 | + "rest", |
| 461 | + **{ |
| 462 | + "uri": TEST_URI, |
| 463 | + "token": existing_token, |
| 464 | + "rest.sigv4-enabled": "true", |
| 465 | + "rest.signing-region": "us-west-2", |
| 466 | + "client.access-key-id": "id", |
| 467 | + "client.secret-access-key": "secret", |
| 468 | + }, |
| 469 | + ) |
| 470 | + |
| 471 | + prepared = catalog._session.prepare_request(Request("GET", f"{TEST_URI}v1/config")) |
| 472 | + adapter = catalog._session.adapters[catalog.uri] |
| 473 | + assert isinstance(adapter, HTTPAdapter) |
| 474 | + adapter.add_headers(prepared) |
| 475 | + |
| 476 | + assert prepared.headers["Authorization"].startswith("AWS4-HMAC-SHA256") |
| 477 | + assert prepared.headers["Original-Authorization"] == f"Bearer {existing_token}" |
| 478 | + assert prepared.headers["x-amz-content-sha256"] == EMPTY_BODY_SHA256 |
| 479 | + |
| 480 | + |
| 481 | +def test_sigv4_sign_request_with_body(rest_mock: Mocker) -> None: |
| 482 | + existing_token = "existing_token" |
| 483 | + |
| 484 | + catalog = RestCatalog( |
| 485 | + "rest", |
| 486 | + **{ |
| 487 | + "uri": TEST_URI, |
| 488 | + "token": existing_token, |
| 489 | + "rest.sigv4-enabled": "true", |
| 490 | + "rest.signing-region": "us-west-2", |
| 491 | + "client.access-key-id": "id", |
| 492 | + "client.secret-access-key": "secret", |
| 493 | + }, |
| 494 | + ) |
| 495 | + |
| 496 | + prepared = catalog._session.prepare_request( |
| 497 | + Request( |
| 498 | + "POST", |
| 499 | + f"{TEST_URI}v1/namespaces", |
| 500 | + data={"namespace": "asdfasd"}, |
| 501 | + ) |
| 502 | + ) |
| 503 | + adapter = catalog._session.adapters[catalog.uri] |
| 504 | + assert isinstance(adapter, HTTPAdapter) |
| 505 | + adapter.add_headers(prepared) |
| 506 | + |
| 507 | + assert prepared.headers["Authorization"].startswith("AWS4-HMAC-SHA256") |
| 508 | + assert prepared.headers["Original-Authorization"] == f"Bearer {existing_token}" |
| 509 | + assert prepared.headers.get("x-amz-content-sha256") != EMPTY_BODY_SHA256 |
| 510 | + |
| 511 | + |
454 | 512 | def test_list_tables_404(rest_mock: Mocker) -> None: |
455 | 513 | namespace = "examples" |
456 | 514 | rest_mock.get( |
|
0 commit comments