From 2ab62bcd54f8aa73a7a77d209133f0da6c73f3ea Mon Sep 17 00:00:00 2001 From: Simon Willison Date: Fri, 17 Apr 2020 16:53:25 -0700 Subject: [PATCH] New columns= parameter for over-riding column types, closes #100 --- docs/python-api.rst | 32 ++++++++++++++++++++++++++++++-- sqlite_utils/db.py | 14 +++++++++++++- tests/test_create.py | 23 +++++++++++++++++++++++ 3 files changed, 66 insertions(+), 3 deletions(-) diff --git a/docs/python-api.rst b/docs/python-api.rst index 8a2ad23..e3c40e6 100644 --- a/docs/python-api.rst +++ b/docs/python-api.rst @@ -157,6 +157,8 @@ If the record does not exist a ``NotFoundError`` will be raised: except NotFoundError: print("Dog not found") +.. _python_api_creating_tables: + Creating tables =============== @@ -197,6 +199,15 @@ You can also specify a primary key by passing the ``pk=`` parameter to the ``.in "is_good_dog": True, }, pk="id") +After inserting a row like this, the ``dogs.last_rowid`` property will return the SQLite ``rowid`` assigned to the most recently inserted record. + +The ``dogs.last_pk`` property will return the last inserted primary key value, if you specified one. This can be very useful when writing code that creates foreign keys or many-to-many relationships. + +.. _python_api_custom_columns: + +Custom column order and column types +------------------------------------ + The order of the columns in the table will be derived from the order of the keys in the dictionary, provided you are using Python 3.6 or later. If you want to explicitly set the order of the columns you can do so using the ``column_order=`` parameter: @@ -213,9 +224,26 @@ If you want to explicitly set the order of the columns you can do so using the ` You don't need to pass all of the columns to the ``column_order`` parameter. If you only pass a subset of the columns the remaining columns will be ordered based on the key order of the dictionary. -After inserting a row like this, the ``dogs.last_rowid`` property will return the SQLite ``rowid`` assigned to the most recently inserted record. +Column types are detected based on the example data provided. Sometimes you may find you need to over-ride these detected types - to create an integer column for data that was provided as a string for example, or to ensure that a table where the first example was ``None`` is created as an ``INTEGER`` rather than a ``TEXT`` column. You can do this using the ``columns=`` parameter: -The ``dogs.last_pk`` property will return the last inserted primary key value, if you specified one. This can be very useful when writing code that creates foreign keys or many-to-many relationships. +.. code-block:: python + + dogs.insert({ + "id": 1, + "name": "Cleo", + "age": "5", + }, pk="id", columns={"age": int, "weight": float}) + +This will create a table with the following schema: + +.. code-block:: sql + + CREATE TABLE [dogs] ( + [id] INTEGER PRIMARY KEY, + [name] TEXT, + [age] INTEGER, + [weight] FLOAT + ) .. _python_api_explicit_create: diff --git a/sqlite_utils/db.py b/sqlite_utils/db.py index d92e92e..9e3987e 100644 --- a/sqlite_utils/db.py +++ b/sqlite_utils/db.py @@ -491,6 +491,7 @@ class Table(Queryable): replace=False, extracts=None, conversions=None, + columns=None, ): super().__init__(db, name) self._defaults = dict( @@ -506,6 +507,7 @@ class Table(Queryable): replace=replace, extracts=extracts, conversions=conversions or {}, + columns=columns, ) def __repr__(self): @@ -932,6 +934,7 @@ class Table(Queryable): replace=DEFAULT, extracts=DEFAULT, conversions=DEFAULT, + columns=DEFAULT, ): return self.insert_all( [record], @@ -946,6 +949,7 @@ class Table(Queryable): replace=replace, extracts=extracts, conversions=conversions, + columns=columns, ) def insert_all( @@ -963,6 +967,7 @@ class Table(Queryable): replace=DEFAULT, extracts=DEFAULT, conversions=DEFAULT, + columns=DEFAULT, upsert=False, ): """ @@ -982,6 +987,7 @@ class Table(Queryable): replace = self.value_or_default("replace", replace) extracts = self.value_or_default("extracts", extracts) conversions = self.value_or_default("conversions", conversions) + columns = self.value_or_default("columns", columns) if upsert and (not pk and not hash_id): raise PrimaryKeyRequired("upsert() requires a pk") @@ -1016,8 +1022,10 @@ class Table(Queryable): if first: if not self.exists(): # Use the first batch to derive the table names + column_types = suggest_column_types(chunk) + column_types.update(columns or {}) self.create( - suggest_column_types(chunk), + column_types, pk, foreign_keys, column_order=column_order, @@ -1154,6 +1162,7 @@ class Table(Queryable): alter=DEFAULT, extracts=DEFAULT, conversions=DEFAULT, + columns=DEFAULT, ): return self.upsert_all( [record], @@ -1166,6 +1175,7 @@ class Table(Queryable): alter=alter, extracts=extracts, conversions=conversions, + columns=columns, ) def upsert_all( @@ -1181,6 +1191,7 @@ class Table(Queryable): alter=DEFAULT, extracts=DEFAULT, conversions=DEFAULT, + columns=DEFAULT, ): return self.insert_all( records, @@ -1194,6 +1205,7 @@ class Table(Queryable): alter=alter, extracts=extracts, conversions=conversions, + columns=columns, upsert=True, ) diff --git a/tests/test_create.py b/tests/test_create.py index 7edcd8d..95b64ab 100644 --- a/tests/test_create.py +++ b/tests/test_create.py @@ -154,6 +154,29 @@ def test_create_table_from_example_with_compound_primary_keys(fresh_db): assert record == table.get(("staff", 2)) +@pytest.mark.parametrize( + "method_name", ("insert", "upsert", "insert_all", "upsert_all") +) +def test_create_table_with_custom_columns(fresh_db, method_name): + table = fresh_db["dogs"] + method = getattr(table, method_name) + record = {"id": 1, "name": "Cleo", "age": "5"} + if method_name.endswith("_all"): + record = [record] + method(record, pk="id", columns={"age": int, "weight": float}) + assert ["dogs"] == fresh_db.table_names() + expected_columns = [ + {"name": "id", "type": "INTEGER"}, + {"name": "name", "type": "TEXT"}, + {"name": "age", "type": "INTEGER"}, + {"name": "weight", "type": "FLOAT"}, + ] + assert expected_columns == [ + {"name": col.name, "type": col.type} for col in table.columns + ] + assert [{"id": 1, "name": "Cleo", "age": 5, "weight": None}] == list(table.rows) + + @pytest.mark.parametrize("use_table_factory", [True, False]) def test_create_table_column_order(fresh_db, use_table_factory): row = collections.OrderedDict(