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