Skip to content

Commit

Permalink
Hack to support reuse of memory command, closes #643
Browse files Browse the repository at this point in the history
  • Loading branch information
simonw committed Nov 8, 2024
1 parent 4c26288 commit 8906f57
Show file tree
Hide file tree
Showing 4 changed files with 65 additions and 2 deletions.
25 changes: 25 additions & 0 deletions docs/plugins.rst
Original file line number Diff line number Diff line change
Expand Up @@ -115,6 +115,31 @@ Example implementation:
"Say hello world"
click.echo("Hello world!")
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:

.. code-block:: python
from contextlib import redirect_stdout
import io
@cli.command()
@click.pass_context
@click.argument(
"paths",
type=click.Path(file_okay=True, dir_okay=False, allow_dash=True),
required=False,
nargs=-1,
)
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")
# Now do something with that database
click.echo(db.schema)
.. _plugins_hooks_prepare_connection:

prepare_connection(conn)
Expand Down
15 changes: 14 additions & 1 deletion sqlite_utils/cli.py
Original file line number Diff line number Diff line change
Expand Up @@ -1815,6 +1815,10 @@ 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",
Expand Down Expand Up @@ -1921,7 +1925,14 @@ def memory(
\b
sqlite-utils memory animals.csv --schema
"""
db = sqlite_utils.Database(memory=True)
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)

# 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:
paths = [sql]
Expand Down Expand Up @@ -1954,6 +1965,7 @@ def memory(
rows = tracker.wrap(rows)
if flatten:
rows = (_flatten(row) for row in rows)

db[file_table].insert_all(rows, alter=True)
if tracker is not None:
db[file_table].transform(types=tracker.types)
Expand All @@ -1964,6 +1976,7 @@ def memory(
for view_name in view_names:
if not db[view_name].exists():
db.create_view(view_name, "select * from [{}]".format(file_table))

if fp:
fp.close()

Expand Down
5 changes: 5 additions & 0 deletions sqlite_utils/db.py
Original file line number Diff line number Diff line change
Expand Up @@ -325,6 +325,8 @@ def __init__(
execute_plugins: bool = True,
strict: bool = False,
):
self.memory_name = None
self.memory = False
assert (filename_or_conn is not None and (not memory and not memory_name)) or (
filename_or_conn is None and (memory or memory_name)
), "Either specify a filename_or_conn or pass memory=True"
Expand All @@ -335,8 +337,11 @@ def __init__(
uri=True,
check_same_thread=False,
)
self.memory = True
self.memory_name = memory_name
elif memory or filename_or_conn == ":memory:":
self.conn = sqlite3.connect(":memory:")
self.memory = True
elif isinstance(filename_or_conn, (str, pathlib.Path)):
if recreate and os.path.exists(filename_or_conn):
try:
Expand Down
22 changes: 21 additions & 1 deletion tests/test_cli_memory.py
Original file line number Diff line number Diff line change
@@ -1,6 +1,6 @@
import json

import pytest
import sys
from click.testing import CliRunner

from sqlite_utils import Database, cli
Expand Down Expand Up @@ -305,3 +305,23 @@ def test_memory_functions():
)
assert result.exit_code == 0
assert result.output.strip() == '[{"hello()": "Hello"}]'


@pytest.mark.parametrize("enabled", (False, True))
def test_memory_named_database_hack(enabled):
# 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

0 comments on commit 8906f57

Please sign in to comment.