From c7234cae8336b8525034e8f917d82dd0699abd42 Mon Sep 17 00:00:00 2001 From: Simon Willison Date: Tue, 15 Jun 2021 22:02:18 -0700 Subject: [PATCH] Initial CSV-only prototype of sqlite-utils memory, refs #272 --- docs/cli.rst | 23 +++++++++++ sqlite_utils/cli.py | 89 +++++++++++++++++++++++++++++++++++++++- tests/test_cli_memory.py | 32 +++++++++++++++ 3 files changed, 142 insertions(+), 2 deletions(-) create mode 100644 tests/test_cli_memory.py diff --git a/docs/cli.rst b/docs/cli.rst index c4dfc0d..060c5b3 100644 --- a/docs/cli.rst +++ b/docs/cli.rst @@ -201,6 +201,29 @@ For example, to retrieve a binary image from a ``BLOB`` column and store it in a $ sqlite-utils photos.db "select contents from photos where id=1" --raw > myphoto.jpg +.. _cli_query_memory: + +Running queries directly against CSV data +========================================= + +If you have data in CSV format you can load it into an in-memory SQLite database and run queries against it directly in a single command using ``sqlite-utils memory``:: + + $ sqlite-utils memory data.csv "select * from data" + +This command supports the same output formats as ``sqlite-utils query`` - so you can use ``--csv`` or ``--tsv`` or ``--nl`` or ``-t`` to format the data. + +You can pass multiple files to the command if you want to run joins between different CSV files:: + + $ sqlite-utils memory one.csv two.csv "select * from one where id in (select id from two)" + +The in-memory tables will be named after the CSV files without their ``.csv`` extension. The tool also sets up aliases for those tables (using SQL views) as ``t1``, ``t2`` and so on, or you can use the alias ``t`` to refer to the first table:: + + $ sqlite-utils memory example.csv "select * from t" + +To read from standard input, use ``-`` as the filename - then use ``stdin`` or ``t`` or ``t1`` as the table name:: + + $ cat example.csv | sqlite-utils memory - "select * from stdin" + .. _cli_rows: Returning all rows in a table diff --git a/sqlite_utils/cli.py b/sqlite_utils/cli.py index 4208830..d6a63c0 100644 --- a/sqlite_utils/cli.py +++ b/sqlite_utils/cli.py @@ -694,11 +694,11 @@ def insert_upsert_implementation( raise click.ClickException("Use just one of --nl, --csv or --tsv") if encoding and not (csv or tsv): raise click.ClickException("--encoding must be used with --csv or --tsv") + if pk and len(pk) == 1: + pk = pk[0] encoding = encoding or "utf-8-sig" buffered = io.BufferedReader(json_file, buffer_size=4096) decoded = io.TextIOWrapper(buffered, encoding=encoding) - if pk and len(pk) == 1: - pk = pk[0] if csv or tsv: if sniff: # Read first 2048 bytes and use that to detect @@ -1087,6 +1087,91 @@ def query( db.attach(alias, attach_path) _load_extensions(db, load_extension) db.register_fts4_bm25() + + _execute_query( + db, sql, param, raw, table, csv, tsv, no_headers, fmt, nl, arrays, json_cols + ) + + +@cli.command() +@click.argument( + "paths", + type=click.Path(file_okay=True, dir_okay=False, allow_dash=True), + required=False, + nargs=-1, +) +@click.argument("sql") +@click.option( + "--attach", + type=(str, click.Path(file_okay=True, dir_okay=False, allow_dash=False)), + multiple=True, + help="Additional databases to attach - specify alias and filepath", +) +@output_options +@click.option("-r", "--raw", is_flag=True, help="Raw output, first column of first row") +@click.option( + "-p", + "--param", + multiple=True, + type=(str, str), + help="Named :parameters for SQL query", +) +@click.option("--dump", is_flag=True, help="Dump SQL for in-memory database") +@load_extension_option +def memory( + paths, + sql, + attach, + nl, + arrays, + csv, + tsv, + no_headers, + table, + fmt, + json_cols, + raw, + param, + dump, + load_extension, +): + "Execute SQL query against an in-memory database, optionally populated by imported data" + db = sqlite_utils.Database(memory=True) + for i, path in enumerate(paths): + if path == "-": + csv_fp = sys.stdin + csv_table = "stdin" + else: + csv_path = pathlib.Path(path) + csv_table = csv_path.stem + csv_fp = csv_path.open() + db[csv_table].insert_all(csv_std.DictReader(csv_fp)) + # Add convenient t / t1 / t2 views + view_names = ["t{}".format(i + 1)] + if i == 0: + view_names.append("t") + for view_name in view_names: + if not db[view_name].exists(): + db.create_view(view_name, "select * from [{}]".format(csv_table)) + + if dump: + for line in db.conn.iterdump(): + click.echo(line) + return + + for alias, attach_path in attach: + db.attach(alias, attach_path) + _load_extensions(db, load_extension) + db.register_fts4_bm25() + + _execute_query( + db, sql, param, raw, table, csv, tsv, no_headers, fmt, nl, arrays, json_cols + ) + + +def _execute_query( + db, sql, param, raw, table, csv, tsv, no_headers, fmt, nl, arrays, json_cols +): with db.conn: try: cursor = db.execute(sql, dict(param)) diff --git a/tests/test_cli_memory.py b/tests/test_cli_memory.py new file mode 100644 index 0000000..df72ec1 --- /dev/null +++ b/tests/test_cli_memory.py @@ -0,0 +1,32 @@ +from sqlite_utils import cli, Database +from click.testing import CliRunner +import pytest + + +def test_memory_basic(): + result = CliRunner().invoke(cli.cli, ["memory", "select 1 + 1"]) + assert result.output.strip() == '[{"1 + 1": 2}]' + + +@pytest.mark.parametrize("sql_from", ("test", "t", "t1")) +@pytest.mark.parametrize("use_stdin", (True, False)) +def test_memory_csv(tmpdir, sql_from, use_stdin): + content = "id,name\n1,Cleo\n2,Bants" + input = None + if use_stdin: + input = content + csv_path = "-" + if sql_from == "test": + sql_from = "stdin" + else: + csv_path = str(tmpdir / "test.csv") + open(csv_path, "w").write(content) + result = CliRunner().invoke( + cli.cli, + ["memory", csv_path, "select * from {}".format(sql_from), "--nl"], + input=input, + ) + assert ( + result.output.strip() + == '{"id": "1", "name": "Cleo"}\n{"id": "2", "name": "Bants"}' + )