Extra parameters for .lookup(), passed to .insert() - closes #342

This commit is contained in:
Simon Willison 2021-11-18 23:26:50 -08:00
commit 93b21c230a
3 changed files with 110 additions and 4 deletions

View file

@ -834,6 +834,17 @@ To create a species record with a note on when it was first seen, you can use th
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.
``.lookup()`` also accepts keyword arguments, which are passed through to the :ref:`insert() method <python_api_creating_tables>` and can be used to influence the shape of the created table. Supported parameters are:
- ``pk`` - which defaults to ``id``
- ``foreign_keys``
- ``column_order``
- ``not_null``
- ``defaults``
- ``extracts``
- ``conversions``
- ``columns``
.. _python_api_extracts:
Populating lookup tables automatically during insert/upsert

View file

@ -2556,7 +2556,7 @@ class Table(Queryable):
ignore = self.value_or_default("ignore", ignore)
replace = self.value_or_default("replace", replace)
extracts = self.value_or_default("extracts", extracts)
conversions = self.value_or_default("conversions", conversions)
conversions = self.value_or_default("conversions", conversions) or {}
columns = self.value_or_default("columns", columns)
if upsert and (not pk and not hash_id):
@ -2718,6 +2718,14 @@ class Table(Queryable):
self,
lookup_values: Dict[str, Any],
extra_values: Optional[Dict[str, Any]] = None,
pk: Optional[str] = "id",
foreign_keys: Optional[ForeignKeysType] = None,
column_order: Optional[List[str]] = None,
not_null: Optional[Set[str]] = None,
defaults: Optional[Dict[str, Any]] = None,
extracts: Optional[Union[Dict[str, str], List[str]]] = None,
conversions: Optional[Dict[str, str]] = None,
columns: Optional[Dict[str, Any]] = None,
):
"""
Create or populate a lookup table with the specified values.
@ -2735,6 +2743,8 @@ class Table(Queryable):
be ignored on subsequent lookup calls for records that already exist.
See :ref:`python_api_lookup_tables` for more details.
All other keyword arguments are passed through to ``.insert()``.
"""
assert isinstance(lookup_values, dict)
if extra_values is not None:
@ -2754,11 +2764,31 @@ class Table(Queryable):
)
)
try:
return rows[0]["id"]
return rows[0][pk]
except IndexError:
return self.insert(combined_values, pk="id").last_pk
return self.insert(
combined_values,
pk=pk,
foreign_keys=foreign_keys,
column_order=column_order,
not_null=not_null,
defaults=defaults,
extracts=extracts,
conversions=conversions,
columns=columns,
).last_pk
else:
pk = self.insert(combined_values, pk="id").last_pk
pk = self.insert(
combined_values,
pk=pk,
foreign_keys=foreign_keys,
column_order=column_order,
not_null=not_null,
defaults=defaults,
extracts=extracts,
conversions=conversions,
columns=columns,
).last_pk
self.create_index(lookup_values.keys(), unique=True)
return pk

View file

@ -86,3 +86,68 @@ def test_lookup_with_extra_values(fresh_db):
"type": "Tree",
"first_seen": "2020-01-01",
}
def test_lookup_with_extra_insert_parameters(fresh_db):
other_table = fresh_db["other_table"]
other_table.insert({"id": 1, "name": "Name"}, pk="id")
species = fresh_db["species"]
id = species.lookup(
{"name": "Palm", "type": "Tree"},
{
"first_seen": "2020-01-01",
"make_not_null": 1,
"fk_to_other": 1,
"default_is_dog": "cat",
"extract_this": "This is extracted",
"convert_to_upper": "upper",
"make_this_integer": "2",
"this_at_front": 1,
},
pk="renamed_id",
foreign_keys=(("fk_to_other", "other_table", "id"),),
column_order=("this_at_front",),
not_null={"make_not_null"},
defaults={"default_is_dog": "dog"},
extracts=["extract_this"],
conversions={"convert_to_upper": "upper(?)"},
columns={"make_this_integer": int},
)
assert species.schema == (
"CREATE TABLE [species] (\n"
" [renamed_id] INTEGER PRIMARY KEY,\n"
" [this_at_front] INTEGER,\n"
" [name] TEXT,\n"
" [type] TEXT,\n"
" [first_seen] TEXT,\n"
" [make_not_null] INTEGER NOT NULL,\n"
" [fk_to_other] INTEGER REFERENCES [other_table]([id]),\n"
" [default_is_dog] TEXT DEFAULT 'dog',\n"
" [extract_this] INTEGER REFERENCES [extract_this]([id]),\n"
" [convert_to_upper] TEXT,\n"
" [make_this_integer] INTEGER\n"
")"
)
assert species.get(id) == {
"renamed_id": id,
"this_at_front": 1,
"name": "Palm",
"type": "Tree",
"first_seen": "2020-01-01",
"make_not_null": 1,
"fk_to_other": 1,
"default_is_dog": "cat",
"extract_this": 1,
"convert_to_upper": "UPPER",
"make_this_integer": 2,
}
assert species.indexes == [
Index(
seq=0,
name="idx_species_name_type",
unique=1,
origin="c",
partial=0,
columns=["name", "type"],
)
]