From ec3d3b2978fbafa90c9cd186d5b62b257984c5e7 Mon Sep 17 00:00:00 2001 From: Simon Willison Date: Tue, 23 Jul 2019 18:26:53 +0200 Subject: [PATCH] extracts= table parameter, refs #46 Still needs docs --- sqlite_utils/db.py | 46 +++++++++++++++++++++++++++++++++---- tests/test_extracts.py | 52 ++++++++++++++++++++++++++++++++++++++++++ 2 files changed, 94 insertions(+), 4 deletions(-) create mode 100644 tests/test_extracts.py diff --git a/sqlite_utils/db.py b/sqlite_utils/db.py index 8cf31ed..a1fcdca 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): + 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..dca45b6 --- /dev/null +++ b/tests/test_extracts.py @@ -0,0 +1,52 @@ +from sqlite_utils.db import Index +import pytest + + +@pytest.mark.parametrize("use_table_factory", [True, False]) +def test_extracts(fresh_db, use_table_factory): + kwargs = dict(extracts={"species_id": "Species"}) + 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 {"Species", "Trees"} == set(fresh_db.table_names()) + assert ( + "CREATE TABLE [Species] (\n [id] INTEGER PRIMARY KEY,\n [value] TEXT\n)" + == fresh_db["Species"].schema + ) + assert ( + "CREATE TABLE [Trees] (\n [id] INTEGER,\n [species_id] INTEGER REFERENCES [Species]([id])\n)" + == fresh_db["Trees"].schema + ) + # Should have unique index on Species + assert [ + Index( + seq=0, + name="idx_Species_value", + unique=1, + origin="c", + partial=0, + columns=["value"], + ) + ] == fresh_db["Species"].indexes + # Finally, check the rows + assert [{"id": 1, "value": "Oak"}, {"id": 2, "value": "Palm"}] == list( + fresh_db["Species"].rows + ) + assert [ + {"id": 1, "species_id": 1}, + {"id": 2, "species_id": 1}, + {"id": 3, "species_id": 2}, + ] == list(fresh_db["Trees"].rows)