diff --git a/docs/cli.rst b/docs/cli.rst index f3854ba..2d439f6 100644 --- a/docs/cli.rst +++ b/docs/cli.rst @@ -354,6 +354,10 @@ For tab-delimited data, use ``--tsv``:: $ sqlite-utils insert dogs.db dogs docs.tsv --tsv +Data is expected to be encoded as Unicode UTF-8. If your data is an another character encoding you can specify it using the ``--encoding`` option:: + + $ sqlite-utils insert dogs.db dogs docs.tsv --tsv --encoding=latin-1 + .. _cli_insert_replace: Insert-replacing data diff --git a/sqlite_utils/cli.py b/sqlite_utils/cli.py index 56dc0c8..cfa7303 100644 --- a/sqlite_utils/cli.py +++ b/sqlite_utils/cli.py @@ -1,5 +1,6 @@ import base64 import click +import codecs from click_default_group import DefaultGroup from datetime import datetime import hashlib @@ -15,6 +16,18 @@ from .utils import sqlite3, decode_base64_values VALID_COLUMN_TYPES = ("INTEGER", "TEXT", "FLOAT", "BLOB") +UNICODE_ERROR = """ +{} + +The input you provided uses a character encoding other than utf-8. + +You can fix this by passing the --encoding= option with the encoding of the file. + +If you do not know the encoding, running 'file filename.csv' may tell you. + +It's often worth trying: --encoding=latin-1 +""".strip() + def output_options(fn): for decorator in reversed( @@ -493,7 +506,7 @@ def insert_upsert_options(fn): required=True, ), click.argument("table"), - click.argument("json_file", type=click.File(), required=True), + click.argument("json_file", type=click.File("rb"), required=True), click.option( "--pk", help="Columns to use as the primary key, e.g. id", multiple=True ), @@ -519,6 +532,10 @@ def insert_upsert_options(fn): type=(str, str), help="Default value that should be set for a column", ), + click.option( + "--encoding", + help="Character encoding for input, defaults to utf-8", + ), ) ): fn = decorator(fn) @@ -541,10 +558,15 @@ def insert_upsert_implementation( truncate=False, not_null=None, default=None, + encoding=None, ): db = sqlite_utils.Database(path) if (nl + csv + tsv) >= 2: raise click.ClickException("Use just one of --nl, --csv or --tsv") + if encoding and not (csv or tsv): + raise click.ClickException("--encoding must be used with --csv or --tsv") + encoding = encoding or "utf-8" + json_file = codecs.getreader(encoding)(json_file) if pk and len(pk) == 1: pk = pk[0] if csv or tsv: @@ -599,6 +621,7 @@ def insert( tsv, batch_size, alter, + encoding, ignore, replace, truncate, @@ -611,49 +634,68 @@ def insert( Input should be a JSON array of objects, unless --nl or --csv is used. """ - insert_upsert_implementation( - path, - table, - json_file, - pk, - nl, - csv, - tsv, - batch_size, - alter=alter, - upsert=False, - ignore=ignore, - replace=replace, - truncate=truncate, - not_null=not_null, - default=default, - ) + try: + insert_upsert_implementation( + path, + table, + json_file, + pk, + nl, + csv, + tsv, + batch_size, + alter=alter, + upsert=False, + ignore=ignore, + replace=replace, + truncate=truncate, + encoding=encoding, + not_null=not_null, + default=default, + ) + except UnicodeDecodeError as ex: + raise click.ClickException(UNICODE_ERROR.format(ex)) @cli.command() @insert_upsert_options def upsert( - path, table, json_file, pk, nl, csv, tsv, batch_size, alter, not_null, default + path, + table, + json_file, + pk, + nl, + csv, + tsv, + batch_size, + alter, + not_null, + default, + encoding, ): """ Upsert records based on their primary key. Works like 'insert' but if an incoming record has a primary key that matches an existing record the existing record will be updated. """ - insert_upsert_implementation( - path, - table, - json_file, - pk, - nl, - csv, - tsv, - batch_size, - alter=alter, - upsert=True, - not_null=not_null, - default=default, - ) + try: + insert_upsert_implementation( + path, + table, + json_file, + pk, + nl, + csv, + tsv, + batch_size, + alter=alter, + upsert=True, + not_null=not_null, + default=default, + encoding=encoding, + ) + except UnicodeDecodeError as ex: + raise click.ClickException(UNICODE_ERROR.format(ex)) @cli.command(name="create-table") diff --git a/tests/test_cli.py b/tests/test_cli.py index c50f87f..d70ff6f 100644 --- a/tests/test_cli.py +++ b/tests/test_cli.py @@ -685,7 +685,9 @@ def test_insert_csv_tsv(content, option, db_path, tmpdir): db = Database(db_path) file_path = str(tmpdir / "insert.csv-tsv") open(file_path, "w").write(content) - result = CliRunner().invoke(cli.cli, ["insert", db_path, "data", file_path, option]) + result = CliRunner().invoke( + cli.cli, ["insert", db_path, "data", file_path, option], catch_exceptions=False + ) assert 0 == result.exit_code assert [{"foo": "1", "bar": "2", "baz": "3"}] == list(db["data"].rows) @@ -1024,7 +1026,9 @@ def test_upsert(db_path, tmpdir): ] open(json_path, "w").write(json.dumps(insert_dogs)) result = CliRunner().invoke( - cli.cli, ["insert", db_path, "dogs", json_path, "--pk", "id"] + cli.cli, + ["insert", db_path, "dogs", json_path, "--pk", "id"], + catch_exceptions=False, ) assert 0 == result.exit_code, result.output assert 2 == db["dogs"].count @@ -1035,7 +1039,9 @@ def test_upsert(db_path, tmpdir): ] open(json_path, "w").write(json.dumps(insert_dogs)) result = CliRunner().invoke( - cli.cli, ["upsert", db_path, "dogs", json_path, "--pk", "id"] + cli.cli, + ["upsert", db_path, "dogs", json_path, "--pk", "id"], + catch_exceptions=False, ) assert 0 == result.exit_code, result.output assert [ @@ -1601,3 +1607,52 @@ def test_extract(db_path, args, expected_table_schema, expected_other_schema): 0 ].schema assert other_schema == expected_other_schema + + +def test_insert_encoding(tmpdir): + db_path = str(tmpdir / "test.db") + latin1_csv = ( + b"date,name,latitude,longitude\n" + b"2020-01-01,Barra da Lagoa,-27.574,-48.422\n" + b"2020-03-04,S\xe3o Paulo,-23.561,-46.645\n" + b"2020-04-05,Salta,-24.793:-65.408" + ) + assert latin1_csv.decode("latin-1").split("\n")[2].split(",")[1] == "São Paulo" + csv_path = str(tmpdir / "test.csv") + open(csv_path, "wb").write(latin1_csv) + # First attempt should error: + bad_result = CliRunner().invoke( + cli.cli, ["insert", db_path, "places", csv_path, "--csv"] + ) + assert bad_result.exit_code == 1 + assert ( + "The input you provided uses a character encoding other than utf-8" + in bad_result.output + ) + # Using --encoding=latin-1 should work + good_result = CliRunner().invoke( + cli.cli, + ["insert", db_path, "places", csv_path, "--encoding", "latin-1", "--csv"], + ) + assert good_result.exit_code == 0 + db = Database(db_path) + assert list(db["places"].rows) == [ + { + "date": "2020-01-01", + "name": "Barra da Lagoa", + "latitude": "-27.574", + "longitude": "-48.422", + }, + { + "date": "2020-03-04", + "name": "São Paulo", + "latitude": "-23.561", + "longitude": "-46.645", + }, + { + "date": "2020-04-05", + "name": "Salta", + "latitude": "-24.793:-65.408", + "longitude": None, + }, + ]