diff --git a/docs/changelog.rst b/docs/changelog.rst index 5d1816a..9e03f39 100644 --- a/docs/changelog.rst +++ b/docs/changelog.rst @@ -11,6 +11,7 @@ Unreleased - ``sqlite-utils query`` can now read the SQL query from standard input by passing ``-`` in place of the query, for example ``echo "select * from dogs" | sqlite-utils query dogs.db -``. (:issue:`765`) - ``sqlite-utils insert`` and ``sqlite-utils upsert`` now accept a ``--code`` option for :ref:`providing a block of Python code ` (or a path to a ``.py`` file) that defines a ``rows()`` function or ``rows`` iterable of rows to insert, as an alternative to importing from a file. (:issue:`684`) +- ``sqlite-utils insert`` and ``sqlite-utils upsert`` now accept ``--type column-name type`` to :ref:`override the type automatically chosen when the table is created `. This is useful for CSV or TSV columns such as ZIP codes that look like integers but should be stored as ``TEXT`` to preserve leading zeros. (:issue:`131`) .. _v4_0: diff --git a/docs/cli-reference.rst b/docs/cli-reference.rst index 5798809..39226ac 100644 --- a/docs/cli-reference.rst +++ b/docs/cli-reference.rst @@ -246,6 +246,9 @@ See :ref:`cli_inserting_data`, :ref:`cli_insert_csv_tsv`, :ref:`cli_insert_unstr - Use --lines to write each incoming line to a column called "line" - Use --text to write the entire input to a column called "text" + Use --type column-name type to override the type automatically chosen when the + table is created. + You can also use --convert to pass a fragment of Python code that will be used to convert each input. @@ -306,6 +309,7 @@ See :ref:`cli_inserting_data`, :ref:`cli_insert_csv_tsv`, :ref:`cli_insert_unstr --alter Alter existing table to add any missing columns --not-null TEXT Columns that should be created as NOT NULL --default ... Default value that should be set for a column + --type ... Column types to use when creating the table --no-detect-types Treat all CSV/TSV columns as TEXT --analyze Run ANALYZE at the end of this operation --load-extension TEXT Path to SQLite extension, with optional :entrypoint @@ -335,6 +339,9 @@ See :ref:`cli_upsert`. If the table already exists and has a primary key, --pk can be omitted. + Use --type column-name type to override the type automatically chosen when the + table is created. + Example: echo '[ @@ -366,6 +373,7 @@ See :ref:`cli_upsert`. --alter Alter existing table to add any missing columns --not-null TEXT Columns that should be created as NOT NULL --default ... Default value that should be set for a column + --type ... Column types to use when creating the table --no-detect-types Treat all CSV/TSV columns as TEXT --analyze Run ANALYZE at the end of this operation --load-extension TEXT Path to SQLite extension, with optional :entrypoint diff --git a/docs/cli.rst b/docs/cli.rst index 4ed92ed..731e93b 100644 --- a/docs/cli.rst +++ b/docs/cli.rst @@ -1366,6 +1366,25 @@ Will produce this schema with automatically detected types: "weight" REAL ); +.. _cli_insert_csv_tsv_column_types: + +Overriding column types +----------------------- + +Use ``--type column-name type`` to override the type automatically chosen when the table is created. This option can be used more than once, and works with both ``insert`` and ``upsert``: + +.. code-block:: bash + + sqlite-utils insert places.db places places.csv --csv \ + --type zipcode text \ + --type score real + +This is useful for values such as ZIP codes, which may look like integers but should be stored as ``TEXT`` to preserve leading zeros. + +The column type should be one of ``TEXT``, ``INTEGER``, ``FLOAT``, ``REAL`` or ``BLOB``. Column types are matched case-insensitively. + +As with detected column types, ``--type`` only affects tables created by the command. If the table already exists, its existing column types are left unchanged. + To disable type detection and treat all columns as TEXT, use ``--no-detect-types``: .. code-block:: bash diff --git a/sqlite_utils/cli.py b/sqlite_utils/cli.py index 4eefc70..cf39ff8 100644 --- a/sqlite_utils/cli.py +++ b/sqlite_utils/cli.py @@ -980,6 +980,16 @@ def insert_upsert_options(*, require_pk=False): type=(str, str), help="Default value that should be set for a column", ), + click.option( + "--type", + "types", + type=( + str, + click.Choice(list(VALID_COLUMN_TYPES), case_sensitive=False), + ), + multiple=True, + help="Column types to use when creating the table", + ), click.option( "--no-detect-types", is_flag=True, @@ -1034,6 +1044,7 @@ def insert_upsert_implementation( truncate=False, not_null=None, default=None, + types=None, no_detect_types=False, analyze=False, load_extension=None, @@ -1047,6 +1058,7 @@ def insert_upsert_implementation( _register_db_for_cleanup(db) _load_extensions(db, load_extension) _maybe_register_functions(db, functions) + column_type_overrides = {column: ctype.upper() for column, ctype in (types or [])} def _insert_docs(docs, tracker=None): extra_kwargs = { @@ -1060,6 +1072,8 @@ def insert_upsert_implementation( extra_kwargs["not_null"] = set(not_null) if default: extra_kwargs["defaults"] = dict(default) + if column_type_overrides: + extra_kwargs["columns"] = column_type_overrides if upsert: extra_kwargs["upsert"] = upsert @@ -1120,7 +1134,9 @@ def insert_upsert_implementation( and not table_existed_before_insert and db.table(table).exists() ): - db.table(table).transform(types=tracker.types) + detected_types = tracker.types + detected_types.update(column_type_overrides) + db.table(table).transform(types=detected_types) if code is not None: if file is not None: @@ -1330,6 +1346,7 @@ def insert( truncate, not_null, default, + types, strict, ): """ @@ -1348,6 +1365,9 @@ def insert( - Use --lines to write each incoming line to a column called "line" - Use --text to write the entire input to a column called "text" + Use --type column-name type to override the type automatically chosen + when the table is created. + You can also use --convert to pass a fragment of Python code that will be used to convert each input. @@ -1420,6 +1440,7 @@ def insert( silent=silent, not_null=not_null, default=default, + types=types, strict=strict, code=code, ) @@ -1454,6 +1475,7 @@ def upsert( alter, not_null, default, + types, no_detect_types, analyze, load_extension, @@ -1467,6 +1489,9 @@ def upsert( If the table already exists and has a primary key, --pk can be omitted. + Use --type column-name type to override the type automatically chosen + when the table is created. + Example: \b @@ -1501,6 +1526,7 @@ def upsert( upsert=True, not_null=not_null, default=default, + types=types, no_detect_types=no_detect_types, analyze=analyze, load_extension=load_extension, diff --git a/tests/test_cli_insert.py b/tests/test_cli_insert.py index e290f97..df6f80c 100644 --- a/tests/test_cli_insert.py +++ b/tests/test_cli_insert.py @@ -664,6 +664,53 @@ def test_insert_csv_detect_types_new_table(db_path): assert db["data"].columns_dict == {"name": str, "age": int, "weight": float} +@pytest.mark.parametrize( + "command,extra_args,input_text,expected_row", + ( + ( + "insert", + [], + "zipcode,score\n01234,9.5\n", + {"zipcode": "01234", "score": 9.5}, + ), + ( + "upsert", + ["--pk", "id"], + "id,zipcode,score\n1,01234,9.5\n", + {"id": 1, "zipcode": "01234", "score": 9.5}, + ), + ), +) +def test_insert_upsert_csv_type_overrides_detected_types( + db_path, command, extra_args, input_text, expected_row +): + result = CliRunner().invoke( + cli.cli, + [ + command, + db_path, + "places", + "-", + "--csv", + ] + + extra_args + + [ + "--type", + "zipcode", + "text", + ], + catch_exceptions=False, + input=input_text, + ) + assert result.exit_code == 0, result.output + db = Database(db_path) + expected_columns = {"zipcode": str, "score": float} + if command == "upsert": + expected_columns = {"id": int, **expected_columns} + assert db["places"].columns_dict == expected_columns + assert list(db["places"].rows) == [expected_row] + + def test_upsert_csv_detect_types_leaves_existing_table_alone(db_path): db = Database(db_path) db["places"].insert({"id": 1, "name": "Boston", "zip": "01234"}, pk="id")