add optional dbt partner source

Author: marcodamore

dbtsl/api/graphql/client/asyncio.py

@@ -49,6 +49,7 @@ def __init__(
         auth_token: str,
         url_format: Optional[str] = None,
         timeout: Optional[Union[TimeoutOptions, float, int]] = None,
+        client_partner_source: Optional[str] = None,
     ):
         """Initialize the metadata client.
 
@@ -65,7 +66,7 @@ def __init__(
         limitations of `gql`'s `aiohttp` transport.
         See: https://github.com/graphql-python/gql/blob/b066e8944b0da0a4bbac6c31f43e5c3c7772cd51/gql/transport/aiohttp.py#L110
         """
-        super().__init__(server_host, environment_id, auth_token, url_format, timeout)
+        super().__init__(server_host, environment_id, auth_token, url_format, timeout, client_partner_source)
 
     @override
     def _create_transport(self, url: str, headers: Dict[str, str]) -> AIOHTTPTransport:

dbtsl/api/graphql/client/base.py

@@ -48,10 +48,11 @@ def _default_backoff(cls) -> ExponentialBackoff:
         )
 
     @classmethod
-    def _extra_headers(cls) -> Dict[str, str]:
-        return {
-            "user-agent": env.PLATFORM.user_agent,
-        }
+    def _extra_headers(cls, client_partner_source: Optional[str] = None) -> Dict[str, str]:
+        headers = {"user-agent": env.PLATFORM.user_agent}
+        if client_partner_source is not None:
+            headers["X-Dbt-Partner-Source"] = client_partner_source
+        return headers
 
     def __init__(  # noqa: D107
         self,
@@ -60,6 +61,7 @@ def __init__(  # noqa: D107
         auth_token: str,
         url_format: Optional[str] = None,
         timeout: Optional[Union[TimeoutOptions, float, int]] = None,
+        client_partner_source: Optional[str] = None,
     ):
         self.environment_id = environment_id
 
@@ -79,7 +81,7 @@ def __init__(  # noqa: D107
 
         headers = {
             "authorization": f"bearer {auth_token}",
-            **self._extra_headers(),
+            **self._extra_headers(client_partner_source),
         }
         transport = self._create_transport(url=server_url, headers=headers)
         self._gql = Client(transport=transport, execute_timeout=self.timeout.execute_timeout)

dbtsl/api/graphql/client/sync.py

@@ -38,6 +38,7 @@ def __init__(
         auth_token: str,
         url_format: Optional[str] = None,
         timeout: Optional[Union[TimeoutOptions, float, int]] = None,
+        client_partner_source: Optional[str] = None,
     ):
         """Initialize the metadata client.
 
@@ -49,11 +50,10 @@ def __init__(
                 into a full URL. If `None`, the default `https://{server_host}/api/graphql`
                 will be assumed.
             timeout: TimeoutOptions or total timeout (in seconds) for all GraphQL requests.
-
         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)
+        super().__init__(server_host, environment_id, auth_token, url_format, timeout, client_partner_source)
 
     @override
     def _create_transport(self, url: str, headers: Dict[str, str]) -> RequestsHTTPTransport:

dbtsl/client/base.py

@@ -40,6 +40,7 @@ def __init__(
         gql_factory: GraphQLClientFactory[TGQLClient],
         adbc_factory: ADBCClientFactory[TADBCClient],
         timeout: Optional[Union[TimeoutOptions, float, int]] = None,
+        requests_headers: Optional[dict[str, str]] = None,
     ) -> None:
         """Initialize the Semantic Layer client.
 
@@ -50,6 +51,7 @@ def __init__(
             gql_factory: class of the underlying GQL client
             adbc_factory: class of the underlying ADBC client
             timeout: `TimeoutOptions` or total timeout for the underlying GraphQL client.
+            requests_headers: additional headers to pass to the requests, optional
         """
         self._has_session = False
 

pyproject.toml

@@ -35,6 +35,7 @@ test = [
   "pytest-asyncio>=0.23.7,<0.24.0",
   "pytest-subtests>=0.12.1,<0.13.0",
   "pytest-mock>=3.14.0,<4.0.0",
+  "python-dotenv>=1.0.0,<2.0.0",
 ]
 
 [tool.hatch.build]

tests/conftest.py

@@ -4,6 +4,7 @@
 
 import pytest
 from gql import Client, gql
+from dotenv import load_dotenv
 from gql.utilities.serialize_variable_values import serialize_variable_values
 
 
@@ -69,6 +70,7 @@ def from_env(cls) -> "Credentials":
         - SL_TOKEN
         - SL_ENV_ID
         """
+        load_dotenv()
         return cls(
             host=os.environ["SL_HOST"],
             token=os.environ["SL_TOKEN"],