AirbyteLib: Add basic secrets management (#34822)

airbyte-lib/README.md

@@ -4,17 +4,44 @@ airbyte-lib is a library that allows to run Airbyte syncs embedded into any Pyth
 
 ## Development
 
-* Make sure [Poetry is installed](https://python-poetry.org/docs/#).
-* Run `poetry install`
-* For examples, check out the `examples` folder. They can be run via `poetry run python examples/<example file>`
-* Unit tests and type checks can be run via `poetry run pytest`
+- Make sure [Poetry is installed](https://python-poetry.org/docs/#).
+- Run `poetry install`
+- For examples, check out the `examples` folder. They can be run via `poetry run python examples/<example file>`
+- Unit tests and type checks can be run via `poetry run pytest`
 
 ## Release
 
-* In your PR:
-   * Bump the version in `pyproject.toml`
-   * Add a changelog entry to the table below
-* Once the PR is merged, go to Github and trigger the `Publish AirbyteLib Manually` workflow. This will publish the new version to PyPI.
+- In your PR:
+  - Bump the version in `pyproject.toml`
+  - Add a changelog entry to the table below
+- Once the PR is merged, go to Github and trigger the `Publish AirbyteLib Manually` workflow. This will publish the new version to PyPI.
+
+## Secrets Management
+
+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).
+
+_Note: Additional secret store options may be supported in the future. [More info here.](https://github.com/airbytehq/airbyte-lib-private-beta/discussions/5)_
+
+### Retrieving Secrets
+
+```python
+from airbyte_lib import get_secret, SecretSource
+
+source = get_connection("source-github")
+source.set_config(
+   "credentials": {
+      "personal_access_token": get_secret("GITHUB_PERSONAL_ACCESS_TOKEN"),
+   }
+)
+```
+
+The `get_secret()` function accepts an optional `source` argument of enum type `SecretSource`. If omitted or set to `SecretSource.ANY`, AirbyteLib will search all available secrets sources. If `source` is set to a specific source, then only that source will be checked. If a list of `SecretSource` entries is passed, then the sources will be checked using the provided ordering.
+
+By default, AirbyteLib will prompt the user for any requested secrets that are not provided via other secret managers. You can disable this prompt by passing `prompt=False` to `get_secret()`.
 
 ### Versioning
 
@@ -24,13 +51,13 @@ Versioning follows [Semantic Versioning](https://semver.org/). For new features,
 
 Regular documentation lives in the `/docs` folder. Based on the doc strings of public methods, we generate API documentation using [pdoc](https://pdoc.dev). To generate the documentation, run `poetry run generate-docs`. The documentation will be generated in the `docs/generate` folder. This needs to be done manually when changing the public interface of the library.
 
-A unit test validates the documentation is up to date. 
+A unit test validates the documentation is up to date.
 
 ## Validating source connectors
 
 To validate a source connector for compliance, the `airbyte-lib-validate-source` script can be used. It can be used like this:
 
-```
+```bash
 airbyte-lib-validate-source —connector-dir . -—sample-config secrets/config.json
 ```
 

airbyte-lib/airbyte_lib/__init__.py

@@ -6,6 +6,7 @@
 from airbyte_lib.caches import DuckDBCache, DuckDBCacheConfig
 from airbyte_lib.datasets import CachedDataset
 from airbyte_lib.results import ReadResult
+from airbyte_lib.secrets import SecretSource, get_secret
 from airbyte_lib.source import Source
 
 
@@ -15,7 +16,9 @@
     "DuckDBCacheConfig",
     "get_connector",
     "get_default_cache",
+    "get_secret",
     "new_local_cache",
     "ReadResult",
+    "SecretSource",
     "Source",
 ]

airbyte-lib/airbyte_lib/exceptions.py

@@ -250,3 +250,16 @@ class AirbyteStreamNotFoundError(AirbyteConnectorError):
 
     stream_name: str | None = None
     available_streams: list[str] | None = None
+
+
+@dataclass
+class AirbyteLibSecretNotFoundError(AirbyteError):
+    """Secret not found."""
+
+    guidance = "Please ensure that the secret is set."
+    help_url = (
+        "https://docs.airbyte.com/using-airbyte/airbyte-lib/getting-started#secrets-management"
+    )
+
+    secret_name: str | None = None
+    sources: list[str] | None = None

airbyte-lib/airbyte_lib/secrets.py

@@ -0,0 +1,86 @@
+# Copyright (c) 2023 Airbyte, Inc., all rights reserved.
+"""Secrets management for AirbyteLib."""
+from __future__ import annotations
+
+import os
+from enum import Enum, auto
+from getpass import getpass
+
+from airbyte_lib import exceptions as exc
+
+
+class SecretSource(Enum):
+    ENV = auto()
+    GOOGLE_COLAB = auto()
+    ANY = auto()
+
+    PROMPT = auto()
+
+
+ALL_SOURCES = [
+    SecretSource.ENV,
+    SecretSource.GOOGLE_COLAB,
+]
+
+try:
+    from google.colab import userdata as colab_userdata
+except ImportError:
+    colab_userdata = None
+
+
+def get_secret(
+    secret_name: str,
+    source: SecretSource | list[SecretSource] = SecretSource.ANY,
+    *,
+    prompt: bool = True,
+) -> str:
+    """Get a secret from the environment.
+
+    The optional `source` argument of enum type `SecretSource` or list of `SecretSource` options.
+    If left blank, the `source` arg will be `SecretSource.ANY`. If `source` is set to a specific
+    source, then only that source will be checked. If a list of `SecretSource` entries is passed,
+    then the sources will be checked using the provided ordering.
+
+    If `prompt` to `True` or if SecretSource.PROMPT is declared in the `source` arg, then the
+    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
+    if SecretSource.ANY 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:
+        if SecretSource.PROMPT in sources:
+            sources.remove(SecretSource.PROMPT)
+
+        sources.append(SecretSource.PROMPT)  # Always check prompt last
+
+    for s in sources:
+        val = _get_secret_from_source(secret_name, s)
+        if val:
+            return val
+
+    raise exc.AirbyteLibSecretNotFoundError(
+        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, None)
+    ):
+        return colab_userdata.get(secret_name)
+
+    if source == SecretSource.PROMPT:
+        return getpass(f"Enter the value for secret '{secret_name}': ")
+
+    return None

airbyte-lib/examples/run_faker.py

@@ -26,4 +26,4 @@
 result = source.read()
 
 for name, records in result.streams.items():
-    print(f"Stream {name}: {len(list(records))} records")
+    print(f"Stream {name}: {len(records)} records")

airbyte-lib/examples/run_github.py

@@ -0,0 +1,28 @@
+# Copyright (c) 2023 Airbyte, Inc., all rights reserved.
+"""A simple test of AirbyteLib, using the Faker source connector.
+
+Usage (from airbyte-lib root directory):
+> poetry run python ./examples/run_faker.py
+
+No setup is needed, but you may need to delete the .venv-source-faker folder
+if your installation gets interrupted or corrupted.
+"""
+from __future__ import annotations
+
+import airbyte_lib as ab
+
+
+GITHUB_TOKEN = ab.get_secret("GITHUB_PERSONAL_ACCESS_TOKEN")
+
+
+source = ab.get_connector("source-github")
+source.set_config(
+    {"repositories": ["airbytehq/airbyte"], "credentials": {"personal_access_token": GITHUB_TOKEN}}
+)
+source.check()
+source.set_streams(["products", "users", "purchases"])
+
+result = source.read()
+
+for name, records in result.streams.items():
+    print(f"Stream {name}: {len(records)} records")