Optional second argument to .lookup() to populate extra columns, closes #339

Author: simonw

docs/python-api.rst

@@ -824,6 +824,16 @@ If you pass in a dictionary with multiple values, both values will be used to in
         })
     })
 
+The ``.lookup()`` method has an optional second argument which can be used to populate other columns in the table but only if the row does not exist yet. These columns will not be included in the unique index.
+
+To create a species record with a note on when it was first seen, you can use this:
+
+.. code-block:: python
+
+    db["Species"].lookup({"name": "Palm"}, {"first_seen": "2021-03-04"})
+
+The first time this is called the record will be created for ``name="Palm"``. Any subsequent calls with that name will ignore the second argument, even if it includes different values.
+
 .. _python_api_extracts:
 
 Populating lookup tables automatically during insert/upsert

sqlite_utils/db.py

@@ -2713,7 +2713,11 @@ def add_missing_columns(self, records: Iterable[Dict[str, Any]]) -> "Table":
                 self.add_column(col_name, col_type)
         return self
 
-    def lookup(self, column_values: Dict[str, Any]):
+    def lookup(
+        self,
+        lookup_values: Dict[str, Any],
+        extra_values: Optional[Dict[str, Any]] = None,
+    ):
         """
         Create or populate a lookup table with the specified values.
 
@@ -2725,28 +2729,36 @@ def lookup(self, column_values: Dict[str, Any]):
         It will then insert a new row with the ``name`` set to ``Palm`` and return the
         new integer primary key value.
 
+        An optional second argument can be provided with more ``name: value`` pairs to
+        be included only if the record is being created for the first time. These will
+        be ignored on subsequent lookup calls for records that already exist.
+
         See :ref:`python_api_lookup_tables` for more details.
         """
-        # lookups is a dictionary - all columns will be used for a unique index
-        assert isinstance(column_values, dict)
+        assert isinstance(lookup_values, dict)
+        if extra_values is not None:
+            assert isinstance(extra_values, dict)
+        combined_values = dict(lookup_values)
+        if extra_values is not None:
+            combined_values.update(extra_values)
         if self.exists():
-            self.add_missing_columns([column_values])
+            self.add_missing_columns([combined_values])
             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]
+            if set(lookup_values.keys()) not in unique_column_sets:
+                self.create_index(lookup_values.keys(), unique=True)
+            wheres = ["[{}] = ?".format(column) for column in lookup_values]
             rows = list(
                 self.rows_where(
-                    " and ".join(wheres), [value for _, value in column_values.items()]
+                    " and ".join(wheres), [value for _, value in lookup_values.items()]
                 )
             )
             try:
                 return rows[0]["id"]
             except IndexError:
-                return self.insert(column_values, pk="id").last_pk
+                return self.insert(combined_values, pk="id").last_pk
         else:
-            pk = self.insert(column_values, pk="id").last_pk
-            self.create_index(column_values.keys(), unique=True)
+            pk = self.insert(combined_values, pk="id").last_pk
+            self.create_index(lookup_values.keys(), unique=True)
             return pk
 
     def m2m(

tests/test_lookup.py

@@ -66,3 +66,23 @@ def test_lookup_fails_if_constraint_cannot_be_added(fresh_db):
     # This will fail because the name column is not unique
     with pytest.raises(Exception, match="UNIQUE constraint failed"):
         species.lookup({"name": "Palm"})
+
+
+def test_lookup_with_extra_values(fresh_db):
+    species = fresh_db["species"]
+    id = species.lookup({"name": "Palm", "type": "Tree"}, {"first_seen": "2020-01-01"})
+    assert species.get(id) == {
+        "id": 1,
+        "name": "Palm",
+        "type": "Tree",
+        "first_seen": "2020-01-01",
+    }
+    # A subsequent lookup() should ignore the second dictionary
+    id2 = species.lookup({"name": "Palm", "type": "Tree"}, {"first_seen": "2021-02-02"})
+    assert id2 == id
+    assert species.get(id2) == {
+        "id": 1,
+        "name": "Palm",
+        "type": "Tree",
+        "first_seen": "2020-01-01",
+    }