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.constants import SECRETS_HYDRATION_PREFIX
 11from airbyte.secrets.base import SecretManager, SecretSourceEnum, SecretString
 12from airbyte.secrets.config import _get_secret_sources
 13
 14
 15def is_secret_available(
 16    secret_name: str,
 17) -> bool:
 18    """Check if a secret is available in any of the configured secret sources.
 19
 20    This function checks all available secret sources for the given secret name.
 21    If the secret is found in any source, it returns `True`; otherwise, it returns `False`.
 22    """
 23    try:
 24        _ = get_secret(secret_name, allow_prompt=False)
 25    except exc.PyAirbyteSecretNotFoundError:
 26        return False
 27    else:
 28        # If no exception was raised, the secret was found.
 29        return True
 30
 31
 32def get_secret(
 33    secret_name: str,
 34    /,
 35    *,
 36    sources: list[SecretManager | SecretSourceEnum] | None = None,
 37    allow_prompt: bool = True,
 38    **kwargs: dict[str, Any],
 39) -> SecretString:
 40    """Get a secret from the environment.
 41
 42    The optional `sources` argument of enum type `SecretSourceEnum` or list of `SecretSourceEnum`
 43    options. If left blank, all available sources will be checked. If a list of `SecretSourceEnum`
 44    entries is passed, then the sources will be checked using the provided ordering.
 45
 46    If `allow_prompt` is `True` or if SecretSourceEnum.PROMPT is declared in the `source` arg, then
 47    the user will be prompted to enter the secret if it is not found in any of the other sources.
 48    """
 49    if secret_name.startswith(SECRETS_HYDRATION_PREFIX):
 50        # If the secret name starts with the hydration prefix, we assume it's a secret reference.
 51        # We strip the prefix and get the actual secret name.
 52        secret_name = secret_name.removeprefix(SECRETS_HYDRATION_PREFIX).lstrip()
 53
 54    if "source" in kwargs:
 55        warnings.warn(
 56            message="The `source` argument is deprecated. Use the `sources` argument instead.",
 57            category=DeprecationWarning,
 58            stacklevel=2,
 59        )
 60        sources = kwargs.pop("source")  # type: ignore [assignment]
 61
 62    available_sources: dict[str, SecretManager] = {}
 63    for available_source in _get_secret_sources():
 64        # Add available sources to the dict. Order matters.
 65        available_sources[available_source.name] = available_source
 66
 67    if sources is None:
 68        # If ANY is in the list, then we don't need to check any other sources.
 69        # This is the default behavior.
 70        sources = list(available_sources.values())
 71
 72    elif not isinstance(sources, list):
 73        sources = [sources]  # type: ignore [unreachable]  # This is a 'just in case' catch.
 74
 75    # Replace any SecretSourceEnum strings with the matching SecretManager object
 76    for source in list(sources):
 77        if isinstance(source, SecretSourceEnum):
 78            if source not in available_sources:
 79                raise exc.PyAirbyteInputError(
 80                    guidance="Invalid secret source name.",
 81                    input_value=source,
 82                    context={
 83                        "Available Sources": list(available_sources.keys()),
 84                    },
 85                )
 86
 87            sources[sources.index(source)] = available_sources[source]
 88
 89    secret_managers = cast("list[SecretManager]", sources)
 90
 91    if SecretSourceEnum.PROMPT in secret_managers:
 92        prompt_source = secret_managers.pop(
 93            # Mis-typed, but okay here since we have equality logic for the enum comparison:
 94            secret_managers.index(SecretSourceEnum.PROMPT),  # type: ignore [arg-type]
 95        )
 96
 97        if allow_prompt:
 98            # Always check prompt last. Add it to the end of the list.
 99            secret_managers.append(prompt_source)
100
101    for secret_mgr in secret_managers:
102        val = secret_mgr.get_secret(secret_name)
103        if val:
104            return SecretString(val)
105
106    raise exc.PyAirbyteSecretNotFoundError(
107        secret_name=secret_name,
108        sources=[str(s) for s in available_sources],
109    )