airbyte.secrets.util

 1# Copyright (c) 2024 Airbyte, Inc., all rights reserved.
 2"""Helper functions for working with secrets."""
 3
 4from __future__ import annotations
 5
 6import warnings
 7from typing import Any, cast
 8
 9from airbyte import exceptions as exc
10from airbyte.secrets.base import SecretManager, SecretSourceEnum, SecretString
11from airbyte.secrets.config import _get_secret_sources
12
13
14def get_secret(
15    secret_name: str,
16    /,
17    *,
18    sources: list[SecretManager | SecretSourceEnum] | None = None,
19    allow_prompt: bool = True,
20    **kwargs: dict[str, Any],
21) -> SecretString:
22    """Get a secret from the environment.
23
24    The optional `sources` argument of enum type `SecretSourceEnum` or list of `SecretSourceEnum`
25    options. If left blank, all available sources will be checked. If a list of `SecretSourceEnum`
26    entries is passed, then the sources will be checked using the provided ordering.
27
28    If `allow_prompt` is `True` or if SecretSourceEnum.PROMPT is declared in the `source` arg, then
29    the user will be prompted to enter the secret if it is not found in any of the other sources.
30    """
31    if "source" in kwargs:
32        warnings.warn(
33            message="The `source` argument is deprecated. Use the `sources` argument instead.",
34            category=DeprecationWarning,
35            stacklevel=2,
36        )
37        sources = kwargs.pop("source")  # type: ignore [assignment]
38
39    available_sources: dict[str, SecretManager] = {}
40    for available_source in _get_secret_sources():
41        # Add available sources to the dict. Order matters.
42        available_sources[available_source.name] = available_source
43
44    if sources is None:
45        # If ANY is in the list, then we don't need to check any other sources.
46        # This is the default behavior.
47        sources = list(available_sources.values())
48
49    elif not isinstance(sources, list):
50        sources = [sources]  # type: ignore [unreachable]  # This is a 'just in case' catch.
51
52    # Replace any SecretSourceEnum strings with the matching SecretManager object
53    for source in list(sources):
54        if isinstance(source, SecretSourceEnum):
55            if source not in available_sources:
56                raise exc.PyAirbyteInputError(
57                    guidance="Invalid secret source name.",
58                    input_value=source,
59                    context={
60                        "Available Sources": list(available_sources.keys()),
61                    },
62                )
63
64            sources[sources.index(source)] = available_sources[source]
65
66    secret_managers = cast("list[SecretManager]", sources)
67
68    if SecretSourceEnum.PROMPT in secret_managers:
69        prompt_source = secret_managers.pop(
70            # Mis-typed, but okay here since we have equality logic for the enum comparison:
71            secret_managers.index(SecretSourceEnum.PROMPT),  # type: ignore [arg-type]
72        )
73
74        if allow_prompt:
75            # Always check prompt last. Add it to the end of the list.
76            secret_managers.append(prompt_source)
77
78    for secret_mgr in secret_managers:
79        val = secret_mgr.get_secret(secret_name)
80        if val:
81            return SecretString(val)
82
83    raise exc.PyAirbyteSecretNotFoundError(
84        secret_name=secret_name,
85        sources=[str(s) for s in available_sources],
86    )