From 2258b431d4f7c543594bf0e7b9d3c6ee4699f8a3 Mon Sep 17 00:00:00 2001
From: Simon Willison <swillison@gmail.com>
Date: Fri, 8 Nov 2024 12:11:17 -0800
Subject: [PATCH] Neater return_db pattern for reusing memory command, closes
 #643

---
 docs/plugins.rst         |  9 ++-------
 sqlite_utils/cli.py      | 16 +++++-----------
 tests/test_cli_memory.py | 29 +++++++++++------------------
 3 files changed, 18 insertions(+), 36 deletions(-)

diff --git a/docs/plugins.rst b/docs/plugins.rst
index b9e4362..9f5c25e 100644
--- a/docs/plugins.rst
+++ b/docs/plugins.rst
@@ -117,13 +117,10 @@ Example implementation:
 
 New commands implemented by plugins can invoke existing commands using the `context.invoke <https://click.palletsprojects.com/en/stable/api/#click.Context.invoke>`__ mechanism.
 
-As a special niche feature, if your plugin needs to import some files and then act against an in-memory database containing those files you can forward to the :ref:`sqlite-utils memory command <cli_memory>` and then access a named in-memory database called ``sqlite_utils_memory`` like this:
+As a special niche feature, if your plugin needs to import some files and then act against an in-memory database containing those files you can forward to the :ref:`sqlite-utils memory command <cli_memory>` and pass it ``return_db=True``:
 
 .. code-block:: python
 
-    from contextlib import redirect_stdout
-    import io
-
     @cli.command()
     @click.pass_context
     @click.argument(
@@ -134,9 +131,7 @@ As a special niche feature, if your plugin needs to import some files and then a
     )
     def show_schema_for_files(ctx, paths):
         from sqlite_utils.cli import memory
-        with redirect_stdout(io.StringIO()):
-            ctx.invoke(memory, paths=paths, sql="select 1")
-        db = sqlite_utils.Database(memory_name="sqlite_utils_memory")
+        db = ctx.invoke(memory, paths=paths, _return_db=True)
         # Now do something with that database
         click.echo(db.schema)
 
diff --git a/sqlite_utils/cli.py b/sqlite_utils/cli.py
index 0e4efbf..cc2d86e 100644
--- a/sqlite_utils/cli.py
+++ b/sqlite_utils/cli.py
@@ -1815,10 +1815,6 @@ def query(
     )
 
 
-# So named memory db "sqlite_utils_memory" isn't dropped when it goes out of scope
-_sqlite_utils_memory_db = []
-
-
 @cli.command()
 @click.argument(
     "paths",
@@ -1898,6 +1894,7 @@ def memory(
     save,
     analyze,
     load_extension,
+    return_db=False,
 ):
     """Execute SQL query against an in-memory database, optionally populated by imported data
 
@@ -1925,13 +1922,7 @@ def memory(
     \b
         sqlite-utils memory animals.csv --schema
     """
-    if getattr(sys, "_sqlite_utils_memory_test", False) or not getattr(
-        sys, "_called_from_test", False
-    ):
-        db = sqlite_utils.Database(memory_name="sqlite_utils_memory")
-        _sqlite_utils_memory_db.append(db)
-    else:
-        db = sqlite_utils.Database(memory=True)
+    db = sqlite_utils.Database(memory=True)
 
     # If --dump or --save or --analyze used but no paths detected, assume SQL query is a path:
     if (dump or save or schema or analyze) and not paths:
@@ -2007,6 +1998,9 @@ def memory(
     if functions:
         _register_functions(db, functions)
 
+    if return_db:
+        return db
+
     _execute_query(
         db,
         sql,
diff --git a/tests/test_cli_memory.py b/tests/test_cli_memory.py
index cd26eab..ac0a177 100644
--- a/tests/test_cli_memory.py
+++ b/tests/test_cli_memory.py
@@ -1,6 +1,6 @@
+import click
 import json
 import pytest
-import sys
 from click.testing import CliRunner
 
 from sqlite_utils import Database, cli
@@ -307,21 +307,14 @@ def test_memory_functions():
     assert result.output.strip() == '[{"hello()": "Hello"}]'
 
 
-@pytest.mark.parametrize("enabled", (False, True))
-def test_memory_named_database_hack(enabled):
+def test_memory_return_db(tmpdir):
     # https://github.com/simonw/sqlite-utils/issues/643
-    sys._sqlite_utils_memory_test = enabled
-    try:
-        result = CliRunner().invoke(
-            cli.cli,
-            ["memory", "-", "--analyze"],
-            input="id,name\n1,Cleo\n2,Bants",
-        )
-        assert result.exit_code == 0
-        db = Database(memory_name="sqlite_utils_memory")
-        if enabled:
-            assert db.table_names() == ["stdin"]
-        else:
-            assert db.table_names() == []
-    finally:
-        sys._sqlite_utils_memory_test = False
+    from sqlite_utils.cli import cli
+
+    path = str(tmpdir / "dogs.csv")
+    open(path, "w").write("id,name\n1,Cleo")
+
+    with click.Context(cli) as ctx:
+        db = ctx.invoke(cli.commands["memory"], paths=(path,), return_db=True)
+
+    assert db.table_names() == ["dogs"]