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

This commit is contained in:
Simon Willison 2021-11-14 18:01:56 -08:00
commit 54a2269e91
3 changed files with 53 additions and 11 deletions

View file

@ -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

View file

@ -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(

View file

@ -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",
}