From d851f58e1d99cce0a295d9ac84d2806f89cfe777 Mon Sep 17 00:00:00 2001 From: Will Deng Date: Mon, 18 Aug 2025 13:36:06 -0400 Subject: [PATCH 1/4] Add ability to pass extra headers as part of the request --- dbtsl/api/graphql/client/base.py | 4 ++++ dbtsl/api/graphql/client/sync.py | 3 ++- dbtsl/client/base.py | 4 +++- dbtsl/client/sync.py | 5 ++++- 4 files changed, 13 insertions(+), 3 deletions(-) diff --git a/dbtsl/api/graphql/client/base.py b/dbtsl/api/graphql/client/base.py index ef25c13..42d18e4 100644 --- a/dbtsl/api/graphql/client/base.py +++ b/dbtsl/api/graphql/client/base.py @@ -62,6 +62,7 @@ def __init__( # noqa: D107 timeout: Optional[Union[TimeoutOptions, float, int]] = None, *, lazy: bool, + extra_headers: Optional[Dict[str, str]] = None, ): self.environment_id = environment_id self.lazy = lazy @@ -83,6 +84,7 @@ def __init__( # noqa: D107 headers = { "authorization": f"bearer {auth_token}", **self._extra_headers(), + **(extra_headers or {}), } transport = self._create_transport(url=server_url, headers=headers) self._gql = Client(transport=transport, execute_timeout=self.timeout.execute_timeout) @@ -163,6 +165,7 @@ def __call__( timeout: Optional[Union[TimeoutOptions, float, int]] = None, *, lazy: bool, + extra_headers: Optional[Dict[str, str]] = None, ) -> TClient: """Initialize the Semantic Layer client. @@ -173,5 +176,6 @@ def __call__( url_format: the URL format string to construct the final URL with timeout: `TimeoutOptions` or total timeout lazy: lazy load large fields + extra_headers: extra headers to be sent with the request. """ pass diff --git a/dbtsl/api/graphql/client/sync.py b/dbtsl/api/graphql/client/sync.py index 1e996fa..7bc58e7 100644 --- a/dbtsl/api/graphql/client/sync.py +++ b/dbtsl/api/graphql/client/sync.py @@ -40,6 +40,7 @@ def __init__( timeout: Optional[Union[TimeoutOptions, float, int]] = None, *, lazy: bool, + extra_headers: Optional[Dict[str, str]] = None, ): """Initialize the metadata client. @@ -56,7 +57,7 @@ def __init__( NOTE: If `timeout` is a `TimeoutOptions`, the `tls_close_timeout` will not be used, since `requests` does not support TLS termination timeouts. """ - super().__init__(server_host, environment_id, auth_token, url_format, timeout, lazy=lazy) + super().__init__(server_host, environment_id, auth_token, url_format, timeout, lazy=lazy, extra_headers=extra_headers) @override def _create_transport(self, url: str, headers: Dict[str, str]) -> RequestsHTTPTransport: diff --git a/dbtsl/client/base.py b/dbtsl/client/base.py index 595ad1a..177ca69 100644 --- a/dbtsl/client/base.py +++ b/dbtsl/client/base.py @@ -1,5 +1,5 @@ from abc import ABC -from typing import Any, Generic, Optional, TypeVar, Union +from typing import Any, Dict, Generic, Optional, TypeVar, Union import dbtsl.env as env from dbtsl.api.adbc.client.base import ADBCClientFactory, BaseADBCClient @@ -43,6 +43,7 @@ def __init__( timeout: Optional[Union[TimeoutOptions, float, int]] = None, *, lazy: bool, + extra_headers: Optional[Dict[str, str]] = None, ) -> None: """Initialize the Semantic Layer client. @@ -66,6 +67,7 @@ def __init__( url_format=env.GRAPHQL_URL_FORMAT, timeout=timeout, lazy=lazy, + extra_headers=extra_headers, ) self._adbc = adbc_factory( server_host=host, diff --git a/dbtsl/client/sync.py b/dbtsl/client/sync.py index f783613..2bcbc9c 100644 --- a/dbtsl/client/sync.py +++ b/dbtsl/client/sync.py @@ -1,5 +1,5 @@ from contextlib import contextmanager -from typing import Iterator, Optional, Union +from typing import Dict, Iterator, Optional, Union from typing_extensions import Self @@ -28,6 +28,7 @@ def __init__( timeout: Optional[Union[TimeoutOptions, float, int]] = None, *, lazy: bool = False, + extra_headers: Optional[Dict[str, str]] = None, ) -> None: """Initialize the Semantic Layer client. @@ -37,6 +38,7 @@ def __init__( host: the Semantic Layer API host timeout: `TimeoutOptions` or total timeout for the underlying GraphQL client. lazy: if true, nested metadata queries will be need to be explicitly populated on-demand. + extra_headers: extra headers to be sent with the request. """ super().__init__( environment_id=environment_id, @@ -46,6 +48,7 @@ def __init__( adbc_factory=SyncADBCClient, timeout=timeout, lazy=lazy, + extra_headers=extra_headers, ) @contextmanager From d6ec4b6273d695265ae7d83b2be61de4e3e64ce7 Mon Sep 17 00:00:00 2001 From: Will Deng Date: Mon, 18 Aug 2025 13:49:26 -0400 Subject: [PATCH 2/4] Add ability to pass header for ADBC requests --- dbtsl/api/adbc/client/asyncio.py | 6 ++++-- dbtsl/api/adbc/client/base.py | 15 +++++++++++++-- dbtsl/api/adbc/client/sync.py | 6 ++++-- dbtsl/client/base.py | 1 + 4 files changed, 22 insertions(+), 6 deletions(-) diff --git a/dbtsl/api/adbc/client/asyncio.py b/dbtsl/api/adbc/client/asyncio.py index 52775c2..551bea4 100644 --- a/dbtsl/api/adbc/client/asyncio.py +++ b/dbtsl/api/adbc/client/asyncio.py @@ -1,6 +1,6 @@ import asyncio from contextlib import asynccontextmanager -from typing import AsyncIterator, Optional +from typing import AsyncIterator, Dict, Optional import pyarrow as pa from typing_extensions import Self, Unpack @@ -18,6 +18,7 @@ def __init__( environment_id: int, auth_token: str, url_format: Optional[str] = None, + extra_headers: Optional[Dict[str, str]] = None, ) -> None: """Initialize the ADBC client. @@ -29,8 +30,9 @@ def __init__( into a full URL. If `None`, the default `grpc+tls://{server_host}:443` will be assumed. + extra_headers: extra headers to be sent with the request. """ - super().__init__(server_host, environment_id, auth_token, url_format) + super().__init__(server_host, environment_id, auth_token, url_format, extra_headers=extra_headers) self._loop = asyncio.get_running_loop() @asynccontextmanager diff --git a/dbtsl/api/adbc/client/base.py b/dbtsl/api/adbc/client/base.py index 5797d66..0657058 100644 --- a/dbtsl/api/adbc/client/base.py +++ b/dbtsl/api/adbc/client/base.py @@ -2,7 +2,7 @@ from contextlib import AbstractContextManager from typing import Dict, Generic, Optional, Protocol, TypeVar, Union -from adbc_driver_flightsql import DatabaseOptions +from adbc_driver_flightsql import DatabaseOptions, StatementOptions from adbc_driver_flightsql.dbapi import Connection from adbc_driver_flightsql.dbapi import connect as adbc_connect # pyright: ignore[reportUnknownVariableType] from adbc_driver_manager import AdbcStatusCode, ProgrammingError @@ -33,11 +33,13 @@ def __init__( # noqa: D107 environment_id: int, auth_token: str, url_format: Optional[str] = None, + extra_headers: Optional[Dict[str, str]] = None, ) -> None: url_format = url_format or self.DEFAULT_URL_FORMAT self._conn_str = url_format.format(server_host=server_host) self._environment_id = environment_id self._auth_token = auth_token + self._extra_headers = extra_headers or {} self._conn_unsafe: Union[Connection, None] = None @@ -48,6 +50,7 @@ def _get_connection_context_manager(self) -> AbstractContextManager[Connection]: DatabaseOptions.AUTHORIZATION_HEADER.value: f"Bearer {self._auth_token}", f"{DatabaseOptions.RPC_CALL_HEADER_PREFIX.value}environmentid": str(self._environment_id), **self._extra_db_kwargs(), + **{f"{StatementOptions.RPC_CALL_HEADER_PREFIX.value}{k}": v for k, v in self._extra_headers.items()} }, ) @@ -88,7 +91,14 @@ def has_session(self) -> bool: class ADBCClientFactory(Protocol, Generic[TClient]): # noqa: D101 @abstractmethod - def __call__(self, server_host: str, environment_id: int, auth_token: str, url_format: str) -> TClient: + def __call__( + self, + server_host: str, + environment_id: int, + auth_token: str, + url_format: str, + extra_headers: Optional[Dict[str, str]] = None, + ) -> TClient: """Initialize the Semantic Layer client. Args: @@ -96,5 +106,6 @@ def __call__(self, server_host: str, environment_id: int, auth_token: str, url_f environment_id: your dbt environment ID auth_token: the API auth token url_format: the URL format string to construct the final URL with + extra_headers: extra headers to be sent with the request. """ pass diff --git a/dbtsl/api/adbc/client/sync.py b/dbtsl/api/adbc/client/sync.py index f38046f..616c037 100644 --- a/dbtsl/api/adbc/client/sync.py +++ b/dbtsl/api/adbc/client/sync.py @@ -1,5 +1,5 @@ from contextlib import contextmanager -from typing import Iterator, Optional +from typing import Dict, Iterator, Optional import pyarrow as pa from typing_extensions import Self, Unpack @@ -17,6 +17,7 @@ def __init__( environment_id: int, auth_token: str, url_format: Optional[str] = None, + extra_headers: Optional[Dict[str, str]] = None, ) -> None: """Initialize the ADBC client. @@ -28,8 +29,9 @@ def __init__( into a full URL. If `None`, the default `grpc+tls://{server_host}:443` will be assumed. + extra_headers: extra headers to be sent with the request. """ - super().__init__(server_host, environment_id, auth_token, url_format) + super().__init__(server_host, environment_id, auth_token, url_format, extra_headers=extra_headers) @contextmanager def session(self) -> Iterator[Self]: diff --git a/dbtsl/client/base.py b/dbtsl/client/base.py index 177ca69..bce7e73 100644 --- a/dbtsl/client/base.py +++ b/dbtsl/client/base.py @@ -74,6 +74,7 @@ def __init__( environment_id=environment_id, auth_token=auth_token, url_format=env.ADBC_URL_FORMAT, + extra_headers=extra_headers, ) @property From 2d7858b2bbaf2d859fca49cba574d570dbe3fd3b Mon Sep 17 00:00:00 2001 From: Will Deng Date: Mon, 18 Aug 2025 14:01:43 -0400 Subject: [PATCH 3/4] lint --- dbtsl/api/adbc/client/base.py | 2 +- dbtsl/api/graphql/client/sync.py | 4 +++- 2 files changed, 4 insertions(+), 2 deletions(-) diff --git a/dbtsl/api/adbc/client/base.py b/dbtsl/api/adbc/client/base.py index 0657058..e39bbcd 100644 --- a/dbtsl/api/adbc/client/base.py +++ b/dbtsl/api/adbc/client/base.py @@ -50,7 +50,7 @@ def _get_connection_context_manager(self) -> AbstractContextManager[Connection]: DatabaseOptions.AUTHORIZATION_HEADER.value: f"Bearer {self._auth_token}", f"{DatabaseOptions.RPC_CALL_HEADER_PREFIX.value}environmentid": str(self._environment_id), **self._extra_db_kwargs(), - **{f"{StatementOptions.RPC_CALL_HEADER_PREFIX.value}{k}": v for k, v in self._extra_headers.items()} + **{f"{StatementOptions.RPC_CALL_HEADER_PREFIX.value}{k}": v for k, v in self._extra_headers.items()}, }, ) diff --git a/dbtsl/api/graphql/client/sync.py b/dbtsl/api/graphql/client/sync.py index 7bc58e7..6778015 100644 --- a/dbtsl/api/graphql/client/sync.py +++ b/dbtsl/api/graphql/client/sync.py @@ -57,7 +57,9 @@ def __init__( NOTE: If `timeout` is a `TimeoutOptions`, the `tls_close_timeout` will not be used, since `requests` does not support TLS termination timeouts. """ - super().__init__(server_host, environment_id, auth_token, url_format, timeout, lazy=lazy, extra_headers=extra_headers) + super().__init__( + server_host, environment_id, auth_token, url_format, timeout, lazy=lazy, extra_headers=extra_headers + ) @override def _create_transport(self, url: str, headers: Dict[str, str]) -> RequestsHTTPTransport: From c5e475e996717f21e76543fad86c4f3f720b3d07 Mon Sep 17 00:00:00 2001 From: Will Deng Date: Mon, 18 Aug 2025 14:02:28 -0400 Subject: [PATCH 4/4] changelog --- .changes/unreleased/Under the Hood-20250818-140222.yaml | 3 +++ 1 file changed, 3 insertions(+) create mode 100644 .changes/unreleased/Under the Hood-20250818-140222.yaml diff --git a/.changes/unreleased/Under the Hood-20250818-140222.yaml b/.changes/unreleased/Under the Hood-20250818-140222.yaml new file mode 100644 index 0000000..9bdf4d7 --- /dev/null +++ b/.changes/unreleased/Under the Hood-20250818-140222.yaml @@ -0,0 +1,3 @@ +kind: Under the Hood +body: Add ability to add extra headers during a request +time: 2025-08-18T14:02:22.72837-04:00