Docs for extracts=

Still needs docs

docs/index.rst

@@ -16,7 +16,7 @@ Contents
 --------
 
 .. toctree::
-   :maxdepth: 2
+   :maxdepth: 3
 
    cli
    python-api

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

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(
+                record_values = []
-                    jsonify_if_needed(
+                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

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)