Repository: simonw/sqlite-utils

Issue:
Utilities for building lookup tables
While building https://github.com/dogsheep/healthkit-to-sqlite I found a need for a neat mechanism for easily building lookup tables - tables where each unique value in a column is replaced by a foreign key to a separate table.

csvs-to-sqlite currently creates those with its "extract" mechanism - but that's written as custom code against Pandas. I'd like to eventually replace Pandas with sqlite-utils there.

See also #42 

Discussion:

Patch:
diff --git a/docs/python-api.rst b/docs/python-api.rst
index 378b5d9..0bbf19d 100644
--- a/docs/python-api.rst
+++ b/docs/python-api.rst
@@ -354,6 +354,32 @@ Note that the ``pk`` and ``column_order`` parameters here are optional if you ar
 
 An ``upsert_all()`` method is also available, which behaves like ``insert_all()`` but performs upserts instead.
 
+.. _python_api_lookup_tables:
+
+Creating 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.
+
+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.
+
+If you call ``.lookup()`` against an existing table without the unique constraint it will attempt to add the constraint, raising an ``IntegrityError`` if the constraint cannot be created.
+
+If you pass in a dictionary with multiple values, both values will be used to insert or retrieve the corresponding ID and any unique constraint that is created will cover all of those columns, for example:
+
+.. code-block:: python
+
+    db["Trees"].insert({
+        "latitude": 49.1265976,
+        "longitude": 2.5496218,
+        "species": db["Species"].lookup({
+            "common_name": "Common Juniper",
+            "latin_name": "Juniperus communis"
+        })
+    })
+
 .. _python_api_add_column:
 
 Adding columns
diff --git a/sqlite_utils/db.py b/sqlite_utils/db.py
index c48483d..a86132b 100644
--- a/sqlite_utils/db.py
+++ b/sqlite_utils/db.py
@@ -224,6 +224,8 @@ class Database:
         single_pk = None
         if isinstance(pk, str):
             single_pk = pk
+            if pk not in [c[0] for c in column_items]:
+                column_items.insert(0, (pk, int))
         for column_name, column_type in column_items:
             column_extras = []
             if column_name == single_pk:
@@ -954,6 +956,29 @@ class Table:
             if col_name not in current_columns:
                 self.add_column(col_name, col_type)
 
+    def lookup(self, column_values):
+        # lookups is a dictionary - all columns will be used for a unique index
+        if self.exists:
+            self.add_missing_columns([column_values])
+            # TODO: Ensure we have a unique index on these
+            unique_column_sets = [set(i.columns) for i in self.indexes]
+            if set(column_values.keys()) not in unique_column_sets:
+                self.create_index(column_values.keys(), unique=True)
+            wheres = ["[{}] = ?".format(column) for column in column_values]
+            rows = list(
+                self.rows_where(
+                    " and ".join(wheres), [value for _, value in column_values.items()]
+                )
+            )
+            try:
+                return rows[0]["id"]
+            except IndexError:
+                return self.insert(column_values, pk="id").last_pk
+        else:
+            pk = self.insert(column_values, pk="id").last_pk
+            self.create_index(column_values.keys(), unique=True)
+            return pk
+
 
 def chunks(sequence, size):
     iterator = iter(sequence)


Test patch:
diff --git a/tests/test_create.py b/tests/test_create.py
index 288a8ee..222a967 100644
--- a/tests/test_create.py
+++ b/tests/test_create.py
@@ -734,3 +734,8 @@ def test_cannot_provide_both_filename_and_memory():
         AssertionError, match="Either specify a filename_or_conn or pass memory=True"
     ):
         Database("/tmp/foo.db", memory=True)
+
+
+def test_creates_id_column(fresh_db):
+    last_pk = fresh_db.table("cats", pk="id").insert({"name": "barry"}).last_pk
+    assert [{"name": "barry", "id": last_pk}] == list(fresh_db["cats"].rows)
diff --git a/tests/test_lookup.py b/tests/test_lookup.py
new file mode 100644
index 0000000..ac8e1c2
--- /dev/null
+++ b/tests/test_lookup.py
@@ -0,0 +1,68 @@
+from sqlite_utils.db import Index
+import pytest
+
+
+def test_lookup_new_table(fresh_db):
+    species = fresh_db["species"]
+    palm_id = species.lookup({"name": "Palm"})
+    oak_id = species.lookup({"name": "Oak"})
+    cherry_id = species.lookup({"name": "Cherry"})
+    assert palm_id == species.lookup({"name": "Palm"})
+    assert oak_id == species.lookup({"name": "Oak"})
+    assert cherry_id == species.lookup({"name": "Cherry"})
+    assert palm_id != oak_id != cherry_id
+    # Ensure the correct indexes were created
+    assert [
+        Index(
+            seq=0,
+            name="idx_species_name",
+            unique=1,
+            origin="c",
+            partial=0,
+            columns=["name"],
+        )
+    ] == species.indexes
+
+
+def test_lookup_new_table_compound_key(fresh_db):
+    species = fresh_db["species"]
+    palm_id = species.lookup({"name": "Palm", "type": "Tree"})
+    oak_id = species.lookup({"name": "Oak", "type": "Tree"})
+    assert palm_id == species.lookup({"name": "Palm", "type": "Tree"})
+    assert oak_id == species.lookup({"name": "Oak", "type": "Tree"})
+    assert [
+        Index(
+            seq=0,
+            name="idx_species_name_type",
+            unique=1,
+            origin="c",
+            partial=0,
+            columns=["name", "type"],
+        )
+    ] == species.indexes
+
+
+def test_lookup_adds_unique_constraint_to_existing_table(fresh_db):
+    species = fresh_db.table("species", pk="id")
+    palm_id = species.insert({"name": "Palm"}).last_pk
+    species.insert({"name": "Oak"})
+    assert [] == species.indexes
+    assert palm_id == species.lookup({"name": "Palm"})
+    assert [
+        Index(
+            seq=0,
+            name="idx_species_name",
+            unique=1,
+            origin="c",
+            partial=0,
+            columns=["name"],
+        )
+    ] == species.indexes
+
+
+def test_lookup_fails_if_constraint_cannot_be_added(fresh_db):
+    species = fresh_db.table("species", pk="id")
+    species.insert_all([{"id": 1, "name": "Palm"}, {"id": 2, "name": "Palm"}])
+    # This will fail because the name column is not unique
+    with pytest.raises(Exception, match="UNIQUE constraint failed"):
+        species.lookup({"name": "Palm"})