sqlite-utils/tests/test_cli_extract.py

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 == []