Skip to main content

prefect_databricks.credentials

Credential classes used to perform authenticated interactions with Databricks

Classes

DatabricksCredentials

Block used to manage Databricks authentication. Supports two authentication methods:
  1. Personal Access Token (PAT): Provide a token field.
  2. Service Principal (OAuth 2.0): Provide client_id, client_secret, and optionally tenant_id for Azure Databricks.
Attributes:
  • databricks_instance: Databricks instance used in formatting the endpoint URL.
  • token: The token to authenticate with Databricks (for PAT authentication).
  • client_id: The service principal client ID for OAuth authentication. If provided, client_secret must also be provided.
  • client_secret: The service principal client secret for OAuth authentication. If provided, client_id must also be provided.
  • tenant_id: The tenant ID for Azure Databricks. Optional, only needed for Azure Databricks service principal authentication.
  • client_kwargs: Additional keyword arguments to pass to AsyncClient.
Examples: Load stored Databricks credentials using PAT:
from prefect_databricks import DatabricksCredentials

databricks_credentials_block = DatabricksCredentials.load("BLOCK_NAME")
Using service principal authentication:
from prefect_databricks import DatabricksCredentials

credentials = DatabricksCredentials(
    databricks_instance="dbc-abc123-def4.cloud.databricks.com",
    client_id="my-client-id",
    client_secret="my-client-secret",
)
client = credentials.get_client()
Methods:

get_client

get_client(self) -> AsyncClient
Gets a Databricks REST AsyncClient. Returns:
  • A Databricks REST AsyncClient.
Examples: Gets a Databricks REST AsyncClient using PAT authentication.
from prefect import flow
from prefect_databricks import DatabricksCredentials

@flow
def example_get_client_flow():
    databricks_credentials = DatabricksCredentials(
        databricks_instance="dbc-abc123-def4.cloud.databricks.com",
        token="my-token",
    )
    client = databricks_credentials.get_client()
    return client

example_get_client_flow()
Gets a Databricks REST AsyncClient using service principal authentication.
from prefect import flow
from prefect_databricks import DatabricksCredentials

@flow
def example_get_client_flow():
    databricks_credentials = DatabricksCredentials(
        databricks_instance="dbc-abc123-def4.cloud.databricks.com",
        client_id="my-client-id",
        client_secret="my-client-secret",
    )
    client = databricks_credentials.get_client()
    return client

example_get_client_flow()

validate_auth_method

validate_auth_method(cls, values: Dict[str, Any]) -> Dict[str, Any]
Validates that either PAT or service principal authentication is configured, but not both. Valid configurations:
  1. token only (PAT authentication)
  2. client_id + client_secret (service principal authentication)
  3. client_id + client_secret + tenant_id (Azure service principal authentication)