diff --git a/docs/python-api.rst b/docs/python-api.rst index 8932656..cc236a2 100644 --- a/docs/python-api.rst +++ b/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 diff --git a/sqlite_utils/db.py b/sqlite_utils/db.py index 828ef68..ecf14be 100644 --- a/sqlite_utils/db.py +++ b/sqlite_utils/db.py @@ -2713,7 +2713,11 @@ class Table(Queryable): 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 @@ class Table(Queryable): 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( diff --git a/tests/test_lookup.py b/tests/test_lookup.py index ac8e1c2..348ef92 100644 --- a/tests/test_lookup.py +++ b/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", + }