diff --git a/sqlite_utils/db.py b/sqlite_utils/db.py index c011d9bb0..4b0636855 100644 --- a/sqlite_utils/db.py +++ b/sqlite_utils/db.py @@ -5163,12 +5163,15 @@ def lookup( if self.exists(): self.add_missing_columns([combined_values]) unique_column_sets = [ - {fold_identifier_case(c) for c in i.columns} for i in self.indexes + {fold_identifier_case(c) for c in i.columns} + for i in self.indexes + if i.unique and not i.partial ] - if { - fold_identifier_case(c) for c in lookup_values - } not in unique_column_sets: - self.create_index(lookup_values.keys(), unique=True) + lookup_column_set = {fold_identifier_case(c) for c in lookup_values} + if lookup_column_set not in unique_column_sets: + self.create_index( + lookup_values.keys(), unique=True, find_unique_name=True + ) # IS rather than = so that null values are matched correctly wheres = [f"{quote_identifier(column)} IS ?" for column in lookup_values] rows = list( diff --git a/tests/test_lookup.py b/tests/test_lookup.py index f96cfefaf..474bc6760 100644 --- a/tests/test_lookup.py +++ b/tests/test_lookup.py @@ -1,3 +1,5 @@ +import sqlite3 + import pytest from sqlite_utils.db import Index @@ -182,3 +184,69 @@ def test_lookup_compound_key_with_null_idempotent(fresh_db): {"id": palm_id, "name": "Palm", "type": None}, {"id": oak_id, "name": "Oak", "type": "Tree"}, ] + + +@pytest.mark.parametrize("index_name", ["custom_species_name", None]) +def test_lookup_adds_full_unique_index_when_only_ordinary_index_exists( + fresh_db, index_name +): + species = fresh_db.table("species", pk="id") + species.insert({"id": 1, "name": "Palm"}) + if index_name is None: + species.create_index(["name"]) + ordinary_index_name = "idx_species_name" + else: + species.create_index(["name"], index_name=index_name) + ordinary_index_name = index_name + + species.lookup({"name": "Palm"}) + + indexes = species.indexes + assert any( + index.unique == 1 and index.partial == 0 and index.columns == ["name"] + for index in indexes + ) + assert any(index.name == ordinary_index_name for index in indexes) + assert len([index for index in indexes if index.columns == ["name"]]) == 2 + + with pytest.raises(sqlite3.IntegrityError): + species.insert({"id": 2, "name": "Palm"}) + + +def test_lookup_rejects_ambiguous_rows_with_ordinary_index(fresh_db): + species = fresh_db.table("species", pk="id") + species.insert_all([{"id": 1, "name": "Palm"}, {"id": 2, "name": "Palm"}]) + species.create_index(["name"], index_name="idx_species_name") + original_rows = list(species.rows) + + with pytest.raises(sqlite3.IntegrityError): + species.lookup({"name": "Palm"}) + + assert list(species.rows) == original_rows + + +def test_partial_unique_index_does_not_count_as_full_lookup_constraint(fresh_db): + species = fresh_db.table("species", pk="id") + species.insert({"id": 1, "name": "Palm"}) + species.db.execute( + 'CREATE UNIQUE INDEX "idx_species_name_partial" ' + 'ON "species" ("name") WHERE "name" = \'Palm\'' + ) + + species.lookup({"name": "Oak"}) + + assert any( + index.unique == 1 and index.partial == 0 and index.columns == ["name"] + for index in species.indexes + ) + + +def test_lookup_reuses_existing_full_unique_index(fresh_db): + species = fresh_db.table("species", pk="id") + species.insert({"id": 1, "name": "Palm"}) + species.create_index(["name"], index_name="existing_unique_name", unique=True) + original_indexes = species.indexes + + species.lookup({"name": "Palm"}) + + assert species.indexes == original_indexes