fix(feature-flags): decouple cache setup from telemetry

Signed-off-by: Cathleen Yan <58714163+cathleeny@users.noreply.github.com>

Author: cathleeny

src/databricks/sql/session.py

@@ -13,6 +13,10 @@
 from databricks.sql.backend.types import SessionId, BackendType
 from databricks.sql.common.unified_http_client import UnifiedHttpClient
 from databricks.sql.common.agent import detect as detect_agent
+from databricks.sql.common.feature_flag import (
+    FeatureFlagsContext,
+    FeatureFlagsContextFactory,
+)
 from databricks.sql.telemetry.telemetry_client import TelemetryClientFactory
 
 if TYPE_CHECKING:
@@ -134,6 +138,7 @@ def __init__(
         # provider when an ``access_token`` is present, and ``None``
         # otherwise (OAuth M2M/U2M resolve purely from the raw kwargs
         # the bridge reads). The Thrift / SEA backends are unchanged.
+        self.feature_flags: Optional[FeatureFlagsContext] = None
         if kwargs.get("use_kernel", False):
             access_token = kwargs.get("access_token")
             self.auth_provider = (
@@ -143,6 +148,13 @@ def __init__(
             self.auth_provider = get_python_sql_connector_auth_provider(
                 server_hostname, http_client=self.http_client, **kwargs
             )
+            self.feature_flags = FeatureFlagsContextFactory.get_instance(
+                self.host,
+                self.http_client,
+                self.auth_provider,
+                self.useragent_header,
+                self.get_spog_headers(),
+            )
 
         self.backend = self._create_backend(
             server_hostname,

src/databricks/sql/telemetry/telemetry_client.py

@@ -40,7 +40,6 @@
 import uuid
 import locale
 from databricks.sql.telemetry.utils import BaseTelemetryClient
-from databricks.sql.common.feature_flag import FeatureFlagsContextFactory
 from databricks.sql.common.unified_http_client import UnifiedHttpClient
 from databricks.sql.common.http import HttpMethod
 from databricks.sql.exc import RequestError
@@ -133,15 +132,9 @@ def is_telemetry_enabled(connection: "Connection") -> bool:
         if not connection.enable_telemetry:
             return False
 
-        # Only fetch feature flags when enable_telemetry=True and not forced
-        session = connection.session
-        context = FeatureFlagsContextFactory.get_instance(
-            session.host,
-            session.http_client,
-            session.auth_provider,
-            session.useragent_header,
-            session.get_spog_headers(),
-        )
+        context = connection.session.feature_flags
+        if context is None:
+            return False
         return context.get_bool(
             TelemetryHelper.TELEMETRY_FEATURE_FLAG_NAME, default_value=False
         )

tests/unit/conftest.py

@@ -0,0 +1,12 @@
+from unittest.mock import patch
+
+import pytest
+
+
+@pytest.fixture(autouse=True)
+def session_feature_flags():
+    # Unit connections must not contact connector-service. Cache/lifecycle tests
+    # exercise the real reader with a mocked HTTP client.
+    with patch("databricks.sql.session.FeatureFlagsContextFactory") as factory:
+        factory.get_instance.return_value.get_bool.return_value = False
+        yield factory

tests/unit/test_feature_flag_lifecycle.py

@@ -0,0 +1,97 @@
+from unittest.mock import MagicMock, patch
+
+import pytest
+
+from databricks import sql
+from databricks.sql.backend.types import BackendType, SessionId
+from databricks.sql.common.feature_flag import FeatureFlagsContextFactory
+from databricks.sql.telemetry.telemetry_client import TelemetryHelper
+
+
+@pytest.mark.parametrize("use_sea", [False, True])
+@pytest.mark.parametrize("oauth", [False, True])
+def test_flags_available_with_telemetry_disabled(session_feature_flags, use_sea, oauth):
+    session_feature_flags.get_instance.side_effect = (
+        FeatureFlagsContextFactory.get_instance
+    )
+    http = MagicMock()
+    http.request.return_value = MagicMock(
+        status=200, data=b'{"flags":[{"name":"sampleLimit","value":"42"}]}'
+    )
+    auth = MagicMock()
+    auth.add_headers.side_effect = lambda headers: headers.update(
+        Authorization="Bearer test-token"
+    )
+    backend = MagicMock()
+    backend.open_session.return_value = SessionId(BackendType.THRIFT, b"1", b"2")
+
+    def create_backend(**kwargs):
+        http.request.assert_not_called()
+        assert kwargs["auth_provider"] is auth
+        return backend
+
+    backend_name = "SeaDatabricksClient" if use_sea else "ThriftDatabricksClient"
+    try:
+        with patch("databricks.sql.client.UnifiedHttpClient", return_value=http), patch(
+            "databricks.sql.session.get_python_sql_connector_auth_provider",
+            return_value=auth,
+        ) as auth_factory, patch(
+            f"databricks.sql.session.{backend_name}", side_effect=create_backend
+        ):
+            conn = sql.connect(
+                "flags.example",
+                "/sql/1.0/warehouses/test?o=123",
+                enable_telemetry=False,
+                use_sea=use_sea,
+                **(
+                    {"auth_type": "databricks-oauth"}
+                    if oauth
+                    else {"access_token": "test-token"}
+                ),
+            )
+            try:
+                assert conn.telemetry_enabled is False
+                http.request.assert_not_called()
+                assert conn.session.feature_flags.get_int32("sampleLimit") == 42
+                assert (
+                    http.request.call_args.kwargs["headers"]["Authorization"]
+                    == "Bearer test-token"
+                )
+                assert (
+                    http.request.call_args.kwargs["headers"]["x-databricks-org-id"]
+                    == "123"
+                )
+                conn.enable_telemetry = True
+                assert TelemetryHelper.is_telemetry_enabled(conn) is False
+                http.request.assert_called_once()
+                auth_factory.assert_called_once()
+            finally:
+                conn.close()
+    finally:
+        FeatureFlagsContextFactory.remove_instance(
+            "flags.example", {"x-databricks-org-id": "123"}
+        )
+
+
+def test_failed_flag_fetch_does_not_block_session(session_feature_flags):
+    session_feature_flags.get_instance.side_effect = (
+        FeatureFlagsContextFactory.get_instance
+    )
+    http = MagicMock()
+    http.request.side_effect = OSError("connector-service unavailable")
+    try:
+        with patch("databricks.sql.client.UnifiedHttpClient", return_value=http), patch(
+            "databricks.sql.session.ThriftDatabricksClient"
+        ) as backend:
+            backend.return_value.open_session.return_value = SessionId(
+                BackendType.THRIFT, b"1", b"2"
+            )
+            with sql.connect(
+                "flags.example", "/test", access_token="test", enable_telemetry=False
+            ) as conn:
+                backend.return_value.open_session.assert_called_once()
+                http.request.assert_not_called()
+                assert conn.session.feature_flags.get_bool("missing") is False
+                http.request.assert_called_once()
+    finally:
+        FeatureFlagsContextFactory.remove_instance("flags.example")

tests/unit/test_session.py

@@ -398,7 +398,9 @@ def _build_session(self, **extra):
             )
             return sess, mock_get_provider
 
-    def test_use_kernel_m2m_does_not_build_connector_provider(self):
+    def test_use_kernel_m2m_does_not_build_connector_provider(
+        self, session_feature_flags
+    ):
         sess, mock_get_provider = self._build_session(
             oauth_client_id="sp-uuid", oauth_client_secret="shh"
         )
@@ -408,15 +410,21 @@ def test_use_kernel_m2m_does_not_build_connector_provider(self):
         # ...and with no access_token, auth_provider is None (M2M
         # resolves in-kernel from the raw kwargs).
         assert sess.auth_provider is None
+        assert sess.feature_flags is None
+        session_feature_flags.get_instance.assert_not_called()
 
-    def test_use_kernel_pat_builds_minimal_access_token_provider(self):
+    def test_use_kernel_pat_builds_minimal_access_token_provider(
+        self, session_feature_flags
+    ):
         from databricks.sql.auth.authenticators import AccessTokenAuthProvider
 
         sess, mock_get_provider = self._build_session(access_token="dapi-xyz")
         mock_get_provider.assert_not_called()
         # PAT path: a minimal AccessTokenAuthProvider, not the
         # federation-wrapped connector provider.
         assert isinstance(sess.auth_provider, AccessTokenAuthProvider)
+        assert sess.feature_flags is None
+        session_feature_flags.get_instance.assert_not_called()
 
 
 class TestKernelTransportOptionsThreading:

tests/unit/test_telemetry.py

@@ -542,6 +542,11 @@ def teardown_method(self):
         TelemetryClientFactory._clients.clear()
         FeatureFlagsContextFactory._context_map.clear()
 
+    def _set_session_flags(self, session):
+        session.feature_flags = FeatureFlagsContextFactory.get_instance(
+            session.host, session.http_client, session.auth_provider, "test-agent"
+        )
+
     def _mock_ff_response(self, mock_http_request, enabled: bool):
         """Helper method to mock feature flag response for unified HTTP client."""
         mock_response = MagicMock()
@@ -577,6 +582,7 @@ def test_telemetry_enabled_when_flag_is_true(self, mock_http_request, MockSessio
         mock_http_client = MagicMock()
         mock_http_client.request = mock_http_request
         mock_session_instance.http_client = mock_http_client
+        self._set_session_flags(mock_session_instance)
 
         conn = sql.client.Connection(
             server_hostname="test",
@@ -609,6 +615,7 @@ def test_telemetry_disabled_when_flag_is_false(
         mock_http_client = MagicMock()
         mock_http_client.request = mock_http_request
         mock_session_instance.http_client = mock_http_client
+        self._set_session_flags(mock_session_instance)
 
         conn = sql.client.Connection(
             server_hostname="test",
@@ -641,6 +648,7 @@ def test_telemetry_disabled_when_flag_request_fails(
         mock_http_client = MagicMock()
         mock_http_client.request = mock_http_request
         mock_session_instance.http_client = mock_http_client
+        self._set_session_flags(mock_session_instance)
 
         conn = sql.client.Connection(
             server_hostname="test",