From 941d281aee6eac20ad64b505511da7e47f697700 Mon Sep 17 00:00:00 2001 From: Simon Willison Date: Tue, 23 Jul 2019 10:00:42 -0700 Subject: [PATCH] extracts= table parameter, closes #46 --- docs/index.rst | 2 +- docs/python-api.rst | 42 ++++++++++++++++++++++++++-- sqlite_utils/db.py | 46 +++++++++++++++++++++++++++--- tests/test_extracts.py | 63 ++++++++++++++++++++++++++++++++++++++++++ 4 files changed, 146 insertions(+), 7 deletions(-) create mode 100644 tests/test_extracts.py diff --git a/docs/index.rst b/docs/index.rst index f50d76e..f05532f 100644 --- a/docs/index.rst +++ b/docs/index.rst @@ -16,7 +16,7 @@ Contents -------- .. toctree:: - :maxdepth: 2 + :maxdepth: 3 cli python-api diff --git a/docs/python-api.rst b/docs/python-api.rst index 0bbf19d..362fee0 100644 --- a/docs/python-api.rst +++ b/docs/python-api.rst @@ -356,11 +356,16 @@ An ``upsert_all()`` method is also available, which behaves like ``insert_all()` .. _python_api_lookup_tables: -Creating lookup tables -====================== +Working with lookup tables +========================== A useful pattern when populating large tables in to break common values out into lookup tables. Consider a table of ``Trees``, where each tree has a species. Ideally these species would be split out into a separate ``Species`` table, with each one assigned an integer primary key that can be referenced from the ``Trees`` table ``species_id`` column. +.. _python_api_explicit_lookup_tables: + +Creating lookup tables explicitly +--------------------------------- + Calling ``db["Species"].lookup({"name": "Palm"})`` creates a table called ``Species`` (if one does not already exist) with two columns: ``id`` and ``name``. It sets up a unique constraint on the ``name`` column to guarantee it will not contain duplicate rows. It then inserts a new row with the ``name`` set to ``Palm`` and returns the new integer primary key value. If the ``Species`` table already exists, it will insert the new row and return the primary key. If a row with that ``name`` already exists, it will return the corresponding primary key value directly. @@ -380,6 +385,39 @@ If you pass in a dictionary with multiple values, both values will be used to in }) }) +.. _python_api_extracts: + +Populating lookup tables automatically during insert/upsert +----------------------------------------------------------- + +A more efficient way to work with lookup tables is to define them using the ``extracts=`` parameter, which is accepted by ``.insert()``, ``.upsert()``, ``.insert_all()``, ``.upsert_all()`` and by the ``.table(...)`` factory function. + +``extracts=`` specifies columns which should be "extracted" out into a separate lookup table during the data insertion. + +It can be either a list of column names, in which case the extracted table names will match the column names exactly, or it can be a dictionary mapping column names to the desired name of the extracted table. + +To extract the ``species`` column out to a separate ``Species`` table, you can do this: + +.. code-block:: python + + # Using the table factory + trees = db.table("Trees", extracts={"species": "Species"}) + trees.insert({ + "latitude": 49.1265976, + "longitude": 2.5496218, + "species": "Common Juniper" + }) + + # If you want the table to be called 'species', you can do this: + trees = db.table("Trees", extracts=["species"]) + + # Using .insert() directly + db["Trees"].insert({ + "latitude": 49.1265976, + "longitude": 2.5496218, + "species": "Common Juniper" + }, extracts={"species": "Species"}) + .. _python_api_add_column: Adding columns diff --git a/sqlite_utils/db.py b/sqlite_utils/db.py index 8cf31ed..ef55976 100644 --- a/sqlite_utils/db.py +++ b/sqlite_utils/db.py @@ -186,9 +186,20 @@ class Database: not_null=None, defaults=None, hash_id=None, + extracts=None, ): foreign_keys = self.resolve_foreign_keys(name, foreign_keys or []) foreign_keys_by_column = {fk.column: fk for fk in foreign_keys} + # any extracts will be treated as integer columns with a foreign key + extracts = resolve_extracts(extracts) + for extract_column, extract_table in extracts.items(): + # Ensure other table exists + if not self[extract_table].exists: + self.create_table(extract_table, {"id": int, "value": str}, pk="id") + columns[extract_column] = int + foreign_keys_by_column[extract_column] = ForeignKey( + name, extract_column, extract_table, "id" + ) # Sanity check not_null, and defaults if provided not_null = not_null or set() defaults = defaults or {} @@ -376,6 +387,7 @@ class Table: hash_id=None, alter=False, ignore=False, + extracts=None, ): self.db = db self.name = name @@ -391,6 +403,7 @@ class Table: hash_id=hash_id, alter=alter, ignore=ignore, + extracts=extracts, ) def __repr__(self): @@ -530,6 +543,7 @@ class Table: not_null=None, defaults=None, hash_id=None, + extracts=None, ): columns = {name: value for (name, value) in columns.items()} with self.db.conn: @@ -542,6 +556,7 @@ class Table: not_null=not_null, defaults=defaults, hash_id=hash_id, + extracts=extracts, ) self.exists = True return self @@ -779,6 +794,7 @@ class Table: hash_id=DEFAULT, alter=DEFAULT, ignore=DEFAULT, + extracts=DEFAULT, ): return self.insert_all( [record], @@ -791,6 +807,7 @@ class Table: hash_id=hash_id, alter=alter, ignore=ignore, + extracts=extracts, ) def insert_all( @@ -806,6 +823,7 @@ class Table: hash_id=DEFAULT, alter=DEFAULT, ignore=DEFAULT, + extracts=DEFAULT, ): """ Like .insert() but takes a list of records and ensures that the table @@ -822,6 +840,7 @@ class Table: hash_id = self.value_or_default("hash_id", hash_id) alter = self.value_or_default("alter", alter) ignore = self.value_or_default("ignore", ignore) + extracts = self.value_or_default("extracts", extracts) assert not (hash_id and pk), "Use either pk= or hash_id=" assert not ( @@ -842,6 +861,7 @@ class Table: not_null=not_null, defaults=defaults, hash_id=hash_id, + extracts=extracts, ) all_columns = set() for record in chunk: @@ -871,13 +891,18 @@ class Table: ), ) values = [] + extracts = resolve_extracts(extracts) for record in chunk: - values.extend( - jsonify_if_needed( + record_values = [] + for key in all_columns: + value = jsonify_if_needed( record.get(key, None if key != hash_id else _hash(record)) ) - for key in all_columns - ) + if key in extracts: + extract_table = extracts[key] + value = self.db[extract_table].lookup({"value": value}) + record_values.append(value) + values.extend(record_values) with self.db.conn: try: result = self.db.conn.execute(sql, values) @@ -911,6 +936,7 @@ class Table: defaults=DEFAULT, hash_id=DEFAULT, alter=DEFAULT, + extracts=DEFAULT, ): return self.insert( record, @@ -922,6 +948,7 @@ class Table: hash_id=hash_id, alter=alter, upsert=True, + extracts=extracts, ) def upsert_all( @@ -935,6 +962,7 @@ class Table: batch_size=DEFAULT, hash_id=DEFAULT, alter=DEFAULT, + extracts=DEFAULT, ): return self.insert_all( records, @@ -947,6 +975,7 @@ class Table: hash_id=hash_id, alter=alter, upsert=True, + extracts=extracts, ) def add_missing_columns(self, records): @@ -958,6 +987,7 @@ class Table: def lookup(self, column_values): # lookups is a dictionary - all columns will be used for a unique index + assert isinstance(column_values, dict) if self.exists: self.add_missing_columns([column_values]) unique_column_sets = [set(i.columns) for i in self.indexes] @@ -998,3 +1028,11 @@ def _hash(record): return hashlib.sha1( json.dumps(record, separators=(",", ":"), sort_keys=True).encode("utf8") ).hexdigest() + + +def resolve_extracts(extracts): + if extracts is None: + extracts = {} + if isinstance(extracts, (list, tuple)): + extracts = {item: item for item in extracts} + return extracts diff --git a/tests/test_extracts.py b/tests/test_extracts.py new file mode 100644 index 0000000..c46b0a5 --- /dev/null +++ b/tests/test_extracts.py @@ -0,0 +1,63 @@ +from sqlite_utils.db import Index +import pytest + + +@pytest.mark.parametrize( + "kwargs,expected_table", + [ + (dict(extracts={"species_id": "Species"}), "Species"), + (dict(extracts=["species_id"]), "species_id"), + (dict(extracts=("species_id",)), "species_id"), + ], +) +@pytest.mark.parametrize("use_table_factory", [True, False]) +def test_extracts(fresh_db, kwargs, expected_table, use_table_factory): + table_kwargs = {} + insert_kwargs = {} + if use_table_factory: + table_kwargs = kwargs + else: + insert_kwargs = kwargs + trees = fresh_db.table("Trees", **table_kwargs) + trees.insert_all( + [ + {"id": 1, "species_id": "Oak"}, + {"id": 2, "species_id": "Oak"}, + {"id": 3, "species_id": "Palm"}, + ], + **insert_kwargs + ) + # Should now have two tables: Trees and Species + assert {expected_table, "Trees"} == set(fresh_db.table_names()) + assert ( + "CREATE TABLE [{}] (\n [id] INTEGER PRIMARY KEY,\n [value] TEXT\n)".format( + expected_table + ) + == fresh_db[expected_table].schema + ) + assert ( + "CREATE TABLE [Trees] (\n [id] INTEGER,\n [species_id] INTEGER REFERENCES [{}]([id])\n)".format( + expected_table + ) + == fresh_db["Trees"].schema + ) + # Should have unique index on Species + assert [ + Index( + seq=0, + name="idx_{}_value".format(expected_table), + unique=1, + origin="c", + partial=0, + columns=["value"], + ) + ] == fresh_db[expected_table].indexes + # Finally, check the rows + assert [{"id": 1, "value": "Oak"}, {"id": 2, "value": "Palm"}] == list( + fresh_db[expected_table].rows + ) + assert [ + {"id": 1, "species_id": 1}, + {"id": 2, "species_id": 1}, + {"id": 3, "species_id": 2}, + ] == list(fresh_db["Trees"].rows)