Repository: datasette/datasette-enrichments (language: python)

Issue:
Mechanism for specifying that an Enrichment subclass uses a secret
Split from:
- #45


Discussion / hints:
Here's how `datasette-enrichments-gpt` does that at the moment: https://github.com/datasette/datasette-enrichments-gpt/blob/1353a46f08f73e9b18a58115d74730fdab7427ad/datasette_enrichments_gpt/__init__.py#L32-L116

```python
    async def get_config_form(self, datasette, db, table):
        class ConfigForm(Form):
            # Select box to pick model from gpt-3.5-turbo or gpt-4-turbo
            model = SelectField(
                "Model",
                choices=[
                    ("gpt-3.5-turbo", "gpt-3.5-turbo"),
                    ("gpt-4-turbo", "gpt-4-turbo"),
                    ("gpt-4-vision", "gpt-4-turbo vision"),
                ],
                default="gpt-3.5-turbo",
            )
            # ... more stuff like that

        def stash_api_key(form, field):
            if not (field.data or "").startswith("sk-"):
                raise ValidationError("API key must start with sk-")
            if not hasattr(datasette, "_enrichments_gpt_stashed_keys"):
                datasette._enrichments_gpt_stashed_keys = {}
            key = secrets.token_urlsafe(16)
            datasette._enrichments_gpt_stashed_keys[key] = field.data
            field.data = key

        class ConfigFormWithKey(ConfigForm):
            api_key = PasswordField(
                "API key",
                description="Your OpenAI API key",
                validators=[
                    DataRequired(message="API key is required."),
                    stash_api_key,
                ],
                render_kw={"autocomplete": "off"},
            )

        plugin_config = datasette.plugin_config("datasette-enrichments-gpt") or {}
        api_key = plugin_config.get("api_key")

        return ConfigForm if api_key else ConfigFormWithKey
```

Then it accesses that stashed key here:

https://github.com/datasette/datasette-enrichments-gpt/blob/1353a46f08f73e9b18a58115d74730fdab7427ad/datasette_enrichments_gpt/__init__.py#L229-L244

```python
def resolve_api_key(datasette, config):
    plugin_config = datasette.plugin_config("datasette-enrichments-gpt") or {}
    api_key = plugin_config.get("api_key")
    if api_key:
        return api_key
    # Look for it in config
    api_key_name = config.get("api_key")
    if not api_key_name:
        raise ApiKeyError("No API key reference found in config")
    # Look it up in the stash
    if not hasattr(datasette, "_enrichments_gpt_stashed_keys"):
        raise ApiKeyError("No API key stash found")
    stashed_keys = datasette._enrichments_gpt_stashed_keys
    if api_key_name not in stashed_keys:
        raise ApiKeyError("No API key found in stash for {}".format(api_key_name))
    return stashed_keys[api_key_name]
```
Since I'm encouraging integration with `datasette-secrets` anyway, I'm going to allow `Enrichment` subclasses to indicate that they need an API key. If they do that, the API key form trick will be applied for them.

One potential design:
```python
class GptEnrichment(Enrichment):
    name = "AI analysis with OpenAI GPT"
    slug = "gpt"
    description = "Analyze data using OpenAI's GPT models"
    runs_in_process = True
    batch_size = 1
    needs_api_key = "OPENAI_API_KEY"
    needs_api_key_label = "OpenAI API Key"
    needs_api_key_obtain_url = "https://..."
```
Or... since we'll be using `register_secrets()` from https://github.com/datasette/datasette-secrets for this anyway, maybe we reuse that `datasette_secrets.Secret()` object here? That would be pretty neat.
I'm going to experiment with `secret = Secret(...)` - and then a pattern where `register_secrets()` returns those secrets.
Maybe don't even tell plugins they have to call `register_secrets()` themselves - the `datasette-enrichments` plugin could automatically return ANY `secrets` property of a subclass as part of its own `register_secrets()` hook implementation.
This depends on `datasette-secrets` and that depends on `datasette>=1.0a13` so this can only be in an alpha release for the moment.
... which means the entire enrichments ecosystem is about to go alpha-only, which isn't ideal.
Shipped `datasette-test` 0.3 with these changes:
- https://github.com/datasette/datasette-test/issues/4
- https://github.com/datasette/datasette-test/issues/5
- And 0.3.1 with a fix: https://github.com/datasette/datasette-test/issues/6
OK, `datasette-secrets` is now compatible with <1.0 and >=1.0a Datasette.

Fix patch:
diff --git a/datasette_enrichments/__init__.py b/datasette_enrichments/__init__.py
index a001fac..cdec022 100644
--- a/datasette_enrichments/__init__.py
+++ b/datasette_enrichments/__init__.py
@@ -2,18 +2,22 @@
 import asyncio
 from datasette import hookimpl
 from datasette.utils import async_call_with_supported_arguments, tilde_encode
+from datasette_secrets import Secret, get_secret
 import json
 import secrets
 import traceback
 import urllib
 from datasette.plugins import pm
+from markupsafe import Markup, escape
 from .views import enrichment_picker, enrichment_view
+from wtforms import PasswordField
+from wtforms.validators import DataRequired
 from .utils import get_with_auth, mark_job_complete, pks_for_rows
 from . import hookspecs
 
 from datasette.utils import await_me_maybe
 
-from typing import TYPE_CHECKING, Any, Tuple, Union, List
+from typing import TYPE_CHECKING, Any, List, Optional, Tuple, Union
 
 if TYPE_CHECKING:
     from datasette.app import Datasette
@@ -65,7 +69,22 @@ async def get_enrichments(datasette):
 """.strip()
 
 
+@hookimpl
+def register_secrets():
+    secrets = []
+    for subclass in Enrichment._subclasses:
+        if subclass.secret:
+            secrets.append(subclass.secret)
+    return secrets
+
+
+class SecretError(Exception):
+    pass
+
+
 class Enrichment(ABC):
+    _subclasses = []
+
     batch_size: int = 100
     # Cancel run after this many errors
     default_max_errors: int = 5
@@ -83,11 +102,36 @@ def name(self):
         # The name of this enrichment
         ...
 
-    description = ""  # Short description of this enrichment
+    description: str = ""  # Short description of this enrichment
+    secret: Optional[Secret] = None
+
+    def __init_subclass__(cls, **kwargs):
+        super().__init_subclass__(**kwargs)
+        cls._subclasses.append(cls)
 
     def __repr__(self):
         return "<Enrichment: {}>".format(self.slug)
 
+    async def get_secret(self, datasette: "Datasette", config: dict):
+        if self.secret is None:
+            breakpoint()
+            raise SecretError("No secret defined for this enrichment")
+        secret = await get_secret(datasette, self.secret.name)
+        if secret is not None:
+            return secret
+        # Try the stashed secrets instead
+        if not hasattr(datasette, "_enrichments_stashed_secrets"):
+            breakpoint()
+            raise SecretError("No secrets have been stashed")
+        stashed_keys = datasette._enrichments_stashed_secrets
+        stash_key = config.get("enrichment_secret")
+        if stash_key not in stashed_keys:
+            breakpoint()
+            raise SecretError(
+                "No secret found in stash for {}".format(self.secret.name)
+            )
+        return stashed_keys[stash_key]
+
     async def log_error(
         self, db: "Database", job_id: int, ids: List[IdType], error: str
     ):
@@ -113,6 +157,51 @@ async def log_error(
     async def get_config_form(self, datasette: "Datasette", db: "Database", table: str):
         return None
 
+    async def _get_config_form(
+        self, datasette: "Datasette", db: "Database", table: str
+    ):
+        # Helper method that adds a `_secret` form field if the enrichment has a secret
+        FormClass = await async_call_with_supported_arguments(
+            self.get_config_form, datasette=datasette, db=db, table=table
+        )
+        if self.secret is None:
+            return FormClass
+        # If secret is already set, return form unmodified
+        if await get_secret(datasette, self.secret.name):
+            return FormClass
+
+        # Otherwise, return form with secret field
+        def stash_api_key(form, field):
+            if not hasattr(datasette, "_enrichments_stashed_secrets"):
+                datasette._enrichments_stashed_secrets = {}
+            key = secrets.token_urlsafe(16)
+            datasette._enrichments_stashed_secrets[key] = field.data
+            field.data = key
+
+        formatted_description = self.secret.description
+        if self.secret.obtain_url and self.secret.obtain_label:
+            html_bits = []
+            if self.secret.description:
+                html_bits.append(escape(self.secret.description))
+                html_bits.append(" - ")
+            html_bits.append(
+                f'<a href="{self.secret.obtain_url}" target="_blank">{self.secret.obtain_label}</a>'
+            )
+            formatted_description = Markup("".join(html_bits))
+
+        class FormWithSecret(FormClass):
+            enrichment_secret = PasswordField(
+                self.secret.name,
+                description=formatted_description,
+                validators=[
+                    DataRequired(message="Secret is required."),
+                    stash_api_key,
+                ],
+                render_kw={"autocomplete": "off"},
+            )
+
+        return FormWithSecret
+
     async def initialize(
         self, datasette: "Datasette", db: "Database", table: str, config: dict
     ):
diff --git a/datasette_enrichments/views.py b/datasette_enrichments/views.py
index a435d67..d9966f2 100644
--- a/datasette_enrichments/views.py
+++ b/datasette_enrichments/views.py
@@ -65,7 +65,7 @@ async def enrichment_view(datasette, request):
 
     form = (
         await async_call_with_supported_arguments(
-            enrichment.get_config_form,
+            enrichment._get_config_form,
             datasette=datasette,
             db=datasette.get_database(database),
             table=table,
@@ -170,7 +170,7 @@ async def enrich_data_post(datasette, request, enrichment, filtered_data):
     post_vars = MultiParams(urllib.parse.parse_qs(body.decode("utf-8")))
 
     Form = await async_call_with_supported_arguments(
-        enrichment.get_config_form, datasette=datasette, db=db, table=table
+        enrichment._get_config_form, datasette=datasette, db=db, table=table
     )
 
     form = Form(post_vars)
diff --git a/docs/developing.md b/docs/developing.md
index 01abc88..33a7952 100644
--- a/docs/developing.md
+++ b/docs/developing.md
@@ -203,6 +203,45 @@ class MyEnrichment(Enrichment):
     log_traceback = True
 ```
 
+## Enrichments that use secrets such as API keys
+
+Enrichments often need to access external APIs, for things like geocoding addresses or accessing Large Language Models. These APIs frequently need some kind of API key.
+
+The enrichments system can work in combination with the [datasette-secrets](https://datasette.io/plugins/datasette-secrets) plugin to manage secret API keys needed for different services.
+
+Examples of enrichments that use this mechanism include [datasette-enrichments-opencage](https://datasette.io/plugins/datasette-enrichments-opencage) and [datasette-enrichments-gpt](https://datasette.io/plugins/datasette-enrichments-gpt).
+
+To define a secret that your plugin needs, add the following code:
+
+```python
+from datasette_enrichments import Enrichment
+from datasette_secrets import Secret
+
+# Then later in your enrichments class
+class TrainEnthusiastsEnrichment(Enrichment):
+    name = "Train Enthusiasts"
+    slug = "train-enthusiasts"
+    description = "Enrich with extra data from the Train Enthusiasts API"
+    secret = Secret(
+        name="TRAIN_ENTHUSIASTS_API_KEY",
+        description="An API key from train-enthusiasts.doesnt.exist",
+        obtain_url="https://train-enthusiasts.doesnt.exist/api-keys",
+        obtain_label="Get an API key"
+    )
+```
+Configuring your enrichment like this will result in the following behavior:
+
+- If a user sets a `DATASETTE_SECRETS_TRAIN_ENTHUSIASTS_API_KEY` environment variable, that value will be used as the API key. You should mention this in your enrichment's documentation.
+- Otherwise, if the user has configured [datasette-secrets](https://datasette.io/plugins/datasette-secrets) to store secrets in the database, admin users with the `manage-secrets` permission will get the option to set that secret as part of the Datasette web UI.
+- If neither of the above are true, any time any user tries to run this enrichment they will be prompted to specify a secret which will be used for that run but will not be stored outside of memory used by the code that runs the enrichment.
+
+With a secret configured in this way, your various class methods such as `initialize()`, `enrich_batch()` and `finalize()` can access the secret value using `await self.get_secret(datasette)` like this:
+
+```python
+api_key = await self.get_secret(datasette, config)
+```
+You must pass both the `datasette` and the `config` arguments that were passed to those methods.
+
 ## Writing tests for enrichments
 
 Take a look at the [test suite for datasette-enrichments-opencage](https://github.com/datasette/datasette-enrichments-opencage/blob/main/tests/test_enrichments_opencage.py) for an example of how to test an enrichment.
diff --git a/setup.py b/setup.py
index e4410a9..347563f 100644
--- a/setup.py
+++ b/setup.py
@@ -32,7 +32,7 @@ def get_long_description():
     version=VERSION,
     packages=["datasette_enrichments"],
     entry_points={"datasette": ["enrichments = datasette_enrichments"]},
-    install_requires=["datasette", "WTForms"],
+    install_requires=["datasette", "WTForms", "datasette-secrets>=0.2"],
     extras_require={
         "test": ["pytest", "pytest-asyncio", "black", "ruff", "packaging"],
         "docs": [


Test patch:
diff --git a/tests/conftest.py b/tests/conftest.py
index 48214c4..5d429c3 100644
--- a/tests/conftest.py
+++ b/tests/conftest.py
@@ -1,7 +1,7 @@
 import asyncio
 from datasette.database import Database
 from typing import List
-from wtforms import Form, SelectField
+from wtforms import Form, SelectField, StringField
 from wtforms.widgets import ListWidget, CheckboxInput
 import pytest
 from datasette.plugins import pm
@@ -16,6 +16,7 @@ class MultiCheckboxField(SelectField):
 @pytest.fixture(autouse=True)
 def load_uppercase_plugin():
     from datasette_enrichments import Enrichment
+    from datasette_secrets import Secret
 
     class UppercaseDemo(Enrichment):
         name = "Convert to uppercase"
@@ -60,15 +61,55 @@ async def enrich_batch(
             # Wait 0.3s
             await asyncio.sleep(0.3)
 
-    class UppercasePlugin:
-        __name__ = "UppercasePlugin"
+    class SecretReplacePlugin(Enrichment):
+        name = "Replace string with a secret"
+        slug = "secretreplace"
+        description = "Replace a string with a secret"
+        secret = Secret(
+            name="STRING_SECRET",
+            description="The secret to use in the replacement",
+        )
+
+        async def get_config_form(self, db, table):
+            choices = [(col, col) for col in await db.table_columns(table)]
+
+            class ConfigForm(Form):
+                column = SelectField("Column", choices=choices)
+                string = StringField("String to be replaced")
+
+            return ConfigForm
+
+        async def enrich_batch(
+            self,
+            datasette,
+            db: Database,
+            table: str,
+            rows: List[dict],
+            pks: List[str],
+            config: dict,
+            job_id: int,
+        ):
+            secret = await self.get_secret(datasette, config)
+            for row in rows:
+                await db.execute_write(
+                    "update [{}] set [{}] = ? where {}".format(
+                        table,
+                        config["column"],
+                        " and ".join('"{}" = ?'.format(pk) for pk in pks),
+                    ),
+                    [row[config["column"]].replace(config["string"], secret)]
+                    + [row[pk] for pk in pks],
+                )
+
+    class EnrichmentsDemoPlugin:
+        __name__ = "EnrichmentsDemoPlugin"
 
         @hookimpl
         def register_enrichments(self):
-            return [UppercaseDemo()]
+            return [UppercaseDemo(), SecretReplacePlugin()]
 
-    pm.register(UppercasePlugin(), name="undo_uppercase")
+    pm.register(EnrichmentsDemoPlugin(), name="undo_EnrichmentsDemoPlugin")
     try:
         yield
     finally:
-        pm.unregister(name="undo_uppercase")
+        pm.unregister(name="undo_EnrichmentsDemoPlugin")
diff --git a/tests/test_enrichments.py b/tests/test_enrichments.py
index 8ebc2bf..b45b56c 100644
--- a/tests/test_enrichments.py
+++ b/tests/test_enrichments.py
@@ -179,3 +179,58 @@ async def test_row_actions(datasette, path, expected_path):
     enrich_page_response = await datasette.client.get(enrich_path, cookies=cookies)
     assert enrich_page_response.status_code == 200
     assert "1 row selected" in enrich_page_response.text
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize("scenario", ("env", "user-input"))
+async def test_enrichment_using_secret(datasette, scenario, monkeypatch):
+    if scenario == "env":
+        monkeypatch.setenv("DATASETTE_SECRETS_STRING_SECRET", "env-secret")
+
+    cookies = {"ds_actor": datasette.sign({"a": {"id": "root"}}, "actor")}
+    response1 = await datasette.client.get("/-/enrich/data/t", cookies=cookies)
+    assert response1.status_code == 200
+    assert (
+        '<a href="/-/enrich/data/t/secretreplace">Replace string with a secret</a>'
+        in response1.text
+    )
+    response2 = await datasette.client.get(
+        "/-/enrich/data/t/secretreplace", cookies=cookies
+    )
+    assert "<h2>Replace string with a secret</h2>" in response2.text
+
+    # name="enrichment_secret" should be present only if not set in env
+    if scenario == "env":
+        assert ' name="enrichment_secret"' not in response2.text
+    else:
+        assert ' name="enrichment_secret"' in response2.text
+
+    # Now try and run it
+    csrftoken = response2.cookies["ds_csrftoken"]
+    cookies["ds_csrftoken"] = csrftoken
+
+    form_data = {"column": "s", "string": "hello", "csrftoken": csrftoken}
+    if scenario == "user-input":
+        form_data["enrichment_secret"] = "user-secret"
+
+    response3 = await datasette.client.post(
+        "/-/enrich/data/t/secretreplace",
+        cookies=cookies,
+        data=form_data,
+    )
+    assert response3.status_code == 302
+    job_id = response3.headers["location"].split("=")[-1]
+
+    # Wait for it to finish and check it worked
+    await wait_for_job(datasette, job_id, "data", timeout=1)
+    # Check for errors
+    job_details = datasette._test_db.execute(
+        "select error_count, done_count from _enrichment_jobs where id = ?", (job_id,)
+    ).fetchone()
+    assert job_details == (0, 2)
+    # Check rows show enrichment ran correctly
+    rows = datasette._test_db.execute("select s from t order by id").fetchall()
+    if scenario == "env":
+        assert rows == [("env-secret",), ("goodbye",)]
+    else:
+        assert rows == [("user-secret",), ("goodbye",)]