From f9cb8240ec7e64e4ffb482ff8d16c4ffe5e7798f Mon Sep 17 00:00:00 2001
From: Sai Shree Pradhan <saishree.pradhan@databricks.com>
Date: Tue, 22 Jul 2025 10:12:30 +0530
Subject: [PATCH 1/9] feature_flag

Signed-off-by: Sai Shree Pradhan <saishree.pradhan@databricks.com>
---
 src/databricks/sql/client.py                  |  17 +-
 src/databricks/sql/common/feature_flag.py     | 177 ++++++++++++++++++
 .../sql/telemetry/telemetry_client.py         |  21 ++-
 3 files changed, 208 insertions(+), 7 deletions(-)
 create mode 100644 src/databricks/sql/common/feature_flag.py

diff --git a/src/databricks/sql/client.py b/src/databricks/sql/client.py
index e4166f117..ecde88533 100755
--- a/src/databricks/sql/client.py
+++ b/src/databricks/sql/client.py
@@ -245,12 +245,6 @@ def read(self) -> Optional[OAuthToken]:
         self.use_cloud_fetch = kwargs.get("use_cloud_fetch", True)
         self._cursors = []  # type: List[Cursor]
 
-        self.server_telemetry_enabled = True
-        self.client_telemetry_enabled = kwargs.get("enable_telemetry", False)
-        self.telemetry_enabled = (
-            self.client_telemetry_enabled and self.server_telemetry_enabled
-        )
-
         self.session = Session(
             server_hostname,
             http_path,
@@ -268,6 +262,17 @@ def read(self) -> Optional[OAuthToken]:
         )
         self.staging_allowed_local_path = kwargs.get("staging_allowed_local_path", None)
 
+        self.client_telemetry_enabled = kwargs.get("enable_telemetry", False)
+        if self.client_telemetry_enabled:
+            self.server_telemetry_enabled = TelemetryHelper.is_server_telemetry_enabled(
+                self
+            )
+            self.telemetry_enabled = (
+                self.client_telemetry_enabled and self.server_telemetry_enabled
+            )
+        else:
+            self.telemetry_enabled = False
+
         TelemetryClientFactory.initialize_telemetry_client(
             telemetry_enabled=self.telemetry_enabled,
             session_id_hex=self.get_session_id_hex(),
diff --git a/src/databricks/sql/common/feature_flag.py b/src/databricks/sql/common/feature_flag.py
new file mode 100644
index 000000000..4ad445e92
--- /dev/null
+++ b/src/databricks/sql/common/feature_flag.py
@@ -0,0 +1,177 @@
+# databricks/sql/common/feature_flags.py
+
+import threading
+import time
+import requests
+from dataclasses import dataclass, field
+from concurrent.futures import ThreadPoolExecutor
+from typing import Dict, Optional, List, Any, TYPE_CHECKING
+
+if TYPE_CHECKING:
+    from databricks.sql.client import Connection
+
+
+@dataclass
+class FeatureFlagEntry:
+    """Represents a single feature flag from the server response."""
+
+    name: str
+    value: str
+
+
+@dataclass
+class FeatureFlagsResponse:
+    """Represents the full JSON response from the feature flag endpoint."""
+
+    flags: List[FeatureFlagEntry] = field(default_factory=list)
+    ttl_seconds: Optional[int] = None
+
+    @classmethod
+    def from_dict(cls, data: Dict[str, Any]) -> "FeatureFlagsResponse":
+        """Factory method to create an instance from a dictionary (parsed JSON)."""
+        flags_data = data.get("flags", [])
+        flags_list = [FeatureFlagEntry(**flag) for flag in flags_data]
+        return cls(flags=flags_list, ttl_seconds=data.get("ttl_seconds"))
+
+
+# --- Constants ---
+FEATURE_FLAGS_ENDPOINT_SUFFIX_FORMAT = (
+    "/api/2.0/connector-service/feature-flags/PYTHON/{}"
+)
+DEFAULT_TTL_SECONDS = 900  # 15 minutes
+REFRESH_BEFORE_EXPIRY_SECONDS = 10  # Start proactive refresh 10s before expiry
+
+
+class FeatureFlagsContext:
+    """
+    Manages fetching and caching of server-side feature flags for a connection.
+
+    1. The very first check for any flag is a synchronous, BLOCKING operation.
+    2. Subsequent refreshes (triggered near TTL expiry) are done asynchronously
+       in the background, returning stale data until the refresh completes.
+    """
+
+    def __init__(self, connection: "Connection", executor: ThreadPoolExecutor):
+        from databricks.sql import __version__
+
+        self._connection = connection
+        self._executor = executor  # Used for ASYNCHRONOUS refreshes
+        self._lock = threading.RLock()
+
+        # Cache state: `None` indicates the cache has never been loaded.
+        self._flags: Optional[Dict[str, str]] = None
+        self._ttl_seconds: int = DEFAULT_TTL_SECONDS
+        self._last_refresh_time: float = 0
+
+        endpoint_suffix = FEATURE_FLAGS_ENDPOINT_SUFFIX_FORMAT.format(__version__)
+        self._feature_flag_endpoint = (
+            f"https://{self._connection.session.host}{endpoint_suffix}"
+        )
+
+    def _is_refresh_needed(self) -> bool:
+        """Checks if the cache is due for a proactive background refresh."""
+        if self._flags is None:
+            return False  # Not eligible for refresh until loaded once.
+
+        refresh_threshold = self._last_refresh_time + (
+            self._ttl_seconds - REFRESH_BEFORE_EXPIRY_SECONDS
+        )
+        return time.monotonic() > refresh_threshold
+
+    def is_feature_enabled(self, name: str, default_value: bool) -> bool:
+        """
+        Checks if a feature is enabled.
+        - BLOCKS on the first call until flags are fetched.
+        - Returns cached values on subsequent calls, triggering non-blocking refreshes if needed.
+        """
+        with self._lock:
+            # If cache has never been loaded, perform a synchronous, blocking fetch.
+            if self._flags is None:
+                self._refresh_flags()
+
+            # If a proactive background refresh is needed, start one. This is non-blocking.
+            elif self._is_refresh_needed():
+                # We don't check for an in-flight refresh; the executor queues the task, which is safe.
+                self._executor.submit(self._refresh_flags)
+
+            # Now, return the value from the populated cache.
+            flag_value = self._flags.get(name)
+            if flag_value is None:
+                return default_value
+            return flag_value.lower() == "true"
+
+    def _refresh_flags(self):
+        """Performs a synchronous network request to fetch and update flags."""
+        headers = {}
+        try:
+            # Authenticate the request
+            self._connection.session.auth_provider.add_headers(headers)
+            headers["User-Agent"] = self._connection.session.useragent_header
+
+            response = requests.get(
+                self._feature_flag_endpoint, headers=headers, timeout=30
+            )
+
+            if response.status_code == 200:
+                ff_response = FeatureFlagsResponse.from_dict(response.json())
+                self._update_cache_from_response(ff_response)
+            else:
+                # On failure, initialize with an empty dictionary to prevent re-blocking.
+                if self._flags is None:
+                    self._flags = {}
+
+        except Exception as e:
+            # On exception, initialize with an empty dictionary to prevent re-blocking.
+            if self._flags is None:
+                self._flags = {}
+
+    def _update_cache_from_response(self, ff_response: FeatureFlagsResponse):
+        """Atomically updates the internal cache state from a successful server response."""
+        with self._lock:
+            self._flags = {flag.name: flag.value for flag in ff_response.flags}
+            if ff_response.ttl_seconds is not None and ff_response.ttl_seconds > 0:
+                self._ttl_seconds = ff_response.ttl_seconds
+            self._last_refresh_time = time.monotonic()
+
+
+class FeatureFlagsContextFactory:
+    """
+    Manages a singleton instance of FeatureFlagsContext per connection session.
+    Also manages a shared ThreadPoolExecutor for all background refresh operations.
+    """
+
+    _context_map: Dict[str, FeatureFlagsContext] = {}
+    _executor: Optional[ThreadPoolExecutor] = None
+    _lock = threading.Lock()
+
+    @classmethod
+    def _initialize(cls):
+        """Initializes the shared executor for async refreshes if it doesn't exist."""
+        if cls._executor is None:
+            cls._executor = ThreadPoolExecutor(
+                max_workers=3, thread_name_prefix="feature-flag-refresher"
+            )
+
+    @classmethod
+    def get_instance(cls, connection: "Connection") -> FeatureFlagsContext:
+        """Gets or creates a FeatureFlagsContext for the given connection."""
+        with cls._lock:
+            cls._initialize()
+            # Use the unique session ID as the key
+            key = connection.get_session_id_hex()
+            if key not in cls._context_map:
+                cls._context_map[key] = FeatureFlagsContext(connection, cls._executor)
+            return cls._context_map[key]
+
+    @classmethod
+    def remove_instance(cls, connection: "Connection"):
+        """Removes the context for a given connection and shuts down the executor if no clients remain."""
+        with cls._lock:
+            key = connection.get_session_id_hex()
+            if key in cls._context_map:
+                cls._context_map.pop(key, None)
+
+            # If this was the last active context, clean up the thread pool.
+            if not cls._context_map and cls._executor is not None:
+                cls._executor.shutdown(wait=False)
+                cls._executor = None
diff --git a/src/databricks/sql/telemetry/telemetry_client.py b/src/databricks/sql/telemetry/telemetry_client.py
index 5eb8c6ed0..dd0285376 100644
--- a/src/databricks/sql/telemetry/telemetry_client.py
+++ b/src/databricks/sql/telemetry/telemetry_client.py
@@ -3,7 +3,7 @@
 import requests
 import logging
 from concurrent.futures import ThreadPoolExecutor
-from typing import Dict, Optional
+from typing import Dict, Optional, TYPE_CHECKING
 from databricks.sql.telemetry.models.event import (
     TelemetryEvent,
     DriverSystemConfiguration,
@@ -30,6 +30,10 @@
 import uuid
 import locale
 from databricks.sql.telemetry.utils import BaseTelemetryClient
+from databricks.sql.common.feature_flag import FeatureFlagsContextFactory
+
+if TYPE_CHECKING:
+    from databricks.sql.client import Connection
 
 logger = logging.getLogger(__name__)
 
@@ -38,6 +42,9 @@ class TelemetryHelper:
     """Helper class for getting telemetry related information."""
 
     _DRIVER_SYSTEM_CONFIGURATION = None
+    TELEMETRY_FEATURE_FLAG_NAME = (
+        "databricks.partnerplatform.clientConfigsFeatureFlags.enableTelemetry"
+    )
 
     @classmethod
     def get_driver_system_configuration(cls) -> DriverSystemConfiguration:
@@ -92,6 +99,18 @@ def get_auth_flow(auth_provider):
         else:
             return None
 
+    @staticmethod
+    def is_server_telemetry_enabled(connection: "Connection") -> bool:
+        """
+        Checks if the server-side feature flag for telemetry is enabled.
+        This is a BLOCKING call on the first check per connection.
+        """
+        context = FeatureFlagsContextFactory.get_instance(connection)
+
+        return context.is_feature_enabled(
+            TelemetryHelper.TELEMETRY_FEATURE_FLAG_NAME, default_value=False
+        )
+
 
 class NoopTelemetryClient(BaseTelemetryClient):
     """

From 7785d4382800709b4744ce15c5f755dacca6ddb9 Mon Sep 17 00:00:00 2001
From: Sai Shree Pradhan <saishree.pradhan@databricks.com>
Date: Tue, 22 Jul 2025 10:21:10 +0530
Subject: [PATCH 2/9] fix static type check

Signed-off-by: Sai Shree Pradhan <saishree.pradhan@databricks.com>
---
 src/databricks/sql/common/feature_flag.py | 2 ++
 1 file changed, 2 insertions(+)

diff --git a/src/databricks/sql/common/feature_flag.py b/src/databricks/sql/common/feature_flag.py
index 4ad445e92..817e6e49f 100644
--- a/src/databricks/sql/common/feature_flag.py
+++ b/src/databricks/sql/common/feature_flag.py
@@ -94,6 +94,8 @@ def is_feature_enabled(self, name: str, default_value: bool) -> bool:
                 # We don't check for an in-flight refresh; the executor queues the task, which is safe.
                 self._executor.submit(self._refresh_flags)
 
+            assert self._flags is not None
+
             # Now, return the value from the populated cache.
             flag_value = self._flags.get(name)
             if flag_value is None:

From f0cdfab0b411c7d941596c5508d2ab4e84e89482 Mon Sep 17 00:00:00 2001
From: Sai Shree Pradhan <saishree.pradhan@databricks.com>
Date: Tue, 22 Jul 2025 10:25:04 +0530
Subject: [PATCH 3/9] fix static type check

Signed-off-by: Sai Shree Pradhan <saishree.pradhan@databricks.com>
---
 src/databricks/sql/common/feature_flag.py | 2 ++
 1 file changed, 2 insertions(+)

diff --git a/src/databricks/sql/common/feature_flag.py b/src/databricks/sql/common/feature_flag.py
index 817e6e49f..2c232f084 100644
--- a/src/databricks/sql/common/feature_flag.py
+++ b/src/databricks/sql/common/feature_flag.py
@@ -159,6 +159,8 @@ def get_instance(cls, connection: "Connection") -> FeatureFlagsContext:
         """Gets or creates a FeatureFlagsContext for the given connection."""
         with cls._lock:
             cls._initialize()
+            assert cls._executor is not None
+
             # Use the unique session ID as the key
             key = connection.get_session_id_hex()
             if key not in cls._context_map:

From 08c0bd87b9411da1318dfdede50955991206f87f Mon Sep 17 00:00:00 2001
From: Sai Shree Pradhan <saishree.pradhan@databricks.com>
Date: Wed, 23 Jul 2025 11:15:42 +0530
Subject: [PATCH 4/9] force enable telemetry

Signed-off-by: Sai Shree Pradhan <saishree.pradhan@databricks.com>
---
 src/databricks/sql/client.py                  | 13 +++---------
 src/databricks/sql/common/feature_flag.py     |  2 +-
 .../sql/telemetry/telemetry_client.py         | 21 ++++++++++---------
 3 files changed, 15 insertions(+), 21 deletions(-)

diff --git a/src/databricks/sql/client.py b/src/databricks/sql/client.py
index ecde88533..ed510fe9f 100755
--- a/src/databricks/sql/client.py
+++ b/src/databricks/sql/client.py
@@ -262,16 +262,9 @@ def read(self) -> Optional[OAuthToken]:
         )
         self.staging_allowed_local_path = kwargs.get("staging_allowed_local_path", None)
 
-        self.client_telemetry_enabled = kwargs.get("enable_telemetry", False)
-        if self.client_telemetry_enabled:
-            self.server_telemetry_enabled = TelemetryHelper.is_server_telemetry_enabled(
-                self
-            )
-            self.telemetry_enabled = (
-                self.client_telemetry_enabled and self.server_telemetry_enabled
-            )
-        else:
-            self.telemetry_enabled = False
+        self.force_enable_telemetry = kwargs.get("force_enable_telemetry", False)
+        self.enable_telemetry = kwargs.get("enable_telemetry", False)
+        self.telemetry_enabled = TelemetryHelper.is_telemetry_enabled(self)
 
         TelemetryClientFactory.initialize_telemetry_client(
             telemetry_enabled=self.telemetry_enabled,
diff --git a/src/databricks/sql/common/feature_flag.py b/src/databricks/sql/common/feature_flag.py
index 2c232f084..ddc8d0466 100644
--- a/src/databricks/sql/common/feature_flag.py
+++ b/src/databricks/sql/common/feature_flag.py
@@ -1,4 +1,4 @@
-# databricks/sql/common/feature_flags.py
+# TODO: Test this when server-side feature flag is available
 
 import threading
 import time
diff --git a/src/databricks/sql/telemetry/telemetry_client.py b/src/databricks/sql/telemetry/telemetry_client.py
index dd0285376..93aeeb464 100644
--- a/src/databricks/sql/telemetry/telemetry_client.py
+++ b/src/databricks/sql/telemetry/telemetry_client.py
@@ -100,16 +100,17 @@ def get_auth_flow(auth_provider):
             return None
 
     @staticmethod
-    def is_server_telemetry_enabled(connection: "Connection") -> bool:
-        """
-        Checks if the server-side feature flag for telemetry is enabled.
-        This is a BLOCKING call on the first check per connection.
-        """
-        context = FeatureFlagsContextFactory.get_instance(connection)
-
-        return context.is_feature_enabled(
-            TelemetryHelper.TELEMETRY_FEATURE_FLAG_NAME, default_value=False
-        )
+    def is_telemetry_enabled(connection: "Connection") -> bool:
+        if connection.force_enable_telemetry:
+            return True
+
+        if connection.enable_telemetry:
+            context = FeatureFlagsContextFactory.get_instance(connection)
+            return context.is_feature_enabled(
+                TelemetryHelper.TELEMETRY_FEATURE_FLAG_NAME, default_value=False
+            )
+        else:
+            return False
 
 
 class NoopTelemetryClient(BaseTelemetryClient):

From 0d300abff1cead09a1991ec92f122e5ffca7434c Mon Sep 17 00:00:00 2001
From: Sai Shree Pradhan <saishree.pradhan@databricks.com>
Date: Thu, 31 Jul 2025 11:02:16 +0530
Subject: [PATCH 5/9] added flag

Signed-off-by: Sai Shree Pradhan <saishree.pradhan@databricks.com>
---
 src/databricks/sql/telemetry/telemetry_client.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/src/databricks/sql/telemetry/telemetry_client.py b/src/databricks/sql/telemetry/telemetry_client.py
index ba9d31e33..8021acabc 100644
--- a/src/databricks/sql/telemetry/telemetry_client.py
+++ b/src/databricks/sql/telemetry/telemetry_client.py
@@ -49,7 +49,7 @@ class TelemetryHelper:
 
     _DRIVER_SYSTEM_CONFIGURATION = None
     TELEMETRY_FEATURE_FLAG_NAME = (
-        "databricks.partnerplatform.clientConfigsFeatureFlags.enableTelemetry"
+        "databricks.partnerplatform.clientConfigsFeatureFlags.enableTelemetryForPythonDriver"
     )
 
     @classmethod

From 6ceb092bf3ce004ceab2edbcf3d5ca2209d46a1d Mon Sep 17 00:00:00 2001
From: Sai Shree Pradhan <saishree.pradhan@databricks.com>
Date: Thu, 31 Jul 2025 11:49:34 +0530
Subject: [PATCH 6/9] linting

Signed-off-by: Sai Shree Pradhan <saishree.pradhan@databricks.com>
---
 src/databricks/sql/common/feature_flag.py        | 2 --
 src/databricks/sql/telemetry/telemetry_client.py | 4 +---
 2 files changed, 1 insertion(+), 5 deletions(-)

diff --git a/src/databricks/sql/common/feature_flag.py b/src/databricks/sql/common/feature_flag.py
index ddc8d0466..f4a55b049 100644
--- a/src/databricks/sql/common/feature_flag.py
+++ b/src/databricks/sql/common/feature_flag.py
@@ -1,5 +1,3 @@
-# TODO: Test this when server-side feature flag is available
-
 import threading
 import time
 import requests
diff --git a/src/databricks/sql/telemetry/telemetry_client.py b/src/databricks/sql/telemetry/telemetry_client.py
index 11629485e..b5a89755d 100644
--- a/src/databricks/sql/telemetry/telemetry_client.py
+++ b/src/databricks/sql/telemetry/telemetry_client.py
@@ -48,9 +48,7 @@ class TelemetryHelper:
     """Helper class for getting telemetry related information."""
 
     _DRIVER_SYSTEM_CONFIGURATION = None
-    TELEMETRY_FEATURE_FLAG_NAME = (
-        "databricks.partnerplatform.clientConfigsFeatureFlags.enableTelemetryForPythonDriver"
-    )
+    TELEMETRY_FEATURE_FLAG_NAME = "databricks.partnerplatform.clientConfigsFeatureFlags.enableTelemetryForPythonDriver"
 
     @classmethod
     def get_driver_system_configuration(cls) -> DriverSystemConfiguration:

From 725cce9137e7dd84134064de242ea0117e4b9b9e Mon Sep 17 00:00:00 2001
From: Sai Shree Pradhan <saishree.pradhan@databricks.com>
Date: Thu, 31 Jul 2025 15:30:59 +0530
Subject: [PATCH 7/9] tests

Signed-off-by: Sai Shree Pradhan <saishree.pradhan@databricks.com>
---
 tests/unit/test_telemetry.py | 91 ++++++++++++++++++++++++++++++++++--
 1 file changed, 88 insertions(+), 3 deletions(-)

diff --git a/tests/unit/test_telemetry.py b/tests/unit/test_telemetry.py
index d0e28c18d..d516a54fe 100644
--- a/tests/unit/test_telemetry.py
+++ b/tests/unit/test_telemetry.py
@@ -7,7 +7,6 @@
     NoopTelemetryClient,
     TelemetryClientFactory,
     TelemetryHelper,
-    BaseTelemetryClient,
 )
 from databricks.sql.telemetry.models.enums import AuthMech, AuthFlow
 from databricks.sql.auth.authenticators import (
@@ -15,6 +14,7 @@
     DatabricksOAuthProvider,
     ExternalAuthProvider,
 )
+from databricks import sql
 
 
 @pytest.fixture
@@ -311,8 +311,6 @@ def test_connection_failure_sends_correct_telemetry_payload(
         mock_session.side_effect = Exception(error_message)
 
         try:
-            from databricks import sql
-
             sql.connect(server_hostname="test-host", http_path="/test-path")
         except Exception as e:
             assert str(e) == error_message
@@ -321,3 +319,90 @@ def test_connection_failure_sends_correct_telemetry_payload(
         call_arguments = mock_export_failure_log.call_args
         assert call_arguments[0][0] == "Exception"
         assert call_arguments[0][1] == error_message
+
+
+@patch("databricks.sql.client.Session")
+class TestTelemetryFeatureFlag:
+    """Tests the interaction between the telemetry feature flag and connection parameters."""
+
+    def _mock_ff_response(self, mock_requests_get, enabled: bool):
+        """Helper to configure the mock response for the feature flag endpoint."""
+        mock_response = MagicMock()
+        mock_response.status_code = 200
+        payload = {
+            "flags": [
+                {
+                    "name": "databricks.partnerplatform.clientConfigsFeatureFlags.enableTelemetryForPythonDriver",
+                    "value": str(enabled).lower(),
+                }
+            ],
+            "ttl_seconds": 3600,
+        }
+        mock_response.json.return_value = payload
+        mock_requests_get.return_value = mock_response
+
+    @patch("databricks.sql.common.feature_flag.requests.get")
+    def test_telemetry_enabled_when_flag_is_true(
+        self, mock_requests_get, MockSession
+    ):
+        """Telemetry should be ON when enable_telemetry=True and server flag is 'true'."""
+        self._mock_ff_response(mock_requests_get, enabled=True)
+        mock_session_instance = MockSession.return_value
+        mock_session_instance.guid_hex = "test-session-ff-true"
+        mock_session_instance.auth_provider = AccessTokenAuthProvider("token")
+
+        conn = sql.client.Connection(
+            server_hostname="test",
+            http_path="test",
+            access_token="test",
+            enable_telemetry=True,
+        )
+
+        assert conn.telemetry_enabled is True
+        mock_requests_get.assert_called_once()
+        client = TelemetryClientFactory.get_telemetry_client("test-session-ff-true")
+        assert isinstance(client, TelemetryClient)
+
+    @patch("databricks.sql.common.feature_flag.requests.get")
+    def test_telemetry_disabled_when_flag_is_false(
+        self, mock_requests_get, MockSession
+    ):
+        """Telemetry should be OFF when enable_telemetry=True but server flag is 'false'."""
+        self._mock_ff_response(mock_requests_get, enabled=False)
+        mock_session_instance = MockSession.return_value
+        mock_session_instance.guid_hex = "test-session-ff-false"
+        mock_session_instance.auth_provider = AccessTokenAuthProvider("token")
+
+        conn = sql.client.Connection(
+            server_hostname="test",
+            http_path="test",
+            access_token="test",
+            enable_telemetry=True,
+        )
+
+        assert conn.telemetry_enabled is False
+        mock_requests_get.assert_called_once()
+        client = TelemetryClientFactory.get_telemetry_client("test-session-ff-false")
+        assert isinstance(client, NoopTelemetryClient)
+
+    @patch("databricks.sql.common.feature_flag.requests.get")
+    def test_telemetry_disabled_when_flag_request_fails(
+        self, mock_requests_get, MockSession
+    ):
+        """Telemetry should default to OFF if the feature flag network request fails."""
+        mock_requests_get.side_effect = Exception("Network is down")
+        mock_session_instance = MockSession.return_value
+        mock_session_instance.guid_hex = "test-session-ff-fail"
+        mock_session_instance.auth_provider = AccessTokenAuthProvider("token")
+
+        conn = sql.client.Connection(
+            server_hostname="test",
+            http_path="test",
+            access_token="test",
+            enable_telemetry=True,
+        )
+
+        assert conn.telemetry_enabled is False
+        mock_requests_get.assert_called_once()
+        client = TelemetryClientFactory.get_telemetry_client("test-session-ff-fail")
+        assert isinstance(client, NoopTelemetryClient)
\ No newline at end of file

From 6f47cbfa8e31175bf2b548856d7974a60742ffd4 Mon Sep 17 00:00:00 2001
From: Sai Shree Pradhan <saishree.pradhan@databricks.com>
Date: Fri, 1 Aug 2025 11:40:17 +0530
Subject: [PATCH 8/9] changed flag value to be of any type

Signed-off-by: Sai Shree Pradhan <saishree.pradhan@databricks.com>
---
 src/databricks/sql/common/feature_flag.py        | 7 ++-----
 src/databricks/sql/telemetry/telemetry_client.py | 3 ++-
 2 files changed, 4 insertions(+), 6 deletions(-)

diff --git a/src/databricks/sql/common/feature_flag.py b/src/databricks/sql/common/feature_flag.py
index f4a55b049..53add9253 100644
--- a/src/databricks/sql/common/feature_flag.py
+++ b/src/databricks/sql/common/feature_flag.py
@@ -76,7 +76,7 @@ def _is_refresh_needed(self) -> bool:
         )
         return time.monotonic() > refresh_threshold
 
-    def is_feature_enabled(self, name: str, default_value: bool) -> bool:
+    def get_flag_value(self, name: str, default_value: Any) -> Any:
         """
         Checks if a feature is enabled.
         - BLOCKS on the first call until flags are fetched.
@@ -95,10 +95,7 @@ def is_feature_enabled(self, name: str, default_value: bool) -> bool:
             assert self._flags is not None
 
             # Now, return the value from the populated cache.
-            flag_value = self._flags.get(name)
-            if flag_value is None:
-                return default_value
-            return flag_value.lower() == "true"
+            return self._flags.get(name, default_value)
 
     def _refresh_flags(self):
         """Performs a synchronous network request to fetch and update flags."""
diff --git a/src/databricks/sql/telemetry/telemetry_client.py b/src/databricks/sql/telemetry/telemetry_client.py
index 2ac22a9d7..55f06c8df 100644
--- a/src/databricks/sql/telemetry/telemetry_client.py
+++ b/src/databricks/sql/telemetry/telemetry_client.py
@@ -110,9 +110,10 @@ def is_telemetry_enabled(connection: "Connection") -> bool:
 
         if connection.enable_telemetry:
             context = FeatureFlagsContextFactory.get_instance(connection)
-            return context.is_feature_enabled(
+            flag_value = context.get_flag_value(
                 TelemetryHelper.TELEMETRY_FEATURE_FLAG_NAME, default_value=False
             )
+            return str(flag_value).lower() == "true"
         else:
             return False
 

From a35d0de80eac5ea00c9d751eca572c83a87052c7 Mon Sep 17 00:00:00 2001
From: Sai Shree Pradhan <saishree.pradhan@databricks.com>
Date: Fri, 1 Aug 2025 12:02:29 +0530
Subject: [PATCH 9/9] test fix

Signed-off-by: Sai Shree Pradhan <saishree.pradhan@databricks.com>
---
 tests/e2e/test_concurrent_telemetry.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/tests/e2e/test_concurrent_telemetry.py b/tests/e2e/test_concurrent_telemetry.py
index cb3aee21f..d924f0569 100644
--- a/tests/e2e/test_concurrent_telemetry.py
+++ b/tests/e2e/test_concurrent_telemetry.py
@@ -76,7 +76,7 @@ def execute_query_worker(thread_id):
 
                 time.sleep(random.uniform(0, 0.05))
                 
-                with self.connection(extra_params={"enable_telemetry": True}) as conn:
+                with self.connection(extra_params={"force_enable_telemetry": True}) as conn:
                     # Capture the session ID from the connection before executing the query
                     session_id_hex = conn.get_session_id_hex()
                     with capture_lock: