Add --extract to CLI insert and upsert

tests/test_cli_extract.py

@@ -0,0 +1,243 @@
+import csv
+import io
+import json
+
+import pytest
+from click.testing import CliRunner
+
+from sqlite_utils import Database, cli
+
+
+@pytest.mark.parametrize("command", ["insert", "upsert"])
+@pytest.mark.parametrize("format", ["json", "nl", "csv", "tsv"])
+@pytest.mark.parametrize("batch_size", [1, 100])
+def test_cli_extract(tmp_path, command, format, batch_size):
+    path = str(tmp_path / "trees.db")
+    rows = [
+        {"id": 1, "species": "oak", "region": "west", "height": 1.5},
+        {"id": 2, "species": "ash", "region": "west", "height": 2.5},
+        {"id": 3, "species": "oak", "region": "east", "height": 3.5},
+    ]
+    if format == "json":
+        input = json.dumps(rows)
+    elif format == "nl":
+        input = "\n".join(json.dumps(row) for row in rows)
+    else:
+        output = io.StringIO()
+        writer = csv.DictWriter(
+            output,
+            fieldnames=list(rows[0]),
+            delimiter="\t" if format == "tsv" else ",",
+        )
+        writer.writeheader()
+        writer.writerows(rows)
+        input = output.getvalue()
+    args = [
+        command,
+        path,
+        "trees",
+        "-",
+        "--pk",
+        "id",
+        "--extract",
+        "species",
+        "--extract",
+        "region",
+        "--batch-size",
+        str(batch_size),
+    ]
+    if format != "json":
+        args.append("--" + format)
+    result = CliRunner().invoke(cli.cli, args, input=input)
+    assert result.exit_code == 0, result.output
+    db = Database(path)
+    assert list(db["trees"].rows) == [
+        {"id": 1, "species": 1, "region": 1, "height": 1.5},
+        {"id": 2, "species": 2, "region": 1, "height": 2.5},
+        {"id": 3, "species": 1, "region": 2, "height": 3.5},
+    ]
+    assert list(db["species"].rows) == [
+        {"id": 1, "value": "oak"},
+        {"id": 2, "value": "ash"},
+    ]
+    assert list(db["region"].rows) == [
+        {"id": 1, "value": "west"},
+        {"id": 2, "value": "east"},
+    ]
+    assert db["trees"].columns_dict == {
+        "id": int,
+        "species": int,
+        "region": int,
+        "height": float,
+    }
+    assert {
+        (fk.column, fk.other_table, fk.other_column) for fk in db["trees"].foreign_keys
+    } == {
+        ("species", "species", "id"),
+        ("region", "region", "id"),
+    }
+    assert list(db.query("pragma foreign_key_check")) == []
+
+
+@pytest.mark.parametrize("command", ["insert", "upsert"])
+def test_cli_extract_code(tmp_path, command):
+    path = str(tmp_path / "code.db")
+    result = CliRunner().invoke(
+        cli.cli,
+        [
+            command,
+            path,
+            "trees",
+            "--code",
+            'rows = [{"id": 1, "species": "oak"}]',
+            "--pk",
+            "id",
+            "--extract",
+            "species",
+        ],
+    )
+    assert result.exit_code == 0, result.output
+    db = Database(path)
+    assert list(db["trees"].rows) == [{"id": 1, "species": 1}]
+    assert list(db["species"].rows) == [{"id": 1, "value": "oak"}]
+
+
+@pytest.mark.parametrize("command", ["insert", "upsert"])
+@pytest.mark.parametrize(
+    "extra_args",
+    [[], ["--strict"], ["--no-detect-types"], ["--type", "species", "TEXT"]],
+)
+def test_cli_extract_csv_types(tmp_path, command, extra_args):
+    path = str(tmp_path / "types.db")
+    result = CliRunner().invoke(
+        cli.cli,
+        [command, path, "trees", "-", "--pk", "id", "--csv", "--extract", "species"]
+        + extra_args,
+        input="id,species,height\n1,oak,1.5\n",
+    )
+    assert result.exit_code == 0, result.output
+    db = Database(path)
+    assert db["trees"].columns_dict["species"] is int
+    assert list(db["trees"].rows)[0]["species"] == 1
+    assert list(db.query("pragma foreign_key_check")) == []
+
+
+@pytest.mark.parametrize("command", ["insert", "upsert"])
+def test_cli_extract_nulls(tmp_path, command):
+    path = str(tmp_path / "nulls.db")
+    result = CliRunner().invoke(
+        cli.cli,
+        [command, path, "trees", "-", "--pk", "id", "--extract", "species"],
+        input=json.dumps(
+            [
+                {"id": 1, "species": None},
+                {"id": 2, "species": "oak"},
+                {"id": 3, "species": None},
+            ]
+        ),
+    )
+    assert result.exit_code == 0, result.output
+    db = Database(path)
+    assert list(db["species"].rows) == [{"id": 1, "value": "oak"}]
+    assert list(db["trees"].rows) == [
+        {"id": 1, "species": None},
+        {"id": 2, "species": 1},
+        {"id": 3, "species": None},
+    ]
+
+
+def test_cli_extract_upsert(tmp_path):
+    path = str(tmp_path / "upsert.db")
+    for command, rows in [
+        ("insert", [{"id": 1, "species": "oak"}, {"id": 2, "species": "ash"}]),
+        ("upsert", [{"id": 1, "species": "ash"}, {"id": 3, "species": "oak"}]),
+    ]:
+        result = CliRunner().invoke(
+            cli.cli,
+            [command, path, "trees", "-", "--pk", "id", "--extract", "species"],
+            input=json.dumps(rows),
+        )
+        assert result.exit_code == 0, result.output
+    db = Database(path)
+    assert list(db["species"].rows) == [
+        {"id": 1, "value": "oak"},
+        {"id": 2, "value": "ash"},
+    ]
+    assert list(db["trees"].rows) == [
+        {"id": 1, "species": 2},
+        {"id": 2, "species": 2},
+        {"id": 3, "species": 1},
+    ]
+
+
+@pytest.mark.parametrize("column", ['tree "species"', "品種", ""])
+def test_cli_extract_identifiers(tmp_path, column):
+    path = str(tmp_path / "identifiers.db")
+    result = CliRunner().invoke(
+        cli.cli,
+        ["insert", path, "trees", "-", "--extract", column],
+        input=json.dumps({column: "oak"}),
+    )
+    assert result.exit_code == 0, result.output
+    db = Database(path)
+    assert list(db["trees"].rows) == [{column: 1}]
+    assert list(db[column].rows) == [{"id": 1, "value": "oak"}]
+
+
+@pytest.mark.parametrize("command", ["insert", "upsert"])
+def test_cli_extract_errors(tmp_path, command):
+    path = str(tmp_path / "invalid.db")
+    result = CliRunner().invoke(
+        cli.cli,
+        [command, path, "trees", "-", "--pk", "id", "--extract", "species"],
+        input="invalid JSON",
+    )
+    assert result.exit_code == 1
+    assert "Invalid JSON" in result.output
+    assert Database(path).table_names() == []
+
+
+@pytest.mark.parametrize("command", ["insert", "upsert"])
+def test_cli_extract_requires_column(tmp_path, command):
+    path = tmp_path / "missing.db"
+    result = CliRunner().invoke(
+        cli.cli, [command, str(path), "trees", "-", "--extract"], input="{}"
+    )
+    assert result.exit_code == 2
+    assert "requires an argument" in result.output
+    assert not path.exists()
+
+
+def test_cli_extract_existing_schema(tmp_path):
+    path = str(tmp_path / "existing.db")
+    db = Database(path)
+    db["trees"].create({"id": int, "species": int}, pk="id")
+    schema = db["trees"].schema
+    result = CliRunner().invoke(
+        cli.cli,
+        ["insert", path, "trees", "-", "--extract", "species"],
+        input='{"id": 1, "species": "oak"}',
+    )
+    assert result.exit_code == 0, result.output
+    assert db["trees"].schema == schema
+    assert db["trees"].foreign_keys == []
+    assert list(db["trees"].rows) == [{"id": 1, "species": 1}]
+    assert list(db["species"].rows) == [{"id": 1, "value": "oak"}]
+
+
+@pytest.mark.parametrize("command", ["insert", "upsert"])
+def test_cli_no_extract_default(tmp_path, command):
+    path = str(tmp_path / "default.db")
+    result = CliRunner().invoke(
+        cli.cli,
+        [command, path, "trees", "-", "--pk", "id"],
+        input='[{"id": 1, "species": "oak"}, {"id": 2, "species": "oak"}]',
+    )
+    assert result.exit_code == 0, result.output
+    db = Database(path)
+    assert db.table_names() == ["trees"]
+    assert list(db["trees"].rows) == [
+        {"id": 1, "species": "oak"},
+        {"id": 2, "species": "oak"},
+    ]
+    assert db["trees"].foreign_keys == []