|
23 | 23 |
|
24 | 24 | import pytest |
25 | 25 | from requests.exceptions import HTTPError |
| 26 | +from requests.adapters import HTTPAdapter |
26 | 27 | from requests_mock import Mocker |
27 | 28 |
|
28 | 29 | import pyiceberg |
|
49 | 50 | from pyiceberg.transforms import IdentityTransform, TruncateTransform |
50 | 51 | from pyiceberg.typedef import RecursiveDict |
51 | 52 | from pyiceberg.utils.config import Config |
| 53 | +from requests import Request |
52 | 54 |
|
53 | 55 | TEST_URI = "https://iceberg-test-catalog/" |
54 | 56 | TEST_CREDENTIALS = "client:secret_with:colon" |
@@ -451,6 +453,66 @@ 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 | + 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 | + |
454 | 516 | def test_list_tables_404(rest_mock: Mocker) -> None: |
455 | 517 | namespace = "examples" |
456 | 518 | rest_mock.get( |
|
0 commit comments