From 49a010c93d90bc68ce1c6fff7639927248912b54 Mon Sep 17 00:00:00 2001 From: Simon Willison Date: Tue, 24 Aug 2021 16:31:13 -0700 Subject: [PATCH] Ability to insert file contents as text, in addition to blob (#321) --- docs/cli.rst | 18 ++++++++----- sqlite_utils/cli.py | 45 +++++++++++++++++++++++++++---- tests/test_insert_files.py | 54 +++++++++++++++++++++++++++++++++++--- 3 files changed, 101 insertions(+), 16 deletions(-) diff --git a/docs/cli.rst b/docs/cli.rst index 3c06b81..0cc20dd 100644 --- a/docs/cli.rst +++ b/docs/cli.rst @@ -331,7 +331,7 @@ The CSV data that was piped into the script is available in the ``stdin`` table, \-\-schema, \-\-analyze, \-\-dump and \-\-save ---------------------------------------------- -To see the in-memory datbase schema that would be used for a file or for multiple files, use ``--schema``:: +To see the in-memory database schema that would be used for a file or for multiple files, use ``--schema``:: % sqlite-utils memory dogs.csv --schema CREATE TABLE [dogs] ( @@ -909,12 +909,10 @@ The command will fail if you reference columns that do not exist on the table. T .. _cli_insert_files: -Inserting binary data from files -================================ +Inserting data from files +========================= -SQLite ``BLOB`` columns can be used to store binary content. It can be useful to insert the contents of files into a SQLite table. - -The ``insert-files`` command can be used to insert the content of files, along with their metadata. +The ``insert-files`` command can be used to insert the content of files, along with their metadata, into a SQLite table. Here's an example that inserts all of the GIF files in the current directory into a ``gifs.db`` database, placing the file contents in an ``images`` table:: @@ -932,6 +930,8 @@ By default this command will create a table with the following schema:: [size] INTEGER ); +Content will be treated as binary by default and stored in a ``BLOB`` column. You can use the ``--text`` option to store that content in a ``TEXT`` column instead. + You can customize the schema using one or more ``-c`` options. For a table schema that includes just the path, MD5 hash and last modification time of the file, you would use this:: $ sqlite-utils insert-files gifs.db images *.gif -c path -c md5 -c mtime --pk=path @@ -944,6 +944,8 @@ This will result in the following schema:: [mtime] FLOAT ); +Note that there's no ``content`` column here at all - if you specify custom columns using ``-c`` you need to include ``-c content`` to create that column. + You can change the name of one of these columns using a ``-c colname:coldef`` parameter. To rename the ``mtime`` column to ``last_modified`` you would use this:: $ sqlite-utils insert-files gifs.db images *.gif \ @@ -967,6 +969,8 @@ The full list of column definitions you can use is as follows: The permission bits of the file, as an integer - you may want to convert this to octal ``content`` The binary file contents, which will be stored as a BLOB +``content_text`` + The text file contents, which will be stored as TEXT ``mtime`` The modification time of the file, as floating point seconds since the Unix epoch ``ctime`` @@ -988,7 +992,7 @@ You can insert data piped from standard input like this:: The ``-`` argument indicates data should be read from standard input. The string passed using the ``--name`` option will be used for the file name and path values. -When inserting data from standard input only the following column definitions are supported: ``name``, ``path``, ``content``, ``sha256``, ``md5`` and ``size``. +When inserting data from standard input only the following column definitions are supported: ``name``, ``path``, ``content``, ``content_text``, ``sha256``, ``md5`` and ``size``. .. _cli_convert: diff --git a/sqlite_utils/cli.py b/sqlite_utils/cli.py index 58cda49..10a11c6 100644 --- a/sqlite_utils/cli.py +++ b/sqlite_utils/cli.py @@ -1811,6 +1811,11 @@ def extract( @click.option("--replace", is_flag=True, help="Replace files with matching primary key") @click.option("--upsert", is_flag=True, help="Upsert files with matching primary key") @click.option("--name", type=str, help="File name to use") +@click.option("--text", is_flag=True, help="Store file content as TEXT, not BLOB") +@click.option( + "--encoding", + help="Character encoding for input, defaults to utf-8", +) @click.option("-s", "--silent", is_flag=True, help="Don't show a progress bar") @load_extension_option def insert_files( @@ -1823,6 +1828,8 @@ def insert_files( replace, upsert, name, + text, + encoding, silent, load_extension, ): @@ -1842,7 +1849,10 @@ def insert_files( --pk name """ if not column: - column = ["path:path", "content:content", "size:size"] + if text: + column = ["path:path", "content_text:content_text", "size:size"] + else: + column = ["path:path", "content:content", "size:size"] if not pk: pk = "path" @@ -1866,7 +1876,16 @@ def insert_files( def to_insert(): for path, relative_path in bar: row = {} - lookups = FILE_COLUMNS + # content_text is special case as it considers 'encoding' + + def _content_text(p): + resolved = p.resolve() + try: + return resolved.read_text(encoding=encoding) + except UnicodeDecodeError as e: + raise UnicodeDecodeErrorForPath(e, resolved) + + lookups = dict(FILE_COLUMNS, content_text=_content_text) if path == "-": stdin_data = sys.stdin.buffer.read() # We only support a subset of columns for this case @@ -1874,6 +1893,9 @@ def insert_files( "name": lambda p: name or "-", "path": lambda p: name or "-", "content": lambda p: stdin_data, + "content_text": lambda p: stdin_data.decode( + encoding or "utf-8" + ), "sha256": lambda p: hashlib.sha256(stdin_data).hexdigest(), "md5": lambda p: hashlib.md5(stdin_data).hexdigest(), "size": lambda p: len(stdin_data), @@ -1899,9 +1921,16 @@ def insert_files( db = sqlite_utils.Database(path) _load_extensions(db, load_extension) - with db.conn: - db[table].insert_all( - to_insert(), pk=pk, alter=alter, replace=replace, upsert=upsert + try: + with db.conn: + db[table].insert_all( + to_insert(), pk=pk, alter=alter, replace=replace, upsert=upsert + ) + except UnicodeDecodeErrorForPath as e: + raise click.ClickException( + UNICODE_ERROR.format( + "Could not read file '{}' as text\n\n{}".format(e.path, e.exception) + ) ) @@ -2149,6 +2178,12 @@ def _render_common(title, values): return "\n".join(lines) +class UnicodeDecodeErrorForPath(Exception): + def __init__(self, exception, path): + self.exception = exception + self.path = path + + FILE_COLUMNS = { "name": lambda p: p.name, "path": lambda p: str(p), diff --git a/tests/test_insert_files.py b/tests/test_insert_files.py index 1e30a8d..86ee4e3 100644 --- a/tests/test_insert_files.py +++ b/tests/test_insert_files.py @@ -3,6 +3,7 @@ from click.testing import CliRunner import os import pathlib import pytest +import sys @pytest.mark.parametrize("silent", (False, True)) @@ -23,6 +24,7 @@ def test_insert_files(silent): "md5", "mode", "content", + "content_text", "mtime", "ctime", "mtime_int", @@ -52,6 +54,7 @@ def test_insert_files(silent): ) assert { "content": b"This is file one", + "content_text": "This is file one", "md5": "556dfb57fce9ca301f914e2273adf354", "name": "one.txt", "path": "one.txt", @@ -60,6 +63,7 @@ def test_insert_files(silent): }.items() <= one.items() assert { "content": b"Two is shorter", + "content_text": "Two is shorter", "md5": "f86f067b083af1911043eb215e74ac70", "name": "two.txt", "path": "two.txt", @@ -68,6 +72,7 @@ def test_insert_files(silent): }.items() <= two.items() assert { "content": b"Three is nested", + "content_text": "Three is nested", "md5": "12580f341781f5a5b589164d3cd39523", "name": "three.txt", "path": os.path.join("nested", "three.txt"), @@ -84,24 +89,65 @@ def test_insert_files(silent): "mtime_iso": str, "mode": int, "fullpath": str, + "content": bytes, + "content_text": str, } for colname, expected_type in expected_types.items(): for row in (one, two, three): assert isinstance(row[colname], expected_type) -def test_insert_files_stdin(): +@pytest.mark.parametrize( + "use_text,encoding,input,expected", + ( + (False, None, "hello world", b"hello world"), + (True, None, "hello world", "hello world"), + (False, None, b"S\xe3o Paulo", b"S\xe3o Paulo"), + (True, "latin-1", b"S\xe3o Paulo", "S\xe3o Paulo"), + ), +) +def test_insert_files_stdin(use_text, encoding, input, expected): runner = CliRunner() with runner.isolated_filesystem(): tmpdir = pathlib.Path(".") db_path = str(tmpdir / "files.db") + args = ["insert-files", db_path, "files", "-", "--name", "stdin-name"] + if use_text: + args += ["--text"] + if encoding is not None: + args += ["--encoding", encoding] result = runner.invoke( cli.cli, - ["insert-files", db_path, "files", "-", "--name", "stdin-name"], + args, catch_exceptions=False, - input="hello world", + input=input, ) assert result.exit_code == 0, result.stdout db = Database(db_path) row = list(db["files"].rows)[0] - assert {"path": "stdin-name", "content": b"hello world", "size": 11} == row + key = "content" + if use_text: + key = "content_text" + assert {"path": "stdin-name", key: expected}.items() <= row.items() + + +@pytest.mark.skipif( + sys.platform.startswith("win"), + reason="Windows has a different way of handling default encodings", +) +def test_insert_files_bad_text_encoding_error(): + runner = CliRunner() + with runner.isolated_filesystem(): + tmpdir = pathlib.Path(".") + latin = tmpdir / "latin.txt" + latin.write_bytes(b"S\xe3o Paulo") + db_path = str(tmpdir / "files.db") + result = runner.invoke( + cli.cli, + ["insert-files", db_path, "files", str(latin), "--text"], + catch_exceptions=False, + ) + assert result.exit_code == 1, result.output + assert result.output.strip().startswith( + "Error: Could not read file '{}' as text".format(str(latin.resolve())) + )