refactor(spark): isolate Databricks PAT authentication

Introduce an authentication strategy boundary while preserving the existing PAT resolution and connector behavior, enabling subsequent auth methods to be reviewed independently.

Signed-off-by: Rohit Sharma <rohitrsh@gmail.com>
Co-authored-by: Cursor <cursoragent@cursor.com>

Author: rohitrsh

plugins/flytekit-spark/flytekitplugins/spark/connector.py

@@ -262,6 +262,8 @@ async def create(
         task_execution_metadata: Optional[TaskExecutionMetadata] = None,
         **kwargs,
     ) -> DatabricksJobMetadata:
+        from .databricks_auth import select_auth
+
         data = json.dumps(_get_databricks_job_spec(task_template))
         databricks_instance = task_template.custom.get(
             "databricksInstance", os.getenv(DEFAULT_DATABRICKS_INSTANCE_ENV_KEY)
@@ -272,22 +274,13 @@ async def create(
                 f"Missing databricks instance. Please set the value through the task config or set the {DEFAULT_DATABRICKS_INSTANCE_ENV_KEY} environment variable in the connector."
             )
 
-        # Get workflow-specific token or fall back to default
         namespace = task_execution_metadata.namespace if task_execution_metadata else None
-
-        # Extract custom secret name from task template (if provided)
-        custom_secret_name = task_template.custom.get("databricksTokenSecret")
-
-        logger.info(f"Creating Databricks job for namespace: {namespace or 'unknown'}")
-        if custom_secret_name:
-            logger.info(f"Using custom secret name: {custom_secret_name}")
-
-        auth_token = get_databricks_token(
-            namespace=namespace, task_template=task_template, secret_name=custom_secret_name
-        )
+        auth = await select_auth(task_template=task_template, namespace=namespace)
+        logger.info("Databricks auth resolved: %s", auth.describe())
         databricks_url = f"https://{databricks_instance}{DATABRICKS_API_ENDPOINT}/runs/submit"
 
         async with aiohttp.ClientSession() as session:
+            auth_token = await auth.get_bearer_token(session)
             async with session.post(databricks_url, headers=get_header(auth_token=auth_token), data=data) as resp:
                 response = await resp.json()
                 if resp.status != http.HTTPStatus.OK:

plugins/flytekit-spark/flytekitplugins/spark/databricks_auth.py

@@ -0,0 +1,83 @@
+"""Authentication strategies for the Databricks connector.
+
+This module introduces a small strategy boundary around the connector's
+existing Personal Access Token (PAT) flow. It intentionally preserves the
+current token resolution behavior so additional authentication methods can be
+added independently.
+"""
+
+from abc import ABC, abstractmethod
+from dataclasses import dataclass
+from typing import Optional
+
+from flytekit import lazy_module
+from flytekit.models.task import TaskTemplate
+
+aiohttp = lazy_module("aiohttp")
+
+
+@dataclass
+class _Settings:
+    """Authentication settings resolved for one Databricks task."""
+
+    task_template: Optional[TaskTemplate]
+    token_secret_name: Optional[str]
+    namespace: Optional[str]
+
+    @staticmethod
+    def from_task(task_template: Optional[TaskTemplate], namespace: Optional[str]) -> "_Settings":
+        custom = task_template.custom if task_template is not None else {}
+        return _Settings(
+            task_template=task_template,
+            token_secret_name=custom.get("databricksTokenSecret"),
+            namespace=namespace,
+        )
+
+
+class DatabricksAuth(ABC):
+    """Interface for obtaining a bearer token for Databricks API calls."""
+
+    auth_type = "unknown"
+    strategy_name = "DatabricksAuth"
+
+    def __init__(self, settings: _Settings):
+        self.settings = settings
+
+    @abstractmethod
+    async def get_bearer_token(self, session: "aiohttp.ClientSession") -> str:  # type: ignore[name-defined]
+        """Return a bearer token for a Databricks API request."""
+
+    def describe(self) -> str:
+        """Return a description that is safe to write to connector logs."""
+        return (
+            f"strategy={self.strategy_name} auth_type={self.auth_type} " f"namespace={self.settings.namespace or 'N/A'}"
+        )
+
+
+class PATAuth(DatabricksAuth):
+    """Delegate to the connector's existing multi-tenant PAT lookup."""
+
+    auth_type = "pat"
+    strategy_name = "PATAuth"
+
+    async def get_bearer_token(self, session: "aiohttp.ClientSession") -> str:  # type: ignore[name-defined]
+        from .connector import get_databricks_token
+
+        return get_databricks_token(
+            namespace=self.settings.namespace,
+            task_template=self.settings.task_template,
+            secret_name=self.settings.token_secret_name,
+        )
+
+
+async def select_auth(
+    task_template: Optional[TaskTemplate],
+    namespace: Optional[str],
+) -> DatabricksAuth:
+    """Select authentication for a task.
+
+    PAT remains the only strategy and therefore the unconditional default in
+    this refactor. Later authentication methods can extend this dispatcher
+    without changing the connector request lifecycle.
+    """
+    return PATAuth(_Settings.from_task(task_template, namespace))

plugins/flytekit-spark/tests/test_databricks_auth.py

@@ -0,0 +1,66 @@
+"""Tests for the Databricks authentication strategy boundary."""
+
+from unittest.mock import AsyncMock, MagicMock, patch
+
+import pytest
+
+from flytekitplugins.spark.databricks_auth import (
+    PATAuth,
+    _Settings,
+    select_auth,
+)
+
+
+def _task_template(**custom):
+    task_template = MagicMock()
+    task_template.custom = custom
+    return task_template
+
+
+def test_settings_use_default_secret_name():
+    settings = _Settings.from_task(task_template=None, namespace="project-a")
+
+    assert settings.token_secret_name is None
+    assert settings.namespace == "project-a"
+
+
+def test_settings_use_task_secret_name():
+    task_template = _task_template()
+    task_template.custom["databricksTokenSecret"] = "custom-token"
+
+    settings = _Settings.from_task(task_template=task_template, namespace="project-a")
+
+    assert settings.token_secret_name == "custom-token"
+
+
+@pytest.mark.asyncio
+async def test_select_auth_returns_pat_by_default():
+    task_template = _task_template()
+    auth = await select_auth(task_template=task_template, namespace="project-a")
+
+    assert isinstance(auth, PATAuth)
+    assert auth.auth_type == "pat"
+
+
+@pytest.mark.asyncio
+async def test_pat_auth_delegates_to_existing_token_lookup():
+    task_template = _task_template()
+    settings = _Settings(
+        task_template=task_template,
+        token_secret_name="custom-token",
+        namespace="project-a",
+    )
+    auth = PATAuth(settings)
+
+    with patch(
+        "flytekitplugins.spark.connector.get_databricks_token",
+        return_value="example-token",
+    ) as get_token:
+        token = await auth.get_bearer_token(AsyncMock())
+
+    assert token == "example-token"
+    get_token.assert_called_once_with(
+        namespace="project-a",
+        task_template=task_template,
+        secret_name="custom-token",
+    )