From 59e2cfbdc12082bac03e8ac6f99c8c41a4bc72ba Mon Sep 17 00:00:00 2001
From: Simon Willison <swillison@gmail.com>
Date: Fri, 26 Aug 2022 22:03:53 -0700
Subject: [PATCH] sqlite-utils memory --functions, refs #471

---
 docs/cli-reference.rst   |  2 ++
 sqlite_utils/cli.py      | 33 ++++++++++++++++++++++-----------
 tests/test_cli_memory.py |  9 +++++++++
 3 files changed, 33 insertions(+), 11 deletions(-)

diff --git a/docs/cli-reference.rst b/docs/cli-reference.rst
index af059f1..9a025ba 100644
--- a/docs/cli-reference.rst
+++ b/docs/cli-reference.rst
@@ -161,6 +161,8 @@ See :ref:`cli_memory`.
           sqlite-utils memory animals.csv --schema
 
     Options:
+      --functions TEXT            Python code defining one or more custom SQL
+                                  functions
       --attach <TEXT FILE>...     Additional databases to attach - specify alias and
                                   filepath
       --flatten                   Flatten nested JSON objects, so {"foo": {"bar":
diff --git a/sqlite_utils/cli.py b/sqlite_utils/cli.py
index 1ac8d78..24c8c03 100644
--- a/sqlite_utils/cli.py
+++ b/sqlite_utils/cli.py
@@ -1669,18 +1669,8 @@ def query(
     _load_extensions(db, load_extension)
     db.register_fts4_bm25()
 
-    # Register any Python functions as SQL functions:
     if functions:
-        sqlite3.enable_callback_tracebacks(True)
-        globals = {}
-        try:
-            exec(functions, globals)
-        except SyntaxError as ex:
-            raise click.ClickException("Error in functions definition: {}".format(ex))
-        # Register all callables in the locals dict:
-        for name, value in globals.items():
-            if callable(value) and not name.startswith("_"):
-                db.register_function(value, name=name)
+        _register_functions(db, functions)
 
     _execute_query(
         db, sql, param, raw, table, csv, tsv, no_headers, fmt, nl, arrays, json_cols
@@ -1695,6 +1685,9 @@ def query(
     nargs=-1,
 )
 @click.argument("sql")
+@click.option(
+    "--functions", help="Python code defining one or more custom SQL functions"
+)
 @click.option(
     "--attach",
     type=(str, click.Path(file_okay=True, dir_okay=False, allow_dash=False)),
@@ -1741,6 +1734,7 @@ def query(
 def memory(
     paths,
     sql,
+    functions,
     attach,
     flatten,
     nl,
@@ -1854,6 +1848,9 @@ def memory(
     _load_extensions(db, load_extension)
     db.register_fts4_bm25()
 
+    if functions:
+        _register_functions(db, functions)
+
     _execute_query(
         db, sql, param, raw, table, csv, tsv, no_headers, fmt, nl, arrays, json_cols
     )
@@ -2993,3 +2990,17 @@ def _load_extensions(db, load_extension):
             if ext == "spatialite" and not os.path.exists(ext):
                 ext = find_spatialite()
             db.conn.load_extension(ext)
+
+
+def _register_functions(db, functions):
+    # Register any Python functions as SQL functions:
+    sqlite3.enable_callback_tracebacks(True)
+    globals = {}
+    try:
+        exec(functions, globals)
+    except SyntaxError as ex:
+        raise click.ClickException("Error in functions definition: {}".format(ex))
+    # Register all callables in the locals dict:
+    for name, value in globals.items():
+        if callable(value) and not name.startswith("_"):
+            db.register_function(value, name=name)
diff --git a/tests/test_cli_memory.py b/tests/test_cli_memory.py
index aee74e7..bb50d23 100644
--- a/tests/test_cli_memory.py
+++ b/tests/test_cli_memory.py
@@ -289,3 +289,12 @@ def test_memory_two_files_with_same_stem(tmpdir):
         ");\n"
         "CREATE VIEW t2 AS select * from [data_2];\n"
     )
+
+
+def test_memory_functions():
+    result = CliRunner().invoke(
+        cli.cli,
+        ["memory", "select hello()", "--functions", "hello = lambda: 'Hello'"],
+    )
+    assert result.exit_code == 0
+    assert result.output.strip() == '[{"hello()": "Hello"}]'