Telemetry: report wrapped provider's auth mech/flow under TokenFederation

TokenFederationProvider wraps an underlying auth provider (PAT, OAuth, M2M)
and only adds a token-exchange step. The telemetry helpers previously fell
through to AuthMech.OTHER with no flow, hiding the actual auth method.
Unwrap and recurse on external_provider so federated PAT reports PAT,
federated M2M reports CLIENT_CREDENTIALS, etc.

Co-authored-by: Isaac

Author: samikshya-db

src/databricks/sql/telemetry/telemetry_client.py

@@ -34,6 +34,7 @@
     DatabricksOAuthProvider,
     ExternalAuthProvider,
 )
+from databricks.sql.auth.token_federation import TokenFederationProvider
 import sys
 import platform
 import uuid
@@ -90,6 +91,8 @@ def get_auth_mechanism(auth_provider):
 
         if not auth_provider:
             return None
+        if isinstance(auth_provider, TokenFederationProvider):
+            return TelemetryHelper.get_auth_mechanism(auth_provider.external_provider)
         if isinstance(auth_provider, AccessTokenAuthProvider):
             return AuthMech.PAT
         elif isinstance(auth_provider, DatabricksOAuthProvider):
@@ -105,6 +108,8 @@ def get_auth_flow(auth_provider):
 
         if not auth_provider:
             return None
+        if isinstance(auth_provider, TokenFederationProvider):
+            return TelemetryHelper.get_auth_flow(auth_provider.external_provider)
         if isinstance(auth_provider, DatabricksOAuthProvider):
             if auth_provider._access_token and auth_provider._refresh_token:
                 return AuthFlow.TOKEN_PASSTHROUGH

tests/unit/test_telemetry.py

@@ -29,6 +29,7 @@
     DatabricksOAuthProvider,
     ExternalAuthProvider,
 )
+from databricks.sql.auth.token_federation import TokenFederationProvider
 from databricks import sql
 
 
@@ -211,6 +212,37 @@ def test_auth_flow_detection(self):
         # Test None auth provider
         assert TelemetryHelper.get_auth_flow(None) is None
 
+    def test_token_federation_unwraps_to_inner_provider(self):
+        """Token federation should report the wrapped provider's mech and flow."""
+
+        def make_federation(inner):
+            fed = MagicMock(spec=TokenFederationProvider)
+            fed.external_provider = inner
+            return fed
+
+        # PAT wrapped by federation
+        pat_inner = AccessTokenAuthProvider("test-token")
+        fed_pat = make_federation(pat_inner)
+        assert TelemetryHelper.get_auth_mechanism(fed_pat) == AuthMech.PAT
+        assert TelemetryHelper.get_auth_flow(fed_pat) is None
+
+        # M2M (ExternalAuthProvider) wrapped by federation
+        m2m_inner = MagicMock(spec=ExternalAuthProvider)
+        fed_m2m = make_federation(m2m_inner)
+        assert TelemetryHelper.get_auth_mechanism(fed_m2m) == AuthMech.OTHER
+        assert TelemetryHelper.get_auth_flow(fed_m2m) == AuthFlow.CLIENT_CREDENTIALS
+
+        # OAuth (browser) wrapped by federation
+        oauth_inner = MagicMock(spec=DatabricksOAuthProvider)
+        oauth_inner._access_token = None
+        oauth_inner._refresh_token = None
+        fed_oauth = make_federation(oauth_inner)
+        assert TelemetryHelper.get_auth_mechanism(fed_oauth) == AuthMech.OAUTH
+        assert (
+            TelemetryHelper.get_auth_flow(fed_oauth)
+            == AuthFlow.BROWSER_BASED_AUTHENTICATION
+        )
+
 
 class TestTelemetryFactory:
     """Tests for TelemetryClientFactory lifecycle and management."""