Skip to content

Commit

Permalink
Add user-agent header with provider version info
Browse files Browse the repository at this point in the history
  • Loading branch information
josh-fell committed Feb 22, 2022
1 parent ae265bf commit 2677bac
Showing 1 changed file with 12 additions and 0 deletions.
12 changes: 12 additions & 0 deletions airflow/providers/dbt/cloud/hooks/dbt.py
Original file line number Diff line number Diff line change
Expand Up @@ -64,13 +64,25 @@ def wrapper(*args, **kwargs) -> Callable:
return wrapper


def _get_provider_info() -> Tuple[str, str]:
from airflow.providers_manager import ProvidersManager

manager = ProvidersManager()
package_name = manager.hooks[DbtCloudHook.conn_type].package_name # type: ignore[union-attr]
provider = manager.providers[package_name]

return package_name, provider.version


class TokenAuth(AuthBase):
"""Helper class for Auth when executing requests."""

def __init__(self, token: str) -> None:
self.token = token

def __call__(self, request: PreparedRequest) -> PreparedRequest:
package_name, provider_version = _get_provider_info()
request.headers["User-Agent"] = f"{package_name}-v{provider_version}"
request.headers["Content-Type"] = "application/json"
request.headers["Authorization"] = f"Token {self.token}"

Expand Down

0 comments on commit 2677bac

Please sign in to comment.