diff --git a/docs/python-api.rst b/docs/python-api.rst index ba903a3..ebfc203 100644 --- a/docs/python-api.rst +++ b/docs/python-api.rst @@ -1001,6 +1001,141 @@ Custom transformations with .transform_sql() If you want to do something more advanced, you can call the ``table.transform_sql(...)`` method with the same arguments that you would have passed to ``table.transform(...)``. This method will return a list of SQL statements that should be executed to implement the change. You can then make modifications to that SQL before executing it yourself. +.. _python_api_extract: + +Extracting columns into a separate table +======================================== + +The ``table.extract()`` method can be used to extract specified columns into a separate table. + +Imagine a ``Trees`` table that looks like this: + +=== ============ ======= + id TreeAddress Species +=== ============ ======= + 1 52 Vine St Palm + 2 12 Draft St Oak + 3 51 Dark Ave Palm + 4 1252 Left St Palm +=== ============ ======= + +The ``Species`` column contains duplicate values. This database could be improved by extracting that column out into a separate ``Species`` table and pointing to it using a foreign key column. + +The schema of the above table is: + +.. code-block:: sql + + CREATE TABLE [Trees] ( + [id] INTEGER PRIMARY KEY, + [TreeAddress] TEXT, + [Species] TEXT + ) + +Here's how to extract the ``Species`` column using ``.extract()``: + +.. code-block:: python + + db["Trees"].extract("Species") + +After running this code the table schema now looks like this: + +.. code-block:: sql + + CREATE TABLE "Trees" ( + [id] INTEGER PRIMARY KEY, + [TreeAddress] TEXT, + [Species_id] INTEGER, + FOREIGN KEY(Species_id) REFERENCES Species(id) + ) + +A new ``Species`` table will have been created with the following schema: + +.. code-block:: sql + + CREATE TABLE [Species] ( + [id] INTEGER PRIMARY KEY, + [Species] TEXT + ) + +The ``.extract()`` method defaults to creating a table with the same name as the column that was extracted, and adding a foreign key column called ``tablename_id``. + +You can specify a custom table name using ``table=``, and a custom foreign key name using ``fk_column=``. This example creates a table called ``tree_species`` and a foreign key column called ``tree_species_id``: + +.. code-block:: python + + db["Trees"].extract("Species", table="tree_species", fk_column="tree_species_id") + +The resulting schema looks like this: + +.. code-block:: sql + + CREATE TABLE "Trees" ( + [id] INTEGER PRIMARY KEY, + [TreeAddress] TEXT, + [tree_species_id] INTEGER, + FOREIGN KEY(tree_species_id) REFERENCES tree_species(id) + ) + + CREATE TABLE [tree_species] ( + [id] INTEGER PRIMARY KEY, + [Species] TEXT + ) + +You can also extract multiple columns into the same external table. Say for example you have a table like this: + +=== ============ ========== ========= + id TreeAddress CommonName LatinName +=== ============ ========== ========= + 1 52 Vine St Palm Arecaceae + 2 12 Draft St Oak Quercus + 3 51 Dark Ave Palm Arecaceae + 4 1252 Left St Palm Arecaceae +=== ============ ========== ========= + +You can pass ``["CommonName", "LatinName"]`` to ``.extract()`` to extract both of those columns: + +.. code-block:: python + + db["Trees"].extract(["CommonName", "LatinName"]) + +This produces the following schema: + +.. code-block:: sql + + CREATE TABLE "Trees" ( + [id] INTEGER PRIMARY KEY, + [TreeAddress] TEXT, + [CommonName_LatinName_id] INTEGER, + FOREIGN KEY(CommonName_LatinName_id) REFERENCES CommonName_LatinName(id) + ) + CREATE TABLE [CommonName_LatinName] ( + [id] INTEGER PRIMARY KEY, + [CommonName] TEXT, + [LatinName] TEXT + ) + +The table name ``CommonName_LatinName`` is derived from the extract columns. You can use ``table=`` and ``fk_column=`` to specify custom names like this: + +.. code-block:: python + + db["Trees"].extract(["CommonName", "LatinName"], table="Species", fk_column="species_id") + +This produces the following schema: + +.. code-block:: sql + + CREATE TABLE "Trees" ( + [id] INTEGER PRIMARY KEY, + [TreeAddress] TEXT, + [species_id] INTEGER, + FOREIGN KEY(species_id) REFERENCES Species(id) + ) + CREATE TABLE [Species] ( + [id] INTEGER PRIMARY KEY, + [CommonName] TEXT, + [LatinName] TEXT + ) + .. _python_api_hash: Setting an ID based on the hash of the row contents diff --git a/sqlite_utils/db.py b/sqlite_utils/db.py index 452c74a..846f97d 100644 --- a/sqlite_utils/db.py +++ b/sqlite_utils/db.py @@ -104,6 +104,10 @@ class PrimaryKeyRequired(Exception): pass +class InvalidColumns(Exception): + pass + + class Database: def __init__( self, @@ -879,6 +883,34 @@ class Table(Queryable): return sqls + def extract(self, columns, table=None, fk_column=None): + if isinstance(columns, str): + columns = [columns] + if not set(columns).issubset(self.columns_dict.keys()): + raise InvalidColumns( + "Invalid columns {} for table with columns {}".format( + columns, list(self.columns_dict.keys()) + ) + ) + table = table or "_".join(columns) + first_column = columns[0] + pks = self.pks + lookup_table = self.db[table] + for row in self.rows: + row_pks = tuple(row[pk] for pk in pks) + lookups = {column: row[column] for column in columns} + self.update(row_pks, {first_column: lookup_table.lookup(lookups)}) + fk_column = fk_column or "{}_id".format(table) + # Now rename first_column and change its type to integer, and drop + # any other extracted columns: + self.transform( + types={first_column: int}, + drop=set(columns[1:]), + rename={first_column: fk_column}, + ) + # And add the foreign key constraint + self.add_foreign_key(fk_column, table, "id") + def create_index(self, columns, index_name=None, unique=False, if_not_exists=False): if index_name is None: index_name = "idx_{}_{}".format( diff --git a/tests/test_extract.py b/tests/test_extract.py new file mode 100644 index 0000000..31b6d3e --- /dev/null +++ b/tests/test_extract.py @@ -0,0 +1,105 @@ +from sqlite_utils.db import Index, InvalidColumns +import itertools +import pytest + + +@pytest.mark.parametrize("table", [None, "Species"]) +@pytest.mark.parametrize("fk_column", [None, "species"]) +def test_extract_single_column(fresh_db, table, fk_column): + expected_table = table or "species" + expected_fk = fk_column or "{}_id".format(expected_table) + iter_species = itertools.cycle(["Palm", "Spruce", "Mangrove", "Oak"]) + fresh_db["tree"].insert_all( + ( + {"id": i, "name": "Tree {}".format(i), "species": next(iter_species)} + for i in range(1, 1001) + ), + pk="id", + ) + fresh_db["tree"].extract("species", table=table, fk_column=fk_column) + assert fresh_db["tree"].schema == ( + 'CREATE TABLE "tree" (\n' + " [id] INTEGER PRIMARY KEY,\n" + " [name] TEXT,\n" + " [{}] INTEGER,\n".format(expected_fk) + + " FOREIGN KEY({}) REFERENCES {}(id)\n".format(expected_fk, expected_table) + + ")" + ) + assert fresh_db[expected_table].schema == ( + "CREATE TABLE [{}] (\n".format(expected_table) + + " [id] INTEGER PRIMARY KEY,\n" + " [species] TEXT\n" + ")" + ) + assert list(fresh_db[expected_table].rows) == [ + {"id": 1, "species": "Palm"}, + {"id": 2, "species": "Spruce"}, + {"id": 3, "species": "Mangrove"}, + {"id": 4, "species": "Oak"}, + ] + assert list(itertools.islice(fresh_db["tree"].rows, 0, 4)) == [ + {"id": 1, "name": "Tree 1", expected_fk: 1}, + {"id": 2, "name": "Tree 2", expected_fk: 2}, + {"id": 3, "name": "Tree 3", expected_fk: 3}, + {"id": 4, "name": "Tree 4", expected_fk: 4}, + ] + + +def test_extract_multiple_columns(fresh_db): + iter_common = itertools.cycle(["Palm", "Spruce", "Mangrove", "Oak"]) + iter_latin = itertools.cycle(["Arecaceae", "Picea", "Rhizophora", "Quercus"]) + fresh_db["tree"].insert_all( + ( + { + "id": i, + "name": "Tree {}".format(i), + "common_name": next(iter_common), + "latin_name": next(iter_latin), + } + for i in range(1, 1001) + ), + pk="id", + ) + + fresh_db["tree"].extract(["common_name", "latin_name"]) + assert fresh_db["tree"].schema == ( + 'CREATE TABLE "tree" (\n' + " [id] INTEGER PRIMARY KEY,\n" + " [name] TEXT,\n" + " [common_name_latin_name_id] INTEGER,\n" + " FOREIGN KEY(common_name_latin_name_id) REFERENCES common_name_latin_name(id)\n" + ")" + ) + assert fresh_db["common_name_latin_name"].schema == ( + "CREATE TABLE [common_name_latin_name] (\n" + " [id] INTEGER PRIMARY KEY,\n" + " [common_name] TEXT,\n" + " [latin_name] TEXT\n" + ")" + ) + assert list(fresh_db["common_name_latin_name"].rows) == [ + {"common_name": "Palm", "id": 1, "latin_name": "Arecaceae"}, + {"common_name": "Spruce", "id": 2, "latin_name": "Picea"}, + {"common_name": "Mangrove", "id": 3, "latin_name": "Rhizophora"}, + {"common_name": "Oak", "id": 4, "latin_name": "Quercus"}, + ] + assert list(itertools.islice(fresh_db["tree"].rows, 0, 4)) == [ + {"id": 1, "name": "Tree 1", "common_name_latin_name_id": 1}, + {"id": 2, "name": "Tree 2", "common_name_latin_name_id": 2}, + {"id": 3, "name": "Tree 3", "common_name_latin_name_id": 3}, + {"id": 4, "name": "Tree 4", "common_name_latin_name_id": 4}, + ] + + +def test_extract_invalid_columns(fresh_db): + fresh_db["tree"].insert( + { + "id": 1, + "name": "Tree 1", + "common_name": "Palm", + "latin_name": "Arecaceae", + }, + pk="id", + ) + with pytest.raises(InvalidColumns): + fresh_db["tree"].extract(["bad_column"])