From 0aa24041af9447b4ba2b14300c17037b8f0a46a8 Mon Sep 17 00:00:00 2001 From: "Aaron (\"AJ\") Steers" Date: Tue, 13 Feb 2024 14:50:37 -0800 Subject: [PATCH] AirbyteLib: support secrets in dotenv files (#35244) --- airbyte-lib/README.md | 5 +- airbyte-lib/airbyte_lib/secrets.py | 104 ++++++++++++++------ airbyte-lib/docs/generated/airbyte_lib.html | 20 +++- airbyte-lib/poetry.lock | 16 ++- airbyte-lib/pyproject.toml | 1 + 5 files changed, 108 insertions(+), 38 deletions(-) diff --git a/airbyte-lib/README.md b/airbyte-lib/README.md index 4becb415477175..868a5483dc6344 100644 --- a/airbyte-lib/README.md +++ b/airbyte-lib/README.md @@ -21,8 +21,9 @@ airbyte-lib is a library that allows to run Airbyte syncs embedded into any Pyth AirbyteLib can auto-import secrets from the following sources: 1. Environment variables. -2. [Google Colab secrets](https://medium.com/@parthdasawant/how-to-use-secrets-in-google-colab-450c38e3ec75). -3. Manual entry via [`getpass`](https://docs.python.org/3.9/library/getpass.html). +2. Variables defined in a local `.env` ("Dotenv") file. +3. [Google Colab secrets](https://medium.com/@parthdasawant/how-to-use-secrets-in-google-colab-450c38e3ec75). +4. Manual entry via [`getpass`](https://docs.python.org/3.9/library/getpass.html). _Note: Additional secret store options may be supported in the future. [More info here.](https://github.com/airbytehq/airbyte-lib-private-beta/discussions/5)_ diff --git a/airbyte-lib/airbyte_lib/secrets.py b/airbyte-lib/airbyte_lib/secrets.py index 2156eab7e1c2fa..6aea9f163d2fcd 100644 --- a/airbyte-lib/airbyte_lib/secrets.py +++ b/airbyte-lib/airbyte_lib/secrets.py @@ -2,30 +2,90 @@ """Secrets management for AirbyteLib.""" from __future__ import annotations +import contextlib import os from enum import Enum, auto from getpass import getpass +from typing import TYPE_CHECKING + +from dotenv import dotenv_values from airbyte_lib import exceptions as exc +if TYPE_CHECKING: + from collections.abc import Callable + + +try: + from google.colab import userdata as colab_userdata +except ImportError: + colab_userdata = None + + class SecretSource(Enum): ENV = auto() + DOTENV = auto() GOOGLE_COLAB = auto() ANY = auto() PROMPT = auto() -ALL_SOURCES = [ - SecretSource.ENV, - SecretSource.GOOGLE_COLAB, -] +def _get_secret_from_env( + secret_name: str, +) -> str | None: + if secret_name not in os.environ: + return None -try: - from google.colab import userdata as colab_userdata -except ImportError: - colab_userdata = None + return os.environ[secret_name] + + +def _get_secret_from_dotenv( + secret_name: str, +) -> str | None: + try: + dotenv_vars: dict[str, str | None] = dotenv_values() + except Exception: + # Can't locate or parse a .env file + return None + + if secret_name not in dotenv_vars: + # Secret not found + return None + + return dotenv_vars[secret_name] + + +def _get_secret_from_colab( + secret_name: str, +) -> str | None: + if colab_userdata is None: + # The module doesn't exist. We probably aren't in Colab. + return None + + try: + return colab_userdata.get(secret_name) + except Exception: + # Secret name not found. Continue. + return None + + +def _get_secret_from_prompt( + secret_name: str, +) -> str | None: + with contextlib.suppress(Exception): + return getpass(f"Enter the value for secret '{secret_name}': ") + + return None + + +_SOURCE_FUNCTIONS: dict[SecretSource, Callable] = { + SecretSource.ENV: _get_secret_from_env, + SecretSource.DOTENV: _get_secret_from_dotenv, + SecretSource.GOOGLE_COLAB: _get_secret_from_colab, + SecretSource.PROMPT: _get_secret_from_prompt, +} def get_secret( @@ -45,8 +105,9 @@ def get_secret( user will be prompted to enter the secret if it is not found in any of the other sources. """ sources = [source] if not isinstance(source, list) else source + all_sources = set(_SOURCE_FUNCTIONS.keys()) - {SecretSource.PROMPT} if SecretSource.ANY in sources: - sources += [s for s in ALL_SOURCES if s not in sources] + sources += [s for s in all_sources if s not in sources] sources.remove(SecretSource.ANY) if prompt or SecretSource.PROMPT in sources: @@ -55,8 +116,9 @@ def get_secret( sources.append(SecretSource.PROMPT) # Always check prompt last - for s in sources: - val = _get_secret_from_source(secret_name, s) + for source in sources: + fn = _SOURCE_FUNCTIONS[source] # Get the matching function for this source + val = fn(secret_name) if val: return val @@ -64,23 +126,3 @@ def get_secret( secret_name=secret_name, sources=[str(s) for s in sources], ) - - -def _get_secret_from_source( - secret_name: str, - source: SecretSource, -) -> str | None: - if source in [SecretSource.ENV, SecretSource.ANY] and secret_name in os.environ: - return os.environ[secret_name] - - if ( - source in [SecretSource.GOOGLE_COLAB, SecretSource.ANY] - and colab_userdata is not None - and colab_userdata.get(secret_name) - ): - return colab_userdata.get(secret_name) - - if source == SecretSource.PROMPT: - return getpass(f"Enter the value for secret '{secret_name}': ") - - return None diff --git a/airbyte-lib/docs/generated/airbyte_lib.html b/airbyte-lib/docs/generated/airbyte_lib.html index 4b3fb85a419a45..5c7d778615234e 100644 --- a/airbyte-lib/docs/generated/airbyte_lib.html +++ b/airbyte-lib/docs/generated/airbyte_lib.html @@ -319,7 +319,7 @@
Inherited Members
def - get_secret( secret_name: str, source: SecretSource | list[SecretSource] = <SecretSource.ANY: 3>, *, prompt: bool = True) -> str: + get_secret( secret_name: str, source: SecretSource | list[SecretSource] = <SecretSource.ANY: 4>, *, prompt: bool = True) -> str:
@@ -475,11 +475,23 @@
Inherited Members
+ +
+
+ DOTENV = +<SecretSource.DOTENV: 2> + + +
+ + + +
GOOGLE_COLAB = -<SecretSource.GOOGLE_COLAB: 2> +<SecretSource.GOOGLE_COLAB: 3>
@@ -491,7 +503,7 @@
Inherited Members
ANY = -<SecretSource.ANY: 3> +<SecretSource.ANY: 4>
@@ -503,7 +515,7 @@
Inherited Members
PROMPT = -<SecretSource.PROMPT: 4> +<SecretSource.PROMPT: 5>
diff --git a/airbyte-lib/poetry.lock b/airbyte-lib/poetry.lock index 30b32b449c82b4..752f802856ebc7 100644 --- a/airbyte-lib/poetry.lock +++ b/airbyte-lib/poetry.lock @@ -1823,6 +1823,20 @@ files = [ [package.dependencies] six = ">=1.5" +[[package]] +name = "python-dotenv" +version = "1.0.1" +description = "Read key-value pairs from a .env file and set them as environment variables" +optional = false +python-versions = ">=3.8" +files = [ + {file = "python-dotenv-1.0.1.tar.gz", hash = "sha256:e324ee90a023d808f1959c46bcbc04446a10ced277783dc6ee09987c37ec10ca"}, + {file = "python_dotenv-1.0.1-py3-none-any.whl", hash = "sha256:f7b63ef50f1b690dddf550d03497b66d609393b40b564ed0d674909a68ebf16a"}, +] + +[package.extras] +cli = ["click (>=5.0)"] + [[package]] name = "python-ulid" version = "2.2.0" @@ -2591,4 +2605,4 @@ files = [ [metadata] lock-version = "2.0" python-versions = "^3.9" -content-hash = "d94dddeda774321ae70d0d4a354c8805956f532619b79d279eca9cf68a7cae47" +content-hash = "15665328452a67b8dfce18573caeae7856425cca2a3bafc5e9e455a619548314" diff --git a/airbyte-lib/pyproject.toml b/airbyte-lib/pyproject.toml index c5f1ff7eef7f98..d367368bd17a23 100644 --- a/airbyte-lib/pyproject.toml +++ b/airbyte-lib/pyproject.toml @@ -30,6 +30,7 @@ pyarrow = "^14.0.2" # psycopg = {extras = ["binary", "pool"], version = "^3.1.16"} rich = "^13.7.0" pendulum = "<=3.0.0" +python-dotenv = "^1.0.1" [tool.poetry.group.dev.dependencies]