Skip to content

prefect_databricks.credentials

Credential classes used to perform authenticated interactions with Databricks

Classes

DatabricksCredentials

Bases: Block

Block used to manage Databricks authentication.

Attributes:

Name Type Description
databricks_instance str

Databricks instance used in formatting the endpoint URL.
token SecretStr

The token to authenticate with Databricks.
client_kwargs Optional[Dict[str, Any]]

Additional keyword arguments to pass to AsyncClient.

Examples:

Load stored Databricks credentials: 

from prefect_databricks import DatabricksCredentials
databricks_credentials_block = DatabricksCredentials.load("BLOCK_NAME")

Source code in prefect_databricks/credentials.py 
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
class DatabricksCredentials(Block):
    """
    Block used to manage Databricks authentication.

    Attributes:
        databricks_instance:
            Databricks instance used in formatting the endpoint URL.
        token: The token to authenticate with Databricks.
        client_kwargs: Additional keyword arguments to pass to AsyncClient.

    Examples:
        Load stored Databricks credentials:
        ```python
        from prefect_databricks import DatabricksCredentials
        databricks_credentials_block = DatabricksCredentials.load("BLOCK_NAME")
        ```
    """

    _block_type_name = "Databricks Credentials"
    _logo_url = "https://cdn.sanity.io/images/3ugk85nk/production/ff9a2573c23954bedd27b0f420465a55b1a99dfd-250x250.png"  # noqa
    _documentation_url = "https://prefecthq.github.io/prefect-databricks/credentials/#prefect_databricks.credentials.DatabricksCredentials"  # noqa

    databricks_instance: str = Field(
        default=...,
        description="Databricks instance used in formatting the endpoint URL.",
    )
    token: SecretStr = Field(
        default=..., description="The token to authenticate with Databricks."
    )
    client_kwargs: Optional[Dict[str, Any]] = Field(
        default=None, description="Additional keyword arguments to pass to AsyncClient."
    )

    def get_client(self) -> AsyncClient:
        """
        Gets an Databricks REST AsyncClient.

        Returns:
            An Databricks REST AsyncClient.

        Example:
            Gets a Databricks REST AsyncClient.
            ```python
            from prefect import flow
            from prefect_databricks import DatabricksCredentials

            @flow
            def example_get_client_flow():
                token = "consumer_key"
                databricks_credentials = DatabricksCredentials(token=token)
                client = databricks_credentials.get_client()
                return client

            example_get_client_flow()
            ```
        """
        base_url = f"https://{self.databricks_instance}/api/"

        client_kwargs = self.client_kwargs or {}
        client_kwargs["headers"] = {
            "Authorization": f"Bearer {self.token.get_secret_value()}"
        }
        client = AsyncClient(base_url=base_url, **client_kwargs)
        return client

Functions

get_client

Gets an Databricks REST AsyncClient.

Returns:

Type Description
AsyncClient

An Databricks REST AsyncClient.
Example

Gets a Databricks REST AsyncClient. 

from prefect import flow
from prefect_databricks import DatabricksCredentials

@flow
def example_get_client_flow():
    token = "consumer_key"
    databricks_credentials = DatabricksCredentials(token=token)
    client = databricks_credentials.get_client()
    return client

example_get_client_flow()

Source code in prefect_databricks/credentials.py 
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
def get_client(self) -> AsyncClient:
    """
    Gets an Databricks REST AsyncClient.

    Returns:
        An Databricks REST AsyncClient.

    Example:
        Gets a Databricks REST AsyncClient.
        ```python
        from prefect import flow
        from prefect_databricks import DatabricksCredentials

        @flow
        def example_get_client_flow():
            token = "consumer_key"
            databricks_credentials = DatabricksCredentials(token=token)
            client = databricks_credentials.get_client()
            return client

        example_get_client_flow()
        ```
    """
    base_url = f"https://{self.databricks_instance}/api/"

    client_kwargs = self.client_kwargs or {}
    client_kwargs["headers"] = {
        "Authorization": f"Bearer {self.token.get_secret_value()}"
    }
    client = AsyncClient(base_url=base_url, **client_kwargs)
    return client