Neater return_db pattern for reusing memory command, closes #643

This commit is contained in:
Simon Willison 2024-11-08 12:11:17 -08:00
commit 2258b431d4
3 changed files with 18 additions and 36 deletions

View file

@ -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)

View file

@ -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,

View file

@ -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"]