From ecfd1423af1e3a733b41ad49ae74bd00a1d75991 Mon Sep 17 00:00:00 2001 From: Rohit Sharma Date: Thu, 30 Jul 2026 14:01:27 +0100 Subject: [PATCH] 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 --- .../flytekitplugins/spark/connector.py | 17 ++-- .../flytekitplugins/spark/databricks_auth.py | 83 +++++++++++++++++++ .../tests/test_databricks_auth.py | 66 +++++++++++++++ 3 files changed, 154 insertions(+), 12 deletions(-) create mode 100644 plugins/flytekit-spark/flytekitplugins/spark/databricks_auth.py create mode 100644 plugins/flytekit-spark/tests/test_databricks_auth.py diff --git a/plugins/flytekit-spark/flytekitplugins/spark/connector.py b/plugins/flytekit-spark/flytekitplugins/spark/connector.py index 06b02048d1..7b47f1b3f5 100644 --- a/plugins/flytekit-spark/flytekitplugins/spark/connector.py +++ b/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: diff --git a/plugins/flytekit-spark/flytekitplugins/spark/databricks_auth.py b/plugins/flytekit-spark/flytekitplugins/spark/databricks_auth.py new file mode 100644 index 0000000000..bda5f2a19d --- /dev/null +++ b/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)) diff --git a/plugins/flytekit-spark/tests/test_databricks_auth.py b/plugins/flytekit-spark/tests/test_databricks_auth.py new file mode 100644 index 0000000000..f1cf1bd9a9 --- /dev/null +++ b/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", + )