diff --git a/docs/cli.rst b/docs/cli.rst index 49d5d92..8fdbbf1 100644 --- a/docs/cli.rst +++ b/docs/cli.rst @@ -232,6 +232,15 @@ You can add a column using the ``add-column`` command:: The last argument here is the type of the column to be created. You can use one of ``text``, ``integer``, ``float`` or ``blob``. If you leave it off, ``text`` will be used. +.. _cli_add_column_alter: + +Adding columns automatically on insert/update +============================================= + +You can use the ``--alter`` option to automatically add new columns if the data you are inserting or upserting is of a different shape:: + + $ sqlite-utils insert dogs.db dogs new-dogs.json --pk=id --alter + .. _cli_add_foreign_key: Adding foreign key constraints diff --git a/docs/python-api.rst b/docs/python-api.rst index 7433c5a..769189d 100644 --- a/docs/python-api.rst +++ b/docs/python-api.rst @@ -266,6 +266,25 @@ If you pass a Python type, it will be mapped to SQLite types as shown here:: np.float32: "FLOAT" np.float64: "FLOAT" +.. _python_api_add_column_alter: + +Adding columns automatically on insert/update +============================================= + +You can insert or update data that includes new columns and have the table automatically altered to fit the new schema using the ``alter=True`` argument. This can be passed to all four of ``.insert()``, ``.upsert()``, ``.insert_all()``and ``.insert_all()``: + +.. code-block:: python + + db["new_table"].insert({"name": "Gareth"}) + # This will throw an exception: + db["new_table"].insert({"name": "Gareth", "age": 32}) + # This will succeed and add a new "age" integer column: + db["new_table"].insert({"name": "Gareth", "age": 32}, alter=True) + # You can see confirm the new column like so: + print(db["new_table"].columns_dict) + # Outputs this: + # {'name': , 'age': } + .. _python_api_add_foreign_key: Adding foreign key constraints @@ -382,6 +401,11 @@ The ``.columns`` property shows the columns in the table:: [Column(cid=0, name='id', type='INTEGER', notnull=0, default_value=None, is_pk=1), Column(cid=1, name='value', type='TEXT', notnull=0, default_value=None, is_pk=0)] +The ``.columns_dict`` property returns a dictionary version of this with just the names and types:: + + >>> db["PlantType"].columns_dict + {'id': , 'value': } + The ``.foreign_keys`` property shows if the table has any foreign key relationships:: >>> db["Street_Tree_List"].foreign_keys diff --git a/sqlite_utils/cli.py b/sqlite_utils/cli.py index 63fcffb..e064716 100644 --- a/sqlite_utils/cli.py +++ b/sqlite_utils/cli.py @@ -262,6 +262,11 @@ def insert_upsert_options(fn): click.option( "--batch-size", type=int, default=100, help="Commit every X records" ), + click.option( + "--alter", + is_flag=True, + help="Alter existing table to add any missing columns", + ), ) ): fn = decorator(fn) @@ -269,7 +274,7 @@ def insert_upsert_options(fn): def insert_upsert_implementation( - path, table, json_file, pk, nl, csv, batch_size, upsert + path, table, json_file, pk, nl, csv, batch_size, alter, upsert ): db = sqlite_utils.Database(path) if nl and csv: @@ -289,12 +294,12 @@ def insert_upsert_implementation( method = db[table].upsert_all else: method = db[table].insert_all - method(docs, pk=pk, batch_size=batch_size) + method(docs, pk=pk, batch_size=batch_size, alter=alter) @cli.command() @insert_upsert_options -def insert(path, table, json_file, pk, nl, csv, batch_size): +def insert(path, table, json_file, pk, nl, csv, batch_size, alter): """ Insert records from JSON file into a table, creating the table if it does not already exist. @@ -302,20 +307,20 @@ def insert(path, table, json_file, pk, nl, csv, batch_size): Input should be a JSON array of objects, unless --nl or --csv is used. """ insert_upsert_implementation( - path, table, json_file, pk, nl, csv, batch_size, upsert=False + path, table, json_file, pk, nl, csv, batch_size, alter=alter, upsert=False ) @cli.command() @insert_upsert_options -def upsert(path, table, json_file, pk, nl, csv, batch_size): +def upsert(path, table, json_file, pk, nl, csv, batch_size, alter): """ 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 replaced. """ insert_upsert_implementation( - path, table, json_file, pk, nl, csv, batch_size, upsert=True + path, table, json_file, pk, nl, csv, batch_size, alter=alter, upsert=True ) diff --git a/sqlite_utils/db.py b/sqlite_utils/db.py index 9b9170f..0e8ee6b 100644 --- a/sqlite_utils/db.py +++ b/sqlite_utils/db.py @@ -59,6 +59,14 @@ if np: ) +REVERSE_COLUMN_TYPE_MAPPING = { + "TEXT": str, + "BLOB": bytes, + "INTEGER": int, + "FLOAT": float, +} + + class AlterError(Exception): pass @@ -182,6 +190,14 @@ class Table: ).fetchall() return [Column(*row) for row in rows] + @property + def columns_dict(self): + "Returns {column: python-type} dictionary" + return { + column.name: REVERSE_COLUMN_TYPE_MAPPING[column.type] + for column in self.columns + } + @property def rows(self): if not self.exists: @@ -235,7 +251,7 @@ class Table: for seqno, cid, name in self.db.conn.execute(column_sql).fetchall(): columns.append(name) row["columns"] = columns - # These coluns may be missing on older SQLite versions: + # These columns may be missing on older SQLite versions: for key, default in {"origin": "c", "partial": 0}.items(): if key not in row: row[key] = default @@ -429,6 +445,7 @@ class Table: upsert=False, column_order=None, hash_id=None, + alter=False, ): return self.insert_all( [record], @@ -437,6 +454,7 @@ class Table: upsert=upsert, column_order=column_order, hash_id=hash_id, + alter=alter, ) def insert_all( @@ -448,6 +466,7 @@ class Table: batch_size=100, column_order=None, hash_id=None, + alter=False, ): """ Like .insert() but takes a list of records and ensures that the table @@ -500,7 +519,15 @@ class Table: for key in all_columns ) with self.db.conn: - result = self.db.conn.execute(sql, values) + try: + result = self.db.conn.execute(sql, values) + except sqlite3.OperationalError as e: + if alter and (" has no column " in e.args[0]): + # Attempt to add any missing columns, then try again + self.add_missing_columns(chunk) + result = self.db.conn.execute(sql, values) + else: + raise self.last_rowid = result.lastrowid self.last_pk = None if hash_id or pk: @@ -513,7 +540,13 @@ class Table: return self def upsert( - self, record, pk=None, foreign_keys=None, column_order=None, hash_id=None + self, + record, + pk=None, + foreign_keys=None, + column_order=None, + hash_id=None, + alter=False, ): return self.insert( record, @@ -522,6 +555,7 @@ class Table: upsert=True, column_order=column_order, hash_id=hash_id, + alter=alter, ) def upsert_all( @@ -532,6 +566,7 @@ class Table: column_order=None, batch_size=100, hash_id=None, + alter=False, ): return self.insert_all( records, @@ -541,8 +576,16 @@ class Table: batch_size=100, upsert=True, hash_id=hash_id, + alter=alter, ) + def add_missing_columns(self, records): + needed_columns = self.detect_column_types(records) + current_columns = self.columns_dict + for col_name, col_type in needed_columns.items(): + if col_name not in current_columns: + self.add_column(col_name, col_type) + def chunks(sequence, size): iterator = iter(sequence) diff --git a/tests/test_cli.py b/tests/test_cli.py index ad2bec7..92bfcf3 100644 --- a/tests/test_cli.py +++ b/tests/test_cli.py @@ -383,6 +383,37 @@ def test_upsert(db_path, tmpdir): ) +def test_insert_alter(db_path, tmpdir): + result = CliRunner().invoke( + cli.cli, + ["insert", db_path, "from_json_nl", "-", "--nl"], + input='{"foo": "bar", "n": 1}\n{"foo": "baz", "n": 2}', + ) + assert 0 == result.exit_code, result.output + # Should get an error with incorrect shaped additional data + result = CliRunner().invoke( + cli.cli, + ["insert", db_path, "from_json_nl", "-", "--nl"], + input='{"foo": "bar", "baz": 5}', + ) + assert 0 != result.exit_code, result.output + # If we run it again with --alter it should work correctly + result = CliRunner().invoke( + cli.cli, + ["insert", db_path, "from_json_nl", "-", "--nl", "--alter"], + input='{"foo": "bar", "baz": 5}', + ) + assert 0 == result.exit_code, result.output + # Sanity check the database itself + db = Database(db_path) + assert {"foo": str, "n": int, "baz": int} == db["from_json_nl"].columns_dict + assert [ + {"foo": "bar", "n": 1, "baz": None}, + {"foo": "baz", "n": 2, "baz": None}, + {"foo": "bar", "baz": 5, "n": None}, + ] == db.execute_returning_dicts("select foo, n, baz from from_json_nl") + + def test_query_csv(db_path): db = Database(db_path) with db.conn: diff --git a/tests/test_create.py b/tests/test_create.py index 24f4e97..2cb2f6b 100644 --- a/tests/test_create.py +++ b/tests/test_create.py @@ -195,6 +195,89 @@ def test_add_foreign_key_error_if_already_exists(fresh_db): assert "Foreign key already exists for author_id => authors.id" == ex.value.args[0] +@pytest.mark.parametrize( + "extra_data,expected_new_columns", + [ + ({"species": "squirrels"}, [{"name": "species", "type": "TEXT"}]), + ( + {"species": "squirrels", "hats": 5}, + [{"name": "species", "type": "TEXT"}, {"name": "hats", "type": "INTEGER"}], + ), + ( + {"hats": 5, "rating": 3.5}, + [{"name": "hats", "type": "INTEGER"}, {"name": "rating", "type": "FLOAT"}], + ), + ], +) +def test_insert_row_alter_table(fresh_db, extra_data, expected_new_columns): + table = fresh_db["books"] + table.insert({"title": "Hedgehogs of the world", "author_id": 1}) + assert [ + {"name": "title", "type": "TEXT"}, + {"name": "author_id", "type": "INTEGER"}, + ] == [{"name": col.name, "type": col.type} for col in table.columns] + record = {"title": "Squirrels of the world", "author_id": 2} + record.update(extra_data) + fresh_db["books"].insert(record, alter=True) + assert [ + {"name": "title", "type": "TEXT"}, + {"name": "author_id", "type": "INTEGER"}, + ] + expected_new_columns == [ + {"name": col.name, "type": col.type} for col in table.columns + ] + + +def test_upsert_rows_alter_table(fresh_db): + table = fresh_db["books"] + table.insert({"id": 1, "title": "Hedgehogs of the world", "author_id": 1}, pk="id") + table.upsert_all( + [ + {"id": 1, "title": "Hedgedogs of the World", "species": "hedgehogs"}, + {"id": 2, "title": "Squirrels of the World", "num_species": 200}, + { + "id": 3, + "title": "Badgers of the World", + "significant_continents": ["Europe", "North America"], + }, + ], + alter=True, + ) + assert { + "author_id": int, + "id": int, + "num_species": int, + "significant_continents": str, + "species": str, + "title": str, + } == table.columns_dict + assert [ + { + "author_id": None, + "id": 1, + "num_species": None, + "significant_continents": None, + "species": "hedgehogs", + "title": "Hedgedogs of the World", + }, + { + "author_id": None, + "id": 2, + "num_species": 200, + "significant_continents": None, + "species": None, + "title": "Squirrels of the World", + }, + { + "author_id": None, + "id": 3, + "num_species": None, + "significant_continents": '["Europe", "North America"]', + "species": None, + "title": "Badgers of the World", + }, + ] == list(table.rows) + + @pytest.mark.parametrize( "columns,index_name,expected_index", (