Skip to content

decorative_secrets.databricks

apply_databricks_secrets_arguments 

apply_databricks_secrets_arguments(
    *args: databricks.sdk.config.Config, **kwargs: str
) -> collections.abc.Callable

This decorator maps parameter names to Databricks secrets. Each key in databricks_secret_arguments represents the name of a parameter in the decorated function which accepts an explicit input, and the corresponding mapped value is a parameter name accepting a tuple with the secret scope and key with which to lookup a secret to pass to the mapped parameter in lieu of an explicitly provided argument.

Parameters:

  • *args (databricks.sdk.config.Config, default: () ) –

    A databricks.sdk.Config instance to configure a workspace client when retrieving secrets remotely (if more than one is provided, only the first is used).

  • **kwargs (str, default: {} ) –

    A mapping of static parameter names to the parameter names of arguments accepting Databricks secret scope + key tuples from which to retrieve a value when the key argument is not explicitly provided.

Example 
from functools import (
    cache,
)
from decorative_secrets.databricks import (
    apply_databricks_secret_arguments,
)
from my_client_sdk import (
    Client,
)


@cache
@apply_databricks_secret_arguments(
    client_id="client_id_databricks_secret",
    client_secret="client_secret_databricks_secret",
)
def get_client(
    client_id: str | None = None,
    client_secret: str = None,
    client_id_databricks_secret: str | None = None,
    client_secret_databricks_secret: str | None = None,
) -> Client:
    return Client(
        oauth2_client_id=client_id,
        oauth2_client_secret=client_secret,
    )


client: Client = get_client(
    client_id_databricks_secret=(
        "client",
        "client-id",
    ),
    client_secret_databricks_secret=(
        "client",
        "client-secret",
    ),
)
Source code in src/decorative_secrets/databricks.py 
 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
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
def apply_databricks_secrets_arguments(
    *args: Config,
    **kwargs: str,
) -> Callable:
    """
    This decorator maps parameter names to Databricks secrets.
    Each key in `databricks_secret_arguments` represents the name of a
    parameter in the decorated function which accepts an explicit input, and
    the corresponding mapped value is a parameter name accepting a tuple with
    the secret scope and key with which to lookup a secret to pass to the
    mapped parameter in lieu of an explicitly provided argument.

    Parameters:
        *args: A `databricks.sdk.Config` instance to configure
            a workspace client when retrieving secrets remotely
            (if more than one  is provided, only the first is used).
        **kwargs: A mapping of static parameter names to the parameter names
            of arguments accepting Databricks secret scope + key tuples
            from which to retrieve a value when the key argument is not
            explicitly provided.

    Example:
        ```python
        from functools import (
            cache,
        )
        from decorative_secrets.databricks import (
            apply_databricks_secret_arguments,
        )
        from my_client_sdk import (
            Client,
        )


        @cache
        @apply_databricks_secret_arguments(
            client_id="client_id_databricks_secret",
            client_secret="client_secret_databricks_secret",
        )
        def get_client(
            client_id: str | None = None,
            client_secret: str = None,
            client_id_databricks_secret: str | None = None,
            client_secret_databricks_secret: str | None = None,
        ) -> Client:
            return Client(
                oauth2_client_id=client_id,
                oauth2_client_secret=client_secret,
            )


        client: Client = get_client(
            client_id_databricks_secret=(
                "client",
                "client-id",
            ),
            client_secret_databricks_secret=(
                "client",
                "client-secret",
            ),
        )
        ```
    """
    config: Config | None = _get_args_config(*args)[1]
    get_scope_key_secret: Callable[[str | tuple[str, str]], str] = partial(
        _get_scope_key_secret,
        config=config,
    )
    return apply_callback_arguments(
        get_scope_key_secret,
        **kwargs,
    )

which_databricks 

which_databricks() -> str

Find the databricks executable, or install the Databricks CLI if not found.

Source code in src/decorative_secrets/databricks.py 
161
162
163
164
165
166
167
168
169
170
171
172
def which_databricks() -> str:
    """
    Find the `databricks` executable, or install the Databricks CLI if not
    found.
    """
    databricks: str = which("databricks") or "databricks"
    try:
        check_output((databricks, "--version"))
    except (CalledProcessError, FileNotFoundError):
        _install_databricks_cli()
        databricks = which("databricks") or "databricks"
    return databricks

databricks_auth_login 

databricks_auth_login(host: str | None = None) -> None

Log in to Databricks using the CLI if not already logged in.

Source code in src/decorative_secrets/databricks.py 
188
189
190
191
192
def databricks_auth_login(host: str | None = None) -> None:
    """
    Log in to Databricks using the CLI if not already logged in.
    """
    return _databricks_auth_login(host or os.getenv("DATABRICKS_HOST"))

get_dbutils 

get_dbutils(
    host: str | None = None,
    account_id: str | None = None,
    username: str | None = None,
    password: str | None = None,
    client_id: str | None = None,
    client_secret: str | None = None,
    token: str | None = None,
    profile: str | None = None,
    config_file: str | None = None,
    azure_workspace_resource_id: str | None = None,
    azure_client_secret: str | None = None,
    azure_client_id: str | None = None,
    azure_tenant_id: str | None = None,
    azure_environment: str | None = None,
    auth_type: str | None = None,
    cluster_id: str | None = None,
    google_credentials: str | None = None,
    google_service_account: str | None = None,
    debug_truncate_bytes: int | None = None,
    debug_headers: bool | None = None,
    product: str = "unknown",
    product_version: str = "0.0.0",
    credentials_strategy: (
        databricks.sdk.credentials_provider.CredentialsStrategy
        | None
    ) = None,
    credentials_provider: (
        databricks.sdk.credentials_provider.CredentialsStrategy
        | None
    ) = None,
    token_audience: str | None = None,
    config: databricks.sdk.config.Config | None = None,
) -> databricks.sdk.dbutils.RemoteDbUtils

Get dbutils using an existing instance from the runtime if found, otherwise, creating one using a workspace client (this requires either having these environment variables set , or providing the equivalent optional arguments.

Source code in src/decorative_secrets/databricks.py 
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359
360
361
362
363
364
365
366
367
368
369
370
371
372
373
374
375
376
377
378
379
380
381
382
383
384
385
386
387
388
389
390
391
392
393
394
395
396
397
398
399
400
401
402
403
404
405
406
407
408
409
410
411
412
413
414
415
416
417
418
419
420
421
def get_dbutils(
    host: str | None = None,
    account_id: str | None = None,
    username: str | None = None,
    password: str | None = None,
    client_id: str | None = None,
    client_secret: str | None = None,
    token: str | None = None,
    profile: str | None = None,
    config_file: str | None = None,
    azure_workspace_resource_id: str | None = None,
    azure_client_secret: str | None = None,
    azure_client_id: str | None = None,
    azure_tenant_id: str | None = None,
    azure_environment: str | None = None,
    auth_type: str | None = None,
    cluster_id: str | None = None,
    google_credentials: str | None = None,
    google_service_account: str | None = None,
    debug_truncate_bytes: int | None = None,
    debug_headers: bool | None = None,
    product: str = "unknown",
    product_version: str = "0.0.0",
    credentials_strategy: CredentialsStrategy | None = None,
    credentials_provider: CredentialsStrategy | None = None,
    token_audience: str | None = None,
    config: Config | None = None,
) -> RemoteDbUtils:  # pragma: no cover - environment dependent
    """
    Get [dbutils](https://docs.databricks.com/dev-tools/databricks-utils.html)
    using an existing instance from the runtime if found, otherwise,
    creating one using a workspace client (this requires either having
    [these environment variables set
    ](https://docs.databricks.com/aws/en/dev-tools/auth#environment-variables),
    or providing the equivalent optional arguments.
    """
    dbutils: RemoteDbUtils | None = None
    with suppress(ImportError):
        from IPython.core.getipython import (  # type: ignore[import-not-found]
            get_ipython,
        )

        if TYPE_CHECKING:
            from IPython.core.interactiveshell import (  # type: ignore[import-not-found]
                InteractiveShell,
            )

        ipython: InteractiveShell = get_ipython()
        if ipython is not None:
            user_namespace_attribute: str
            user_namespace: dict
            for user_namespace_attribute in "user_ns", "user_global_ns":
                dbutils = getattr(ipython, user_namespace_attribute, {}).get(
                    "dbutils", None
                )
                if dbutils is not None:
                    return dbutils
    dbutils = globals().get("dbutils")
    if dbutils is not None:
        return dbutils
    databricks_workspace_client: WorkspaceClient = (
        _get_databricks_workspace_client(
            host=host,
            account_id=account_id,
            username=username,
            password=password,
            client_id=client_id,
            client_secret=client_secret,
            token=token,
            profile=profile,
            config_file=config_file,
            azure_workspace_resource_id=azure_workspace_resource_id,
            azure_client_secret=azure_client_secret,
            azure_client_id=azure_client_id,
            azure_tenant_id=azure_tenant_id,
            azure_environment=azure_environment,
            auth_type=auth_type,
            cluster_id=cluster_id,
            google_credentials=google_credentials,
            google_service_account=google_service_account,
            debug_truncate_bytes=debug_truncate_bytes,
            debug_headers=debug_headers,
            product=product,
            product_version=product_version,
            credentials_strategy=credentials_strategy,
            credentials_provider=credentials_provider,
            token_audience=token_audience,
            config=config,
        )
    )
    return databricks_workspace_client.dbutils

get_databricks_secret 

get_databricks_secret(
    scope: str,
    key: str,
    config: databricks.sdk.config.Config | None = None,
) -> str

Get a secret from Databricks.

Source code in src/decorative_secrets/databricks.py 
441
442
443
444
445
446
447
448
449
450
451
452
453
454
def get_databricks_secret(
    scope: str,
    key: str,
    config: Config | None = None,
) -> str:
    """
    Get a secret from Databricks.
    """
    return _get_secret(
        scope,
        key,
        config=config,
        **os.environ,
    )

main 

main() -> None

Run a command: - install: Install the Databricks CLI if not already installed - get: Get a secret from Databricks and print it to stdout

Source code in src/decorative_secrets/databricks.py 
499
500
501
502
503
504
505
506
507
508
509
510
511
512
513
514
515
516
517
518
519
520
521
522
523
524
525
526
527
528
529
530
531
532
533
534
535
536
537
538
539
540
541
542
543
544
545
546
547
548
549
550
551
552
553
554
555
556
557
558
559
560
561
562
563
564
565
566
567
568
569
570
571
572
573
574
575
576
577
578
579
580
581
582
583
584
585
586
def main() -> None:
    """
    Run a command:
    -   install: Install the Databricks CLI if not already installed
    -   get: Get a secret from Databricks and print it to stdout
    """
    command = _get_command()
    if command in ("--help", "-h"):
        _print_help()
        return
    parser: argparse.ArgumentParser
    if command == "install":
        parser = argparse.ArgumentParser(
            prog="decorative-secrets databricks install",
            description="Install the Databricks CLI",
        )
        parser.parse_args()
        _install_databricks_cli()
    elif command == "get":
        parser = argparse.ArgumentParser(
            prog="decorative-secrets databricks get",
            description="Get a secret from Databricks",
        )
        parser.add_argument(
            "scope",
            type=str,
        )
        parser.add_argument(
            "key",
            type=str,
        )
        parser.add_argument(
            "--host",
            default=None,
            type=str,
            help="A Databricks workspace host URL",
        )
        parser.add_argument(
            "-cid",
            "--client-id",
            default=None,
            type=str,
            help="A Databricks OAuth2 Client ID",
        )
        parser.add_argument(
            "-cs",
            "--client-secret",
            default=None,
            type=str,
            help="A Databricks OAuth2 Client Secret",
        )
        parser.add_argument(
            "-t",
            "--token",
            default=None,
            type=str,
            help="A Databricks Personal Access Token",
        )
        parser.add_argument(
            "-p",
            "--profile",
            default=None,
            type=str,
            help="A Databricks Configuration Profile",
        )
        namespace: argparse.Namespace = parser.parse_args()
        config: Config | None = None
        if (
            namespace.host
            or namespace.client_id
            or namespace.client_secret
            or namespace.token
            or namespace.profile
        ):
            config = Config(
                host=namespace.host,
                client_id=namespace.client_id,
                client_secret=namespace.client_secret,
                token=namespace.token,
                profile=namespace.profile,
            )
        print(  # noqa: T201
            get_databricks_secret(
                namespace.scope,
                namespace.key,
                config=config,
            )
        )