From 00e4bd5ff18ef4c3db6c1d67e2b974131c80d65c Mon Sep 17 00:00:00 2001 From: Simon Willison Date: Fri, 18 Jun 2021 20:11:54 -0700 Subject: [PATCH] TSV and JSON support for sqlite-utils memory Closes #281, closes #279, refs #272 --- docs/cli.rst | 51 +++++++++++++++++-------- sqlite_utils/cli.py | 24 ++++++++---- sqlite_utils/utils.py | 62 ++++++++++++++++++++++++++++++- tests/test_cli_memory.py | 80 ++++++++++++++++++++++++++++++++++++++-- 4 files changed, 191 insertions(+), 26 deletions(-) diff --git a/docs/cli.rst b/docs/cli.rst index 69d00a0..c6067ac 100644 --- a/docs/cli.rst +++ b/docs/cli.rst @@ -234,21 +234,21 @@ This example attaches the ``books.db`` database under the alias ``books`` and th sqlite-utils dogs.db --attach books books.db \ 'select * from sqlite_master union all select * from books.sqlite_master' -.. _cli_query_memory: +.. _cli_memory: -Querying CSV data directly using an in-memory database -====================================================== +Querying data directly using an in-memory database +================================================== The ``sqlite-utils memory`` command works similar to ``sqlite-utils query``, but allows you to execute queries against an in-memory database. -You can also pass this command CSV files which will be loaded into a temporary in-memory table, allowing you to execute SQL against that data without a separate step to first convert it to SQLite. +You can also pass this command CSV or JSON files which will be loaded into a temporary in-memory table, allowing you to execute SQL against that data without a separate step to first convert it to SQLite. Without any extra arguments, this command executes SQL against the in-memory database directly:: $ sqlite-utils memory 'select sqlite_version()' [{"sqlite_version()": "3.35.5"}] -It takes all of the same formatting options as :ref:`sqlite-utils query `: ``--csv`` and ``--csv`` and ``--table`` and ``--nl``:: +It takes all of the same output formatting options as :ref:`sqlite-utils query `: ``--csv`` and ``--csv`` and ``--table`` and ``--nl``:: $ sqlite-utils memory 'select sqlite_version()' --csv sqlite_version() @@ -260,24 +260,28 @@ It takes all of the same formatting options as :ref:`sqlite-utils query ` - so either a single JSON object (treated as a single row) or a list of JSON objects. + +CSV data can be comma- or tab- delimited. + +The in-memory tables will be named after the files without their extensions. 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:: +To read from standard input, use either ``-`` or ``stdin`` as the filename - then use ``stdin`` or ``t`` or ``t1`` as the table name:: $ cat example.csv | sqlite-utils memory - "select * from stdin" @@ -287,7 +291,24 @@ Incoming CSV data will be assumed to use ``utf-8``. If your data uses a differen If you are joining across multiple CSV files they must all use the same encoding. -.. _cli_query_memory_attach: +.. _cli_memory_explicit: + +Explicitly specifying the format +-------------------------------- + +By default, ``sqlite-utils memory`` will attempt to detect the incoming data format (JSON, TSV or CSV) automatically. + +You can instead specify an explicit format by adding a ``:csv``, ``:tsv``, ``:json`` or ``:nl`` (for newline-delimited JSON) suffix to the filename. For example:: + + $ sqlite-utils memory one.dat:csv two.dat:nl "select * from one union select * from two" + +Here the contents of ``one.dat`` will be treated as CSV and the contents of ``two.dat`` will be treated as newline-delimited JSON. + +To explicitly specify the format for data piped into the tool on standard input, use ``stdin:format`` - for example:: + + $ cat one.dat | sqlite-utils memory stdin:csv "select * from stdin" + +.. _cli_memory_attach: Joining in-memory data against existing databases using \-\-attach ------------------------------------------------------------------ @@ -303,7 +324,7 @@ Here the ``--attach trees trees.db`` option makes the ``trees.db`` database avai The CSV data that was piped into the script is available in the ``stdin`` table, so ``... where rowid in (select id from stdin)`` can be used to return rows from the ``trees`` table that match IDs that were piped in as CSV content. -.. _cli_query_memory_dump_save: +.. _cli_memory_dump_save: \-\-dump and \-\-save --------------------- diff --git a/sqlite_utils/cli.py b/sqlite_utils/cli.py index c5176f8..a4c1789 100644 --- a/sqlite_utils/cli.py +++ b/sqlite_utils/cli.py @@ -14,7 +14,14 @@ import os import sys import csv as csv_std import tabulate -from .utils import file_progress, find_spatialite, sqlite3, decode_base64_values +from .utils import ( + file_progress, + find_spatialite, + sqlite3, + decode_base64_values, + rows_from_file, + Format, +) CONTEXT_SETTINGS = dict(help_option_names=["-h", "--help"]) @@ -1175,18 +1182,21 @@ def memory( paths = [sql] sql = None for i, path in enumerate(paths): - if path == "-": + # Path may have a :format suffix + if ":" in path and path.rsplit(":", 1)[-1].upper() in Format.__members__: + path, suffix = path.rsplit(":", 1) + format = Format[suffix.upper()] + else: + format = None + if path in ("-", "stdin"): csv_fp = sys.stdin.buffer csv_table = "stdin" else: csv_path = pathlib.Path(path) csv_table = csv_path.stem csv_fp = csv_path.open("rb") - - encoding = encoding or "utf-8-sig" - decoded_fp = io.TextIOWrapper(csv_fp, encoding=encoding) - - db[csv_table].insert_all(csv_std.DictReader(decoded_fp)) + rows = rows_from_file(csv_fp, format=format, encoding=encoding) + db[csv_table].insert_all(rows, alter=True) # Add convenient t / t1 / t2 views view_names = ["t{}".format(i + 1)] if i == 0: diff --git a/sqlite_utils/utils.py b/sqlite_utils/utils.py index 9d5dac6..8a0b9ab 100644 --- a/sqlite_utils/utils.py +++ b/sqlite_utils/utils.py @@ -1,8 +1,13 @@ import base64 -import click import contextlib +import csv +import enum import io +import json import os +from typing import Generator + +import click try: import pysqlite3 as sqlite3 @@ -111,3 +116,58 @@ def file_progress(file, silent=False, **kwargs): file_length = os.path.getsize(file.name) with click.progressbar(length=file_length, **kwargs) as bar: yield UpdateWrapper(file, bar.update) + + +class Format(enum.Enum): + CSV = 1 + TSV = 2 + JSON = 3 + NL = 4 + + +class RowsFromFileError(Exception): + pass + + +class RowsFromFileBadJSON(RowsFromFileError): + pass + + +def rows_from_file( + fp, + format=None, + dialect=None, + encoding=None, +) -> Generator[dict, None, None]: + if format == Format.JSON: + decoded = json.load(fp) + if isinstance(decoded, dict): + decoded = [decoded] + if not isinstance(decoded, list): + raise RowsFromFileBadJSON("JSON must be a list or a dictionary") + yield from decoded + elif format == Format.NL: + yield from (json.loads(line) for line in fp if line.strip()) + elif format == Format.CSV: + decoded_fp = io.TextIOWrapper(fp, encoding=encoding or "utf-8-sig") + yield from csv.DictReader(decoded_fp, dialect=dialect) + elif format == Format.TSV: + yield from rows_from_file( + fp, format=Format.CSV, dialect=csv.excel_tab, encoding=encoding + ) + elif format is None: + # Detect the format, then call this recursively + buffered = io.BufferedReader(fp, buffer_size=4096) + first_bytes = buffered.peek(2048).strip() + if first_bytes.startswith(b"[") or first_bytes.startswith(b"{"): + # TODO: Detect newline-JSON + yield from rows_from_file(buffered, format=Format.JSON) + else: + dialect = csv.Sniffer().sniff( + first_bytes.decode(encoding or "utf-8-sig", "ignore") + ) + yield from rows_from_file( + buffered, format=Format.CSV, dialect=dialect, encoding=encoding + ) + else: + raise RowsFromFileError("Bad format") diff --git a/tests/test_cli_memory.py b/tests/test_cli_memory.py index f45abdf..c91beaf 100644 --- a/tests/test_cli_memory.py +++ b/tests/test_cli_memory.py @@ -1,8 +1,10 @@ -from sqlite_utils import cli, Database -from click.testing import CliRunner -import pytest import json +import pytest +from click.testing import CliRunner + +from sqlite_utils import Database, cli + def test_memory_basic(): result = CliRunner().invoke(cli.cli, ["memory", "select 1 + 1"]) @@ -35,6 +37,78 @@ def test_memory_csv(tmpdir, sql_from, use_stdin): ) +@pytest.mark.parametrize("use_stdin", (True, False)) +def test_memory_tsv(tmpdir, use_stdin): + data = "id\tname\n1\tCleo\n2\tBants" + if use_stdin: + input = data + path = "stdin:tsv" + sql_from = "stdin" + else: + input = None + path = str(tmpdir / "chickens.tsv") + open(path, "w").write(data) + path = path + ":tsv" + sql_from = "chickens" + result = CliRunner().invoke( + cli.cli, + ["memory", path, "select * from {}".format(sql_from)], + input=data, + ) + assert result.exit_code == 0, result.output + assert json.loads(result.output.strip()) == [ + {"id": "1", "name": "Cleo"}, + {"id": "2", "name": "Bants"}, + ] + + +@pytest.mark.parametrize("use_stdin", (True, False)) +def test_memory_json(tmpdir, use_stdin): + data = '[{"name": "Bants"}, {"name": "Dori", "age": 1}]' + if use_stdin: + input = data + path = "stdin:json" + sql_from = "stdin" + else: + input = None + path = str(tmpdir / "chickens.json") + open(path, "w").write(data) + path = path + ":json" + sql_from = "chickens" + result = CliRunner().invoke( + cli.cli, + ["memory", path, "select * from {}".format(sql_from)], + input=input, + ) + assert result.exit_code == 0, result.output + assert json.loads(result.output.strip()) == [ + {"name": "Bants", "age": None}, + {"name": "Dori", "age": 1}, + ] + + +@pytest.mark.parametrize("use_stdin", (True, False)) +def test_memory_json_nl(tmpdir, use_stdin): + data = '{"name": "Bants"}\n\n{"name": "Dori"}' + if use_stdin: + input = data + path = "stdin:nl" + sql_from = "stdin" + else: + input = None + path = str(tmpdir / "chickens.json") + open(path, "w").write(data) + path = path + ":nl" + sql_from = "chickens" + result = CliRunner().invoke( + cli.cli, + ["memory", path, "select * from {}".format(sql_from)], + input=data, + ) + assert result.exit_code == 0, result.output + assert json.loads(result.output.strip()) == [{"name": "Bants"}, {"name": "Dori"}] + + @pytest.mark.parametrize("use_stdin", (True, False)) def test_memory_csv_encoding(tmpdir, use_stdin): latin1_csv = (