prefect_databricks.credentials

Credential classes used to perform authenticated interactions with Databricks

Classes

DatabricksCredentials ^{<a href="https://github.com/PrefectHQ/prefect/blob/main/src/integrations/prefect-databricks/prefect_databricks/credentials.py#L15" target="_blank"><Icon icon="github" style="width: 14px; height: 14px;" /></a>}

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 ^{<a href="https://github.com/PrefectHQ/prefect/blob/main/src/integrations/prefect-databricks/prefect_databricks/credentials.py#L219" target="_blank"><Icon icon="github" style="width: 14px; height: 14px;" /></a>}

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 ^{<a href="https://github.com/PrefectHQ/prefect/blob/main/src/integrations/prefect-databricks/prefect_databricks/credentials.py#L103" target="_blank"><Icon icon="github" style="width: 14px; height: 14px;" /></a>}

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)