Skip to content

Commit

Permalink
Add telemetry headers
Browse files Browse the repository at this point in the history
- Capture telemetry via a header at SDK init
- Enable users to opt-out of telemetry via setting env_var
  • Loading branch information
marwan116 committed Jun 17, 2024
1 parent 922c2eb commit 275e75e
Showing 1 changed file with 8 additions and 2 deletions.
10 changes: 8 additions & 2 deletions anyscale_provider/hooks/anyscale.py
Original file line number Diff line number Diff line change
Expand Up @@ -120,8 +120,14 @@ def __init__(self, conn_id: str = default_conn_name, **kwargs: Any) -> None:
if not token:
raise AirflowException(f"Missing API token for connection id {self.conn_id}")

self.sdk = Anyscale(auth_token=token)

# Add custom headers if telemetry is enabled - by default telemetry is enabled.
headers = {}
telemetry_env = os.getenv("ANYSCALE__AIRFLOW_TELEMETRY_ENABLED", "true")
telemetry_enabled = telemetry_env.lower() in ["true", "1", "yes", "on"]
if telemetry_enabled:
headers["X-Anyscale-Source"] = "astronomer-sdk"

self.sdk = Anyscale(auth_token=token, headers=headers)

@classmethod
def get_ui_field_behaviour(cls) -> Dict[str, Any]:
Expand Down

0 comments on commit 275e75e

Please sign in to comment.