diff --git a/docs/cli.rst b/docs/cli.rst index d896a6e..66c93cc 100644 --- a/docs/cli.rst +++ b/docs/cli.rst @@ -1740,6 +1740,14 @@ By default this command will create a table with the following schema: 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. +Pass ``--sqlar`` to store the content zlib-compressed instead, using the same ``name``, ``mode``, ``mtime``, ``sz`` and ``data`` columns as `SQLite's own sqlar archive format `__: + +.. code-block:: bash + + sqlite-utils insert-files archive.db sqlar *.gif --sqlar + +Content is only stored compressed if doing so makes it smaller - otherwise the original bytes are stored as-is, matching the behaviour of SQLite's ``sqlar_compress()`` function. + 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: .. code-block:: bash @@ -1792,6 +1800,8 @@ The full list of column definitions you can use is as follows: The binary file contents, which will be stored as a BLOB ``content_text`` The text file contents, which will be stored as TEXT +``content_sqlar`` + The file contents zlib-compressed for storage in a ``sqlar``-compatible BLOB column, matching the behaviour of ``--sqlar`` - the content is left uncompressed if compression would not make it smaller ``mtime`` The modification time of the file, as floating point seconds since the Unix epoch ``ctime`` @@ -1819,7 +1829,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``, ``content_text``, ``sha256``, ``md5`` and ``size``. +When inserting data from standard input only the following column definitions are supported: ``name``, ``path``, ``content``, ``content_text``, ``content_sqlar``, ``sha256``, ``md5`` and ``size``. .. _cli_convert: diff --git a/sqlite_utils/cli.py b/sqlite_utils/cli.py index a3c37fd..5609d65 100644 --- a/sqlite_utils/cli.py +++ b/sqlite_utils/cli.py @@ -11,6 +11,7 @@ import pathlib import pdb # noqa: T100 import sys import textwrap +import zlib from datetime import datetime, timezone from runpy import run_module from typing import Any @@ -2913,6 +2914,11 @@ def extract( @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( + "--sqlar", + is_flag=True, + help="Store file content zlib-compressed, compatible with SQLite's sqlar format", +) @click.option( "--encoding", help="Character encoding for input, defaults to utf-8", @@ -2930,6 +2936,7 @@ def insert_files( upsert, name, text, + sqlar, encoding, silent, load_extension, @@ -2949,13 +2956,23 @@ def insert_files( -c size:size \\ --pk name """ + if text and sqlar: + raise click.ClickException("Cannot use --text and --sqlar together") if not column: if text: column = ["path:path", "content_text:content_text", "size:size"] + elif sqlar: + column = [ + "name:name", + "mode:mode", + "mtime:mtime_int", + "sz:size", + "data:content_sqlar", + ] else: column = ["path:path", "content:content", "size:size"] if not pks: - pks = ["path"] + pks = ["name"] if sqlar else ["path"] def yield_paths_and_relative_paths(): for f_or_d in file_or_dir: @@ -2997,6 +3014,9 @@ def insert_files( "content_text": lambda p, data=stdin_data: data.decode( encoding or "utf-8" ), + "content_sqlar": lambda p, data=stdin_data: sqlar_compress( + data + ), "sha256": lambda p, data=stdin_data: hashlib.sha256( data ).hexdigest(), @@ -3688,6 +3708,14 @@ class UnicodeDecodeErrorForPath(Exception): self.path = path +def sqlar_compress(data): + # Matches SQLite's sqlar_compress(): use the zlib-compressed blob only + # if it is actually smaller than the original, otherwise store as-is. + # https://sqlite.org/sqlar.html + compressed = zlib.compress(data) + return compressed if len(compressed) < len(data) else data + + FILE_COLUMNS = { "name": lambda p: p.name, "path": lambda p: str(p), @@ -3696,6 +3724,7 @@ FILE_COLUMNS = { "md5": lambda p: hashlib.md5(p.resolve().read_bytes()).hexdigest(), "mode": lambda p: p.stat().st_mode, "content": lambda p: p.resolve().read_bytes(), + "content_sqlar": lambda p: sqlar_compress(p.resolve().read_bytes()), "mtime": lambda p: p.stat().st_mtime, "ctime": lambda p: p.stat().st_ctime, "mtime_int": lambda p: int(p.stat().st_mtime), diff --git a/tests/test_insert_files.py b/tests/test_insert_files.py index 8894aea..9a5278d 100644 --- a/tests/test_insert_files.py +++ b/tests/test_insert_files.py @@ -1,6 +1,7 @@ import os import pathlib import sys +import zlib import pytest from click.testing import CliRunner @@ -171,6 +172,55 @@ def test_insert_files_fixed_value_column_bad_type(): assert "colname:text:value" in result.output +def test_insert_files_sqlar(): + runner = CliRunner() + with runner.isolated_filesystem(): + tmpdir = pathlib.Path(".") + db_path = str(tmpdir / "files.db") + # Compresses well - lots of repetition + compressible = b"abcdefgh" * 1000 + # Does not compress smaller than the original + incompressible = os.urandom(20) + (tmpdir / "one.txt").write_bytes(compressible) + (tmpdir / "two.bin").write_bytes(incompressible) + result = runner.invoke( + cli.cli, + ["insert-files", db_path, "sqlar", str(tmpdir), "--sqlar"], + catch_exceptions=False, + ) + assert result.exit_code == 0, result.output + db = Database(db_path) + assert db["sqlar"].columns_dict == { + "name": str, + "mode": int, + "mtime": int, + "sz": int, + "data": bytes, + } + assert db["sqlar"].pks == ["name"] + rows_by_name = {r["name"]: r for r in db["sqlar"].rows} + one, two = rows_by_name["one.txt"], rows_by_name["two.bin"] + assert one["sz"] == len(compressible) + assert len(one["data"]) < len(compressible) + assert zlib.decompress(one["data"]) == compressible + assert two["sz"] == len(incompressible) + assert two["data"] == incompressible + + +def test_insert_files_sqlar_and_text_conflict(): + runner = CliRunner() + with runner.isolated_filesystem(): + tmpdir = pathlib.Path(".") + db_path = str(tmpdir / "files.db") + (tmpdir / "one.txt").write_text("hello", "utf-8") + result = runner.invoke( + cli.cli, + ["insert-files", db_path, "files", str(tmpdir), "--text", "--sqlar"], + ) + assert result.exit_code == 1 + assert "Cannot use --text and --sqlar together" in result.output + + @pytest.mark.parametrize( "use_text,encoding,input,expected", (