mirror of
https://github.com/simonw/sqlite-utils.git
synced 2026-07-23 01:14:31 +02:00
table.extract() method, refs #42
This commit is contained in:
Simon Willison
parent
f29f6821f2
commit
f8553799d3
3 changed files with 272 additions and 0 deletions
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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(
|
||||
|
|
|
|||
105
tests/test_extract.py
Normal file
105
tests/test_extract.py
Normal file
|
|
@ -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"])
|
||||
Loading…
Add table
Add a link
Reference in a new issue