From 211c2c679189469a95da55877ed6df53b129a14a Mon Sep 17 00:00:00 2001
From: Simon Willison <swillison@gmail.com>
Date: Sun, 19 Nov 2023 20:58:14 -0800
Subject: [PATCH] Fix for config dict with strings vs lists, closes #18

Also includes datasette fix from #16
---
 datasette_enrichments/views.py           | 4 ++--
 example-enrichments/jinja_sandbox.py     | 6 +++---
 example-enrichments/openai_embeddings.py | 5 +++--
 example-enrichments/uppercase.py         | 5 +++--
 tests/conftest.py                        | 1 +
 tests/test_enrichments.py                | 2 +-
 6 files changed, 13 insertions(+), 10 deletions(-)

diff --git a/datasette_enrichments/views.py b/datasette_enrichments/views.py
index cdbf5bb..bc0af82 100644
--- a/datasette_enrichments/views.py
+++ b/datasette_enrichments/views.py
@@ -144,6 +144,7 @@ async def enrich_data_post(datasette, request, enrichment, stuff):
     Form = await enrichment.get_config_form(db, table)
 
     form = Form(post_vars)
+
     if not form.validate():
         # TODO: Fix this
         return Response.html(
@@ -169,8 +170,7 @@ async def enrich_data_post(datasette, request, enrichment, stuff):
             )
         )
 
-    config = post_vars._data.copy()
-    config.pop("csrftoken", None)
+    config = {field.name: field.data for field in form}
 
     # Initialize any necessary tables
     await enrichment.initialize(datasette, db, table, config)
diff --git a/example-enrichments/jinja_sandbox.py b/example-enrichments/jinja_sandbox.py
index 3e43ca6..1192a3d 100644
--- a/example-enrichments/jinja_sandbox.py
+++ b/example-enrichments/jinja_sandbox.py
@@ -38,7 +38,7 @@ class ConfigForm(Form):
 
     async def initialize(self, datasette, db, table, config):
         # Ensure column exists
-        output_column = config["output_column"][0]
+        output_column = config["output_column"]
 
         def add_column_if_not_exists(conn):
             db = sqlite_utils.Database(conn)
@@ -58,8 +58,8 @@ async def enrich_batch(
         job_id: int,
     ):
         env = SandboxedEnvironment(enable_async=True)
-        template = env.from_string(config["template"][0])
-        output_column = config["output_column"][0]
+        template = env.from_string(config["template"])
+        output_column = config["output_column"]
         for row in rows:
             output = await template.render_async({"row": row})
             await db.execute_write(
diff --git a/example-enrichments/openai_embeddings.py b/example-enrichments/openai_embeddings.py
index bcdd537..85c6ea7 100644
--- a/example-enrichments/openai_embeddings.py
+++ b/example-enrichments/openai_embeddings.py
@@ -87,6 +87,7 @@ async def initialize(self, datasette, db, table, config):
 
     async def enrich_batch(
         self,
+        datasette,
         db: Database,
         table: str,
         rows: List[dict],
@@ -94,9 +95,9 @@ async def enrich_batch(
         config: dict,
         job_id: int,
     ):
-        template = SpaceTemplate(config["template"][0])
+        template = SpaceTemplate(config["template"])
         texts = [template.safe_substitute(row) for row in rows]
-        token = config["api_token"][0]
+        token = config["api_token"]
 
         async with httpx.AsyncClient() as client:
             response = await client.post(
diff --git a/example-enrichments/uppercase.py b/example-enrichments/uppercase.py
index 3240baf..0ed494a 100644
--- a/example-enrichments/uppercase.py
+++ b/example-enrichments/uppercase.py
@@ -1,12 +1,12 @@
 from datasette_enrichments import Enrichment
 from datasette.database import Database
 from typing import List
-from wtforms import Form, SelectField
+from wtforms import Form, SelectMultipleField
 from wtforms.widgets import ListWidget, CheckboxInput
 import asyncio
 
 
-class MultiCheckboxField(SelectField):
+class MultiCheckboxField(SelectMultipleField):
     widget = ListWidget(prefix_label=False)
     option_widget = CheckboxInput()
 
@@ -27,6 +27,7 @@ class ConfigForm(Form):
 
     async def enrich_batch(
         self,
+        datasette,
         db: Database,
         table: str,
         rows: List[dict],
diff --git a/tests/conftest.py b/tests/conftest.py
index c347590..502c4f2 100644
--- a/tests/conftest.py
+++ b/tests/conftest.py
@@ -32,6 +32,7 @@ class ConfigForm(Form):
 
         async def enrich_batch(
             self,
+            datasette,
             db: Database,
             table: str,
             rows: List[dict],
diff --git a/tests/test_enrichments.py b/tests/test_enrichments.py
index 3211e3b..b15c1b5 100644
--- a/tests/test_enrichments.py
+++ b/tests/test_enrichments.py
@@ -54,4 +54,4 @@ async def test_uppercase_plugin(tmpdir, is_root):
     # It should be queued up
     assert db.execute(
         "select status, enrichment, database_name, table_name, config from _enrichment_jobs"
-    ).fetchall() == [("p", "uppercasedemo", "data", "t", '{"columns": ["s"]}')]
+    ).fetchall() == [("p", "uppercasedemo", "data", "t", '{"columns": "s"}')]