From daef55fe289e377d5f8d0cce6dc71bf7a5703ae5 Mon Sep 17 00:00:00 2001
From: Marwan Sarieddine <sarieddine.marwan@gmail.com>
Date: Fri, 28 Jun 2024 17:55:57 +0300
Subject: [PATCH] use airflow config instead of env vars for setting telemetry

---
 README.md                           | 2 +-
 anyscale_provider/hooks/anyscale.py | 5 +++--
 2 files changed, 4 insertions(+), 3 deletions(-)

diff --git a/README.md b/README.md
index ad384d1..3695376 100644
--- a/README.md
+++ b/README.md
@@ -193,7 +193,7 @@ We will not sell data or buy data about you.
 
 You will always be able to disable the usage stats collection.
 
-To disable usage stats collection, set the environment variable `ANYSCALE__AIRFLOW_TELEMETRY_ENABLED` to `false`.
+To disable usage stats collection, set the configuration anyscale.telemetry_enabled to False in your airflow.cfg file.
 
 
 ### Changelog
diff --git a/anyscale_provider/hooks/anyscale.py b/anyscale_provider/hooks/anyscale.py
index 7ab9049..ab45a08 100644
--- a/anyscale_provider/hooks/anyscale.py
+++ b/anyscale_provider/hooks/anyscale.py
@@ -5,8 +5,10 @@
 from functools import cached_property
 from typing import Any
 
+from airflow.configuration import conf
 from airflow.exceptions import AirflowException
 from airflow.hooks.base import BaseHook
+
 from anyscale import Anyscale
 from anyscale.job.models import JobConfig, JobStatus
 from anyscale.service.models import ServiceConfig, ServiceStatus
@@ -44,8 +46,7 @@ def client(self) -> Anyscale:
 
         # 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"]
+        telemetry_enabled = conf.getboolean("anyscale", "telemetry_enabled", fallback=True)
         if telemetry_enabled:
             headers["X-Anyscale-Source"] = "airflow"