force enable telemetry

Signed-off-by: Sai Shree Pradhan <saishree.pradhan@databricks.com>

Author: saishreeeee

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,

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

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):