table.extract() method, refs #42

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

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(

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"])