From 580502431614d3653c93249988290265f3163d4b Mon Sep 17 00:00:00 2001 From: Simon Willison Date: Tue, 23 Jul 2019 06:06:59 -0700 Subject: [PATCH] Implemented table.lookup(...), closes #44 * Add pk column if missing from insert * Implemented table.lookup(...) --- docs/python-api.rst | 26 +++++++++++++++++ sqlite_utils/db.py | 25 ++++++++++++++++ tests/test_create.py | 5 ++++ tests/test_lookup.py | 68 ++++++++++++++++++++++++++++++++++++++++++++ 4 files changed, 124 insertions(+) create mode 100644 tests/test_lookup.py diff --git a/docs/python-api.rst b/docs/python-api.rst index 378b5d9..0bbf19d 100644 --- a/docs/python-api.rst +++ b/docs/python-api.rst @@ -354,6 +354,32 @@ Note that the ``pk`` and ``column_order`` parameters here are optional if you ar An ``upsert_all()`` method is also available, which behaves like ``insert_all()`` but performs upserts instead. +.. _python_api_lookup_tables: + +Creating lookup tables +====================== + +A useful pattern when populating large tables in to break common values out into lookup tables. Consider a table of ``Trees``, where each tree has a species. Ideally these species would be split out into a separate ``Species`` table, with each one assigned an integer primary key that can be referenced from the ``Trees`` table ``species_id`` column. + +Calling ``db["Species"].lookup({"name": "Palm"})`` creates a table called ``Species`` (if one does not already exist) with two columns: ``id`` and ``name``. It sets up a unique constraint on the ``name`` column to guarantee it will not contain duplicate rows. It then inserts a new row with the ``name`` set to ``Palm`` and returns the new integer primary key value. + +If the ``Species`` table already exists, it will insert the new row and return the primary key. If a row with that ``name`` already exists, it will return the corresponding primary key value directly. + +If you call ``.lookup()`` against an existing table without the unique constraint it will attempt to add the constraint, raising an ``IntegrityError`` if the constraint cannot be created. + +If you pass in a dictionary with multiple values, both values will be used to insert or retrieve the corresponding ID and any unique constraint that is created will cover all of those columns, for example: + +.. code-block:: python + + db["Trees"].insert({ + "latitude": 49.1265976, + "longitude": 2.5496218, + "species": db["Species"].lookup({ + "common_name": "Common Juniper", + "latin_name": "Juniperus communis" + }) + }) + .. _python_api_add_column: Adding columns diff --git a/sqlite_utils/db.py b/sqlite_utils/db.py index c48483d..a86132b 100644 --- a/sqlite_utils/db.py +++ b/sqlite_utils/db.py @@ -224,6 +224,8 @@ class Database: single_pk = None if isinstance(pk, str): single_pk = pk + if pk not in [c[0] for c in column_items]: + column_items.insert(0, (pk, int)) for column_name, column_type in column_items: column_extras = [] if column_name == single_pk: @@ -954,6 +956,29 @@ class Table: if col_name not in current_columns: self.add_column(col_name, col_type) + def lookup(self, column_values): + # lookups is a dictionary - all columns will be used for a unique index + if self.exists: + self.add_missing_columns([column_values]) + # TODO: Ensure we have a unique index on these + 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] + rows = list( + self.rows_where( + " and ".join(wheres), [value for _, value in column_values.items()] + ) + ) + try: + return rows[0]["id"] + except IndexError: + return self.insert(column_values, pk="id").last_pk + else: + pk = self.insert(column_values, pk="id").last_pk + self.create_index(column_values.keys(), unique=True) + return pk + def chunks(sequence, size): iterator = iter(sequence) diff --git a/tests/test_create.py b/tests/test_create.py index 288a8ee..222a967 100644 --- a/tests/test_create.py +++ b/tests/test_create.py @@ -734,3 +734,8 @@ def test_cannot_provide_both_filename_and_memory(): AssertionError, match="Either specify a filename_or_conn or pass memory=True" ): Database("/tmp/foo.db", memory=True) + + +def test_creates_id_column(fresh_db): + last_pk = fresh_db.table("cats", pk="id").insert({"name": "barry"}).last_pk + assert [{"name": "barry", "id": last_pk}] == list(fresh_db["cats"].rows) diff --git a/tests/test_lookup.py b/tests/test_lookup.py new file mode 100644 index 0000000..ac8e1c2 --- /dev/null +++ b/tests/test_lookup.py @@ -0,0 +1,68 @@ +from sqlite_utils.db import Index +import pytest + + +def test_lookup_new_table(fresh_db): + species = fresh_db["species"] + palm_id = species.lookup({"name": "Palm"}) + oak_id = species.lookup({"name": "Oak"}) + cherry_id = species.lookup({"name": "Cherry"}) + assert palm_id == species.lookup({"name": "Palm"}) + assert oak_id == species.lookup({"name": "Oak"}) + assert cherry_id == species.lookup({"name": "Cherry"}) + assert palm_id != oak_id != cherry_id + # Ensure the correct indexes were created + assert [ + Index( + seq=0, + name="idx_species_name", + unique=1, + origin="c", + partial=0, + columns=["name"], + ) + ] == species.indexes + + +def test_lookup_new_table_compound_key(fresh_db): + species = fresh_db["species"] + palm_id = species.lookup({"name": "Palm", "type": "Tree"}) + oak_id = species.lookup({"name": "Oak", "type": "Tree"}) + assert palm_id == species.lookup({"name": "Palm", "type": "Tree"}) + assert oak_id == species.lookup({"name": "Oak", "type": "Tree"}) + assert [ + Index( + seq=0, + name="idx_species_name_type", + unique=1, + origin="c", + partial=0, + columns=["name", "type"], + ) + ] == species.indexes + + +def test_lookup_adds_unique_constraint_to_existing_table(fresh_db): + species = fresh_db.table("species", pk="id") + palm_id = species.insert({"name": "Palm"}).last_pk + species.insert({"name": "Oak"}) + assert [] == species.indexes + assert palm_id == species.lookup({"name": "Palm"}) + assert [ + Index( + seq=0, + name="idx_species_name", + unique=1, + origin="c", + partial=0, + columns=["name"], + ) + ] == species.indexes + + +def test_lookup_fails_if_constraint_cannot_be_added(fresh_db): + species = fresh_db.table("species", pk="id") + species.insert_all([{"id": 1, "name": "Palm"}, {"id": 2, "name": "Palm"}]) + # This will fail because the name column is not unique + with pytest.raises(Exception, match="UNIQUE constraint failed"): + species.lookup({"name": "Palm"})