diff --git a/docs/python-api.rst b/docs/python-api.rst index cc236a2..1c0dc60 100644 --- a/docs/python-api.rst +++ b/docs/python-api.rst @@ -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 ` 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 diff --git a/sqlite_utils/db.py b/sqlite_utils/db.py index fd58c40..aebca89 100644 --- a/sqlite_utils/db.py +++ b/sqlite_utils/db.py @@ -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 diff --git a/tests/test_lookup.py b/tests/test_lookup.py index 348ef92..31be414 100644 --- a/tests/test_lookup.py +++ b/tests/test_lookup.py @@ -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"], + ) + ]