Skip to content

Commit

Permalink
Add Support for Custom Header Configurations in RESTCatalog (#467)
Browse files Browse the repository at this point in the history
  • Loading branch information
geruh authored Feb 27, 2024
1 parent c5735dc commit a1099d8
Show file tree
Hide file tree
Showing 3 changed files with 65 additions and 4 deletions.
13 changes: 13 additions & 0 deletions mkdocs/docs/configuration.md
Original file line number Diff line number Diff line change
Expand Up @@ -156,6 +156,19 @@ catalog:
| rest.signing-name | execute-api | The service signing name to use when SigV4 signing a request |
| rest.authorization-url | https://auth-service/cc | Authentication URL to use for client credentials authentication (default: uri + 'v1/oauth/tokens') |

### Headers in RESTCatalog

To configure custom headers in RESTCatalog, include them in the catalog properties with the prefix `header.`. This
ensures that all HTTP requests to the REST service include the specified headers.

```yaml
catalog:
default:
uri: http://rest-catalog/ws/
credential: t-1234:secret
header.content-type: application/vnd.api+json
```

## SQL Catalog

The SQL catalog requires a database for its backend. PyIceberg supports PostgreSQL and SQLite through psycopg2. The database connection has to be configured using the `uri` property. See SQLAlchemy's [documentation for URL format](https://docs.sqlalchemy.org/en/20/core/engines.html#backend-specific-urls):
Expand Down
17 changes: 13 additions & 4 deletions pyiceberg/catalog/rest.py
Original file line number Diff line number Diff line change
Expand Up @@ -115,6 +115,7 @@ class Endpoints:
SIGV4_REGION = "rest.signing-region"
SIGV4_SERVICE = "rest.signing-name"
AUTH_URL = "rest.authorization-url"
HEADER_PREFIX = "header."

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

Expand Down Expand Up @@ -242,10 +243,7 @@ def _create_session(self) -> Session:
self._refresh_token(session, self.properties.get(TOKEN))

# Set HTTP headers
session.headers["Content-type"] = "application/json"
session.headers["X-Client-Version"] = ICEBERG_REST_SPEC_VERSION
session.headers["User-Agent"] = f"PyIceberg/{__version__}"
session.headers["X-Iceberg-Access-Delegation"] = "vended-credentials"
self._config_headers(session)

# Configure SigV4 Request Signing
if str(self.properties.get(SIGV4, False)).lower() == "true":
Expand Down Expand Up @@ -458,6 +456,17 @@ def _refresh_token(self, session: Optional[Session] = None, new_token: Optional[
if token := self.properties.get(TOKEN):
session.headers[AUTHORIZATION_HEADER] = f"{BEARER_PREFIX} {token}"

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

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)}

@retry(**_RETRY_ARGS)
def create_table(
self,
Expand Down
39 changes: 39 additions & 0 deletions tests/catalog/test_rest.py
Original file line number Diff line number Diff line change
Expand Up @@ -167,6 +167,45 @@ def test_config_200(requests_mock: Mocker) -> None:
assert history[1].url == "https://iceberg-test-catalog/v1/config?warehouse=s3%3A%2F%2Fsome-bucket"


def test_properties_sets_headers(requests_mock: Mocker) -> None:
requests_mock.get(
f"{TEST_URI}v1/config",
json={"defaults": {}, "overrides": {}},
status_code=200,
)

catalog = RestCatalog(
"rest", uri=TEST_URI, warehouse="s3://some-bucket", **{"header.Content-Type": "application/vnd.api+json"}
)

assert (
catalog._session.headers.get("Content-type") == "application/vnd.api+json"
), "Expected 'Content-Type' header to be 'application/vnd.api+json'"

assert (
requests_mock.last_request.headers["Content-type"] == "application/vnd.api+json"
), "Config request did not include expected 'Content-Type' header"


def test_config_sets_headers(requests_mock: Mocker) -> None:
namespace = "leden"
requests_mock.get(
f"{TEST_URI}v1/config",
json={"defaults": {"header.Content-Type": "application/vnd.api+json"}, "overrides": {}},
status_code=200,
)
requests_mock.post(f"{TEST_URI}v1/namespaces", json={"namespace": [namespace], "properties": {}}, status_code=200)
catalog = RestCatalog("rest", uri=TEST_URI, warehouse="s3://some-bucket")
catalog.create_namespace(namespace)

assert (
catalog._session.headers.get("Content-type") == "application/vnd.api+json"
), "Expected 'Content-Type' header to be 'application/vnd.api+json'"
assert (
requests_mock.last_request.headers["Content-type"] == "application/vnd.api+json"
), "Create namespace request did not include expected 'Content-Type' header"


def test_token_400(rest_mock: Mocker) -> None:
rest_mock.post(
f"{TEST_URI}v1/oauth/tokens",
Expand Down

0 comments on commit a1099d8

Please sign in to comment.