mirror of
https://github.com/simonw/sqlite-utils.git
synced 2026-07-23 09:24:31 +02:00
extracts= table parameter, closes #46
This commit is contained in:
Simon Willison
GitHub
parent
e22cfcd953
commit
941d281aee
4 changed files with 146 additions and 7 deletions
|
|
@ -16,7 +16,7 @@ Contents
|
|||
--------
|
||||
|
||||
.. toctree::
|
||||
:maxdepth: 2
|
||||
:maxdepth: 3
|
||||
|
||||
cli
|
||||
python-api
|
||||
|
|
|
|||
|
|
@ -356,11 +356,16 @@ An ``upsert_all()`` method is also available, which behaves like ``insert_all()`
|
|||
|
||||
.. _python_api_lookup_tables:
|
||||
|
||||
Creating lookup tables
|
||||
======================
|
||||
Working with 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.
|
||||
|
||||
.. _python_api_explicit_lookup_tables:
|
||||
|
||||
Creating lookup tables explicitly
|
||||
---------------------------------
|
||||
|
||||
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.
|
||||
|
|
@ -380,6 +385,39 @@ If you pass in a dictionary with multiple values, both values will be used to in
|
|||
})
|
||||
})
|
||||
|
||||
.. _python_api_extracts:
|
||||
|
||||
Populating lookup tables automatically during insert/upsert
|
||||
-----------------------------------------------------------
|
||||
|
||||
A more efficient way to work with lookup tables is to define them using the ``extracts=`` parameter, which is accepted by ``.insert()``, ``.upsert()``, ``.insert_all()``, ``.upsert_all()`` and by the ``.table(...)`` factory function.
|
||||
|
||||
``extracts=`` specifies columns which should be "extracted" out into a separate lookup table during the data insertion.
|
||||
|
||||
It can be either a list of column names, in which case the extracted table names will match the column names exactly, or it can be a dictionary mapping column names to the desired name of the extracted table.
|
||||
|
||||
To extract the ``species`` column out to a separate ``Species`` table, you can do this:
|
||||
|
||||
.. code-block:: python
|
||||
|
||||
# Using the table factory
|
||||
trees = db.table("Trees", extracts={"species": "Species"})
|
||||
trees.insert({
|
||||
"latitude": 49.1265976,
|
||||
"longitude": 2.5496218,
|
||||
"species": "Common Juniper"
|
||||
})
|
||||
|
||||
# If you want the table to be called 'species', you can do this:
|
||||
trees = db.table("Trees", extracts=["species"])
|
||||
|
||||
# Using .insert() directly
|
||||
db["Trees"].insert({
|
||||
"latitude": 49.1265976,
|
||||
"longitude": 2.5496218,
|
||||
"species": "Common Juniper"
|
||||
}, extracts={"species": "Species"})
|
||||
|
||||
.. _python_api_add_column:
|
||||
|
||||
Adding columns
|
||||
|
|
|
|||
|
|
@ -186,9 +186,20 @@ class Database:
|
|||
not_null=None,
|
||||
defaults=None,
|
||||
hash_id=None,
|
||||
extracts=None,
|
||||
):
|
||||
foreign_keys = self.resolve_foreign_keys(name, foreign_keys or [])
|
||||
foreign_keys_by_column = {fk.column: fk for fk in foreign_keys}
|
||||
# any extracts will be treated as integer columns with a foreign key
|
||||
extracts = resolve_extracts(extracts)
|
||||
for extract_column, extract_table in extracts.items():
|
||||
# Ensure other table exists
|
||||
if not self[extract_table].exists:
|
||||
self.create_table(extract_table, {"id": int, "value": str}, pk="id")
|
||||
columns[extract_column] = int
|
||||
foreign_keys_by_column[extract_column] = ForeignKey(
|
||||
name, extract_column, extract_table, "id"
|
||||
)
|
||||
# Sanity check not_null, and defaults if provided
|
||||
not_null = not_null or set()
|
||||
defaults = defaults or {}
|
||||
|
|
@ -376,6 +387,7 @@ class Table:
|
|||
hash_id=None,
|
||||
alter=False,
|
||||
ignore=False,
|
||||
extracts=None,
|
||||
):
|
||||
self.db = db
|
||||
self.name = name
|
||||
|
|
@ -391,6 +403,7 @@ class Table:
|
|||
hash_id=hash_id,
|
||||
alter=alter,
|
||||
ignore=ignore,
|
||||
extracts=extracts,
|
||||
)
|
||||
|
||||
def __repr__(self):
|
||||
|
|
@ -530,6 +543,7 @@ class Table:
|
|||
not_null=None,
|
||||
defaults=None,
|
||||
hash_id=None,
|
||||
extracts=None,
|
||||
):
|
||||
columns = {name: value for (name, value) in columns.items()}
|
||||
with self.db.conn:
|
||||
|
|
@ -542,6 +556,7 @@ class Table:
|
|||
not_null=not_null,
|
||||
defaults=defaults,
|
||||
hash_id=hash_id,
|
||||
extracts=extracts,
|
||||
)
|
||||
self.exists = True
|
||||
return self
|
||||
|
|
@ -779,6 +794,7 @@ class Table:
|
|||
hash_id=DEFAULT,
|
||||
alter=DEFAULT,
|
||||
ignore=DEFAULT,
|
||||
extracts=DEFAULT,
|
||||
):
|
||||
return self.insert_all(
|
||||
[record],
|
||||
|
|
@ -791,6 +807,7 @@ class Table:
|
|||
hash_id=hash_id,
|
||||
alter=alter,
|
||||
ignore=ignore,
|
||||
extracts=extracts,
|
||||
)
|
||||
|
||||
def insert_all(
|
||||
|
|
@ -806,6 +823,7 @@ class Table:
|
|||
hash_id=DEFAULT,
|
||||
alter=DEFAULT,
|
||||
ignore=DEFAULT,
|
||||
extracts=DEFAULT,
|
||||
):
|
||||
"""
|
||||
Like .insert() but takes a list of records and ensures that the table
|
||||
|
|
@ -822,6 +840,7 @@ class Table:
|
|||
hash_id = self.value_or_default("hash_id", hash_id)
|
||||
alter = self.value_or_default("alter", alter)
|
||||
ignore = self.value_or_default("ignore", ignore)
|
||||
extracts = self.value_or_default("extracts", extracts)
|
||||
|
||||
assert not (hash_id and pk), "Use either pk= or hash_id="
|
||||
assert not (
|
||||
|
|
@ -842,6 +861,7 @@ class Table:
|
|||
not_null=not_null,
|
||||
defaults=defaults,
|
||||
hash_id=hash_id,
|
||||
extracts=extracts,
|
||||
)
|
||||
all_columns = set()
|
||||
for record in chunk:
|
||||
|
|
@ -871,13 +891,18 @@ class Table:
|
|||
),
|
||||
)
|
||||
values = []
|
||||
extracts = resolve_extracts(extracts)
|
||||
for record in chunk:
|
||||
values.extend(
|
||||
jsonify_if_needed(
|
||||
record_values = []
|
||||
for key in all_columns:
|
||||
value = jsonify_if_needed(
|
||||
record.get(key, None if key != hash_id else _hash(record))
|
||||
)
|
||||
for key in all_columns
|
||||
)
|
||||
if key in extracts:
|
||||
extract_table = extracts[key]
|
||||
value = self.db[extract_table].lookup({"value": value})
|
||||
record_values.append(value)
|
||||
values.extend(record_values)
|
||||
with self.db.conn:
|
||||
try:
|
||||
result = self.db.conn.execute(sql, values)
|
||||
|
|
@ -911,6 +936,7 @@ class Table:
|
|||
defaults=DEFAULT,
|
||||
hash_id=DEFAULT,
|
||||
alter=DEFAULT,
|
||||
extracts=DEFAULT,
|
||||
):
|
||||
return self.insert(
|
||||
record,
|
||||
|
|
@ -922,6 +948,7 @@ class Table:
|
|||
hash_id=hash_id,
|
||||
alter=alter,
|
||||
upsert=True,
|
||||
extracts=extracts,
|
||||
)
|
||||
|
||||
def upsert_all(
|
||||
|
|
@ -935,6 +962,7 @@ class Table:
|
|||
batch_size=DEFAULT,
|
||||
hash_id=DEFAULT,
|
||||
alter=DEFAULT,
|
||||
extracts=DEFAULT,
|
||||
):
|
||||
return self.insert_all(
|
||||
records,
|
||||
|
|
@ -947,6 +975,7 @@ class Table:
|
|||
hash_id=hash_id,
|
||||
alter=alter,
|
||||
upsert=True,
|
||||
extracts=extracts,
|
||||
)
|
||||
|
||||
def add_missing_columns(self, records):
|
||||
|
|
@ -958,6 +987,7 @@ class Table:
|
|||
|
||||
def lookup(self, column_values):
|
||||
# lookups is a dictionary - all columns will be used for a unique index
|
||||
assert isinstance(column_values, dict)
|
||||
if self.exists:
|
||||
self.add_missing_columns([column_values])
|
||||
unique_column_sets = [set(i.columns) for i in self.indexes]
|
||||
|
|
@ -998,3 +1028,11 @@ def _hash(record):
|
|||
return hashlib.sha1(
|
||||
json.dumps(record, separators=(",", ":"), sort_keys=True).encode("utf8")
|
||||
).hexdigest()
|
||||
|
||||
|
||||
def resolve_extracts(extracts):
|
||||
if extracts is None:
|
||||
extracts = {}
|
||||
if isinstance(extracts, (list, tuple)):
|
||||
extracts = {item: item for item in extracts}
|
||||
return extracts
|
||||
|
|
|
|||
63
tests/test_extracts.py
Normal file
63
tests/test_extracts.py
Normal file
|
|
@ -0,0 +1,63 @@
|
|||
from sqlite_utils.db import Index
|
||||
import pytest
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"kwargs,expected_table",
|
||||
[
|
||||
(dict(extracts={"species_id": "Species"}), "Species"),
|
||||
(dict(extracts=["species_id"]), "species_id"),
|
||||
(dict(extracts=("species_id",)), "species_id"),
|
||||
],
|
||||
)
|
||||
@pytest.mark.parametrize("use_table_factory", [True, False])
|
||||
def test_extracts(fresh_db, kwargs, expected_table, use_table_factory):
|
||||
table_kwargs = {}
|
||||
insert_kwargs = {}
|
||||
if use_table_factory:
|
||||
table_kwargs = kwargs
|
||||
else:
|
||||
insert_kwargs = kwargs
|
||||
trees = fresh_db.table("Trees", **table_kwargs)
|
||||
trees.insert_all(
|
||||
[
|
||||
{"id": 1, "species_id": "Oak"},
|
||||
{"id": 2, "species_id": "Oak"},
|
||||
{"id": 3, "species_id": "Palm"},
|
||||
],
|
||||
**insert_kwargs
|
||||
)
|
||||
# Should now have two tables: Trees and Species
|
||||
assert {expected_table, "Trees"} == set(fresh_db.table_names())
|
||||
assert (
|
||||
"CREATE TABLE [{}] (\n [id] INTEGER PRIMARY KEY,\n [value] TEXT\n)".format(
|
||||
expected_table
|
||||
)
|
||||
== fresh_db[expected_table].schema
|
||||
)
|
||||
assert (
|
||||
"CREATE TABLE [Trees] (\n [id] INTEGER,\n [species_id] INTEGER REFERENCES [{}]([id])\n)".format(
|
||||
expected_table
|
||||
)
|
||||
== fresh_db["Trees"].schema
|
||||
)
|
||||
# Should have unique index on Species
|
||||
assert [
|
||||
Index(
|
||||
seq=0,
|
||||
name="idx_{}_value".format(expected_table),
|
||||
unique=1,
|
||||
origin="c",
|
||||
partial=0,
|
||||
columns=["value"],
|
||||
)
|
||||
] == fresh_db[expected_table].indexes
|
||||
# Finally, check the rows
|
||||
assert [{"id": 1, "value": "Oak"}, {"id": 2, "value": "Palm"}] == list(
|
||||
fresh_db[expected_table].rows
|
||||
)
|
||||
assert [
|
||||
{"id": 1, "species_id": 1},
|
||||
{"id": 2, "species_id": 1},
|
||||
{"id": 3, "species_id": 2},
|
||||
] == list(fresh_db["Trees"].rows)
|
||||
Loading…
Add table
Add a link
Reference in a new issue