mirror of
https://github.com/simonw/sqlite-utils.git
synced 2026-07-26 10:54:32 +02:00
Both of these subclass a common Queryable class. Also updated documentation to cover the new View class. And added view.drop() method.
1176 lines
40 KiB
Python
1176 lines
40 KiB
Python
from .utils import sqlite3, OperationalError
|
|
from collections import namedtuple
|
|
import datetime
|
|
import hashlib
|
|
import itertools
|
|
import json
|
|
import pathlib
|
|
|
|
SQLITE_MAX_VARS = 999
|
|
|
|
try:
|
|
import numpy as np
|
|
except ImportError:
|
|
np = None
|
|
|
|
Column = namedtuple(
|
|
"Column", ("cid", "name", "type", "notnull", "default_value", "is_pk")
|
|
)
|
|
ForeignKey = namedtuple(
|
|
"ForeignKey", ("table", "column", "other_table", "other_column")
|
|
)
|
|
Index = namedtuple("Index", ("seq", "name", "unique", "origin", "partial", "columns"))
|
|
DEFAULT = object()
|
|
|
|
COLUMN_TYPE_MAPPING = {
|
|
float: "FLOAT",
|
|
int: "INTEGER",
|
|
bool: "INTEGER",
|
|
str: "TEXT",
|
|
bytes.__class__: "BLOB",
|
|
bytes: "BLOB",
|
|
datetime.datetime: "TEXT",
|
|
datetime.date: "TEXT",
|
|
datetime.time: "TEXT",
|
|
None.__class__: "TEXT",
|
|
# SQLite explicit types
|
|
"TEXT": "TEXT",
|
|
"INTEGER": "INTEGER",
|
|
"FLOAT": "FLOAT",
|
|
"BLOB": "BLOB",
|
|
"text": "TEXT",
|
|
"integer": "INTEGER",
|
|
"float": "FLOAT",
|
|
"blob": "BLOB",
|
|
}
|
|
# If numpy is available, add more types
|
|
if np:
|
|
COLUMN_TYPE_MAPPING.update(
|
|
{
|
|
np.int8: "INTEGER",
|
|
np.int16: "INTEGER",
|
|
np.int32: "INTEGER",
|
|
np.int64: "INTEGER",
|
|
np.uint8: "INTEGER",
|
|
np.uint16: "INTEGER",
|
|
np.uint32: "INTEGER",
|
|
np.uint64: "INTEGER",
|
|
np.float16: "FLOAT",
|
|
np.float32: "FLOAT",
|
|
np.float64: "FLOAT",
|
|
}
|
|
)
|
|
|
|
|
|
REVERSE_COLUMN_TYPE_MAPPING = {
|
|
"": str, # Columns in views sometimes have type = ''
|
|
"TEXT": str,
|
|
"BLOB": bytes,
|
|
"INTEGER": int,
|
|
"FLOAT": float,
|
|
}
|
|
|
|
|
|
class AlterError(Exception):
|
|
pass
|
|
|
|
|
|
class NoObviousTable(Exception):
|
|
pass
|
|
|
|
|
|
class BadPrimaryKey(Exception):
|
|
pass
|
|
|
|
|
|
class NotFoundError(Exception):
|
|
pass
|
|
|
|
|
|
class Database:
|
|
def __init__(self, filename_or_conn=None, memory=False):
|
|
assert (filename_or_conn is not None and not memory) or (
|
|
filename_or_conn is None and memory
|
|
), "Either specify a filename_or_conn or pass memory=True"
|
|
if memory:
|
|
self.conn = sqlite3.connect(":memory:")
|
|
elif isinstance(filename_or_conn, str):
|
|
self.conn = sqlite3.connect(filename_or_conn)
|
|
elif isinstance(filename_or_conn, pathlib.Path):
|
|
self.conn = sqlite3.connect(str(filename_or_conn))
|
|
else:
|
|
self.conn = filename_or_conn
|
|
|
|
def __getitem__(self, table_name):
|
|
return self.table(table_name)
|
|
|
|
def __repr__(self):
|
|
return "<Database {}>".format(self.conn)
|
|
|
|
def table(self, table_name, **kwargs):
|
|
klass = View if table_name in self.view_names() else Table
|
|
return klass(self, table_name, **kwargs)
|
|
|
|
def escape(self, value):
|
|
# Normally we would use .execute(sql, [params]) for escaping, but
|
|
# occasionally that isn't available - most notable when we need
|
|
# to include a "... DEFAULT 'value'" in a column definition.
|
|
return self.conn.execute(
|
|
# Use SQLite itself to correctly escape this string:
|
|
"SELECT quote(:value)",
|
|
{"value": value},
|
|
).fetchone()[0]
|
|
|
|
def table_names(self, fts4=False, fts5=False):
|
|
where = ["type = 'table'"]
|
|
if fts4:
|
|
where.append("sql like '%FTS4%'")
|
|
if fts5:
|
|
where.append("sql like '%FTS5%'")
|
|
sql = "select name from sqlite_master where {}".format(" AND ".join(where))
|
|
return [r[0] for r in self.conn.execute(sql).fetchall()]
|
|
|
|
def view_names(self):
|
|
return [
|
|
r[0]
|
|
for r in self.conn.execute(
|
|
"select name from sqlite_master where type = 'view'"
|
|
).fetchall()
|
|
]
|
|
|
|
@property
|
|
def tables(self):
|
|
return [self[name] for name in self.table_names()]
|
|
|
|
@property
|
|
def views(self):
|
|
return [self[name] for name in self.view_names()]
|
|
|
|
def execute_returning_dicts(self, sql, params=None):
|
|
cursor = self.conn.execute(sql, params or tuple())
|
|
keys = [d[0] for d in cursor.description]
|
|
return [dict(zip(keys, row)) for row in cursor.fetchall()]
|
|
|
|
def resolve_foreign_keys(self, name, foreign_keys):
|
|
# foreign_keys may be a list of strcolumn names, a list of ForeignKey tuples,
|
|
# a list of tuple-pairs or a list of tuple-triples. We want to turn
|
|
# it into a list of ForeignKey tuples
|
|
if all(isinstance(fk, ForeignKey) for fk in foreign_keys):
|
|
return foreign_keys
|
|
if all(isinstance(fk, str) for fk in foreign_keys):
|
|
# It's a list of columns
|
|
fks = []
|
|
for column in foreign_keys:
|
|
other_table = self[name].guess_foreign_table(column)
|
|
other_column = self[name].guess_foreign_column(other_table)
|
|
fks.append(ForeignKey(name, column, other_table, other_column))
|
|
return fks
|
|
assert all(
|
|
isinstance(fk, (tuple, list)) for fk in foreign_keys
|
|
), "foreign_keys= should be a list of tuples"
|
|
fks = []
|
|
for tuple_or_list in foreign_keys:
|
|
assert len(tuple_or_list) in (
|
|
2,
|
|
3,
|
|
), "foreign_keys= should be a list of tuple pairs or triples"
|
|
if len(tuple_or_list) == 3:
|
|
fks.append(
|
|
ForeignKey(
|
|
name, tuple_or_list[0], tuple_or_list[1], tuple_or_list[2]
|
|
)
|
|
)
|
|
else:
|
|
# Guess the primary key
|
|
fks.append(
|
|
ForeignKey(
|
|
name,
|
|
tuple_or_list[0],
|
|
tuple_or_list[1],
|
|
self[name].guess_foreign_column(tuple_or_list[1]),
|
|
)
|
|
)
|
|
return fks
|
|
|
|
def create_table(
|
|
self,
|
|
name,
|
|
columns,
|
|
pk=None,
|
|
foreign_keys=None,
|
|
column_order=None,
|
|
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 {}
|
|
assert all(
|
|
n in columns for n in not_null
|
|
), "not_null set {} includes items not in columns {}".format(
|
|
repr(not_null), repr(set(columns.keys()))
|
|
)
|
|
assert all(
|
|
n in columns for n in defaults
|
|
), "defaults set {} includes items not in columns {}".format(
|
|
repr(set(defaults)), repr(set(columns.keys()))
|
|
)
|
|
column_items = list(columns.items())
|
|
if column_order is not None:
|
|
column_items.sort(
|
|
key=lambda p: column_order.index(p[0]) if p[0] in column_order else 999
|
|
)
|
|
if hash_id:
|
|
column_items.insert(0, (hash_id, str))
|
|
pk = hash_id
|
|
# Sanity check foreign_keys point to existing tables
|
|
for fk in foreign_keys:
|
|
if not any(
|
|
c for c in self[fk.other_table].columns if c.name == fk.other_column
|
|
):
|
|
raise AlterError(
|
|
"No such column: {}.{}".format(fk.other_table, fk.other_column)
|
|
)
|
|
|
|
column_defs = []
|
|
# ensure pk is a tuple
|
|
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:
|
|
column_extras.append("PRIMARY KEY")
|
|
if column_name in not_null:
|
|
column_extras.append("NOT NULL")
|
|
if column_name in defaults:
|
|
column_extras.append(
|
|
"DEFAULT {}".format(self.escape(defaults[column_name]))
|
|
)
|
|
if column_name in foreign_keys_by_column:
|
|
column_extras.append(
|
|
"REFERENCES [{other_table}]([{other_column}])".format(
|
|
other_table=foreign_keys_by_column[column_name].other_table,
|
|
other_column=foreign_keys_by_column[column_name].other_column,
|
|
)
|
|
)
|
|
column_defs.append(
|
|
" [{column_name}] {column_type}{column_extras}".format(
|
|
column_name=column_name,
|
|
column_type=COLUMN_TYPE_MAPPING[column_type],
|
|
column_extras=(" " + " ".join(column_extras))
|
|
if column_extras
|
|
else "",
|
|
)
|
|
)
|
|
extra_pk = ""
|
|
if single_pk is None and pk and len(pk) > 1:
|
|
extra_pk = ",\n PRIMARY KEY ({pks})".format(
|
|
pks=", ".join(["[{}]".format(p) for p in pk])
|
|
)
|
|
columns_sql = ",\n".join(column_defs)
|
|
sql = """CREATE TABLE [{table}] (
|
|
{columns_sql}{extra_pk}
|
|
);
|
|
""".format(
|
|
table=name, columns_sql=columns_sql, extra_pk=extra_pk
|
|
)
|
|
self.conn.execute(sql)
|
|
return self.table(
|
|
name,
|
|
pk=pk,
|
|
foreign_keys=foreign_keys,
|
|
column_order=column_order,
|
|
not_null=not_null,
|
|
defaults=defaults,
|
|
hash_id=hash_id,
|
|
)
|
|
|
|
def create_view(self, name, sql):
|
|
self.conn.execute(
|
|
"""
|
|
CREATE VIEW {name} AS {sql}
|
|
""".format(
|
|
name=name, sql=sql
|
|
)
|
|
)
|
|
|
|
def m2m_table_candidates(self, table, other_table):
|
|
"Returns potential m2m tables for arguments, based on FKs"
|
|
candidates = []
|
|
tables = {table, other_table}
|
|
for table in self.tables:
|
|
# Does it have foreign keys to both table and other_table?
|
|
has_fks_to = {fk.other_table for fk in table.foreign_keys}
|
|
if has_fks_to.issuperset(tables):
|
|
candidates.append(table.name)
|
|
return candidates
|
|
|
|
def add_foreign_keys(self, foreign_keys):
|
|
# foreign_keys is a list of explicit 4-tuples
|
|
assert all(
|
|
len(fk) == 4 and isinstance(fk, (list, tuple)) for fk in foreign_keys
|
|
), "foreign_keys must be a list of 4-tuples, (table, column, other_table, other_column)"
|
|
|
|
foreign_keys_to_create = []
|
|
|
|
# Verify that all tables and columns exist
|
|
for table, column, other_table, other_column in foreign_keys:
|
|
if not self[table].exists:
|
|
raise AlterError("No such table: {}".format(table))
|
|
if column not in self[table].columns_dict:
|
|
raise AlterError("No such column: {} in {}".format(column, table))
|
|
if not self[other_table].exists:
|
|
raise AlterError("No such other_table: {}".format(other_table))
|
|
if (
|
|
other_column != "rowid"
|
|
and other_column not in self[other_table].columns_dict
|
|
):
|
|
raise AlterError(
|
|
"No such other_column: {} in {}".format(other_column, other_table)
|
|
)
|
|
# We will silently skip foreign keys that exist already
|
|
if not any(
|
|
fk
|
|
for fk in self[table].foreign_keys
|
|
if fk.column == column
|
|
and fk.other_table == other_table
|
|
and fk.other_column == other_column
|
|
):
|
|
foreign_keys_to_create.append(
|
|
(table, column, other_table, other_column)
|
|
)
|
|
|
|
# Construct SQL for use with "UPDATE sqlite_master SET sql = ? WHERE name = ?"
|
|
table_sql = {}
|
|
for table, column, other_table, other_column in foreign_keys_to_create:
|
|
old_sql = table_sql.get(table, self[table].schema)
|
|
extra_sql = ",\n FOREIGN KEY({column}) REFERENCES {other_table}({other_column})\n".format(
|
|
column=column, other_table=other_table, other_column=other_column
|
|
)
|
|
# Stick that bit in at the very end just before the closing ')'
|
|
last_paren = old_sql.rindex(")")
|
|
new_sql = old_sql[:last_paren].strip() + extra_sql + old_sql[last_paren:]
|
|
table_sql[table] = new_sql
|
|
|
|
# And execute it all within a single transaction
|
|
with self.conn:
|
|
cursor = self.conn.cursor()
|
|
schema_version = cursor.execute("PRAGMA schema_version").fetchone()[0]
|
|
cursor.execute("PRAGMA writable_schema = 1")
|
|
for table_name, new_sql in table_sql.items():
|
|
cursor.execute(
|
|
"UPDATE sqlite_master SET sql = ? WHERE name = ?",
|
|
(new_sql, table_name),
|
|
)
|
|
cursor.execute("PRAGMA schema_version = %d" % (schema_version + 1))
|
|
cursor.execute("PRAGMA writable_schema = 0")
|
|
# Have to VACUUM outside the transaction to ensure .foreign_keys property
|
|
# can see the newly created foreign key.
|
|
self.vacuum()
|
|
|
|
def index_foreign_keys(self):
|
|
for table_name in self.table_names():
|
|
table = self[table_name]
|
|
existing_indexes = {
|
|
i.columns[0] for i in table.indexes if len(i.columns) == 1
|
|
}
|
|
for fk in table.foreign_keys:
|
|
if fk.column not in existing_indexes:
|
|
table.create_index([fk.column])
|
|
|
|
def vacuum(self):
|
|
self.conn.execute("VACUUM;")
|
|
|
|
|
|
class Queryable:
|
|
exists = False
|
|
|
|
def __init__(self, db, name):
|
|
self.db = db
|
|
self.name = name
|
|
|
|
@property
|
|
def count(self):
|
|
return self.db.conn.execute(
|
|
"select count(*) from [{}]".format(self.name)
|
|
).fetchone()[0]
|
|
|
|
@property
|
|
def rows(self):
|
|
return self.rows_where()
|
|
|
|
def rows_where(self, where=None, where_args=None):
|
|
if not self.exists:
|
|
return []
|
|
sql = "select * from [{}]".format(self.name)
|
|
if where is not None:
|
|
sql += " where " + where
|
|
cursor = self.db.conn.execute(sql, where_args or [])
|
|
columns = [c[0] for c in cursor.description]
|
|
for row in cursor:
|
|
yield dict(zip(columns, row))
|
|
|
|
@property
|
|
def columns(self):
|
|
if not self.exists:
|
|
return []
|
|
rows = self.db.conn.execute(
|
|
"PRAGMA table_info([{}])".format(self.name)
|
|
).fetchall()
|
|
return [Column(*row) for row in rows]
|
|
|
|
@property
|
|
def columns_dict(self):
|
|
"Returns {column: python-type} dictionary"
|
|
return {
|
|
column.name: REVERSE_COLUMN_TYPE_MAPPING[column.type]
|
|
for column in self.columns
|
|
}
|
|
|
|
@property
|
|
def schema(self):
|
|
return self.db.conn.execute(
|
|
"select sql from sqlite_master where name = ?", (self.name,)
|
|
).fetchone()[0]
|
|
|
|
|
|
class Table(Queryable):
|
|
def __init__(
|
|
self,
|
|
db,
|
|
name,
|
|
pk=None,
|
|
foreign_keys=None,
|
|
column_order=None,
|
|
not_null=None,
|
|
defaults=None,
|
|
upsert=False,
|
|
batch_size=100,
|
|
hash_id=None,
|
|
alter=False,
|
|
ignore=False,
|
|
extracts=None,
|
|
):
|
|
super().__init__(db, name)
|
|
self.exists = self.name in self.db.table_names()
|
|
self._defaults = dict(
|
|
pk=pk,
|
|
foreign_keys=foreign_keys,
|
|
column_order=column_order,
|
|
not_null=not_null,
|
|
defaults=defaults,
|
|
upsert=upsert,
|
|
batch_size=batch_size,
|
|
hash_id=hash_id,
|
|
alter=alter,
|
|
ignore=ignore,
|
|
extracts=extracts,
|
|
)
|
|
|
|
def __repr__(self):
|
|
return "<Table {}{}>".format(
|
|
self.name,
|
|
" (does not exist yet)"
|
|
if not self.exists
|
|
else " ({})".format(", ".join(c.name for c in self.columns)),
|
|
)
|
|
|
|
@property
|
|
def pks(self):
|
|
names = [column.name for column in self.columns if column.is_pk]
|
|
if not names:
|
|
names = ["rowid"]
|
|
return names
|
|
|
|
def get(self, pk_values):
|
|
if not isinstance(pk_values, (list, tuple)):
|
|
pk_values = [pk_values]
|
|
pks = self.pks
|
|
last_pk = pk_values[0] if len(pks) == 1 else pk_values
|
|
if len(pks) != len(pk_values):
|
|
raise NotFoundError(
|
|
"Need {} primary key value{}".format(
|
|
len(pks), "" if len(pks) == 1 else "s"
|
|
)
|
|
)
|
|
|
|
wheres = ["[{}] = ?".format(pk_name) for pk_name in pks]
|
|
rows = self.rows_where(" and ".join(wheres), pk_values)
|
|
try:
|
|
row = list(rows)[0]
|
|
self.last_pk = last_pk
|
|
return row
|
|
except IndexError:
|
|
raise NotFoundError
|
|
|
|
@property
|
|
def foreign_keys(self):
|
|
fks = []
|
|
for row in self.db.conn.execute(
|
|
"PRAGMA foreign_key_list([{}])".format(self.name)
|
|
).fetchall():
|
|
if row is not None:
|
|
id, seq, table_name, from_, to_, on_update, on_delete, match = row
|
|
fks.append(
|
|
ForeignKey(
|
|
table=self.name,
|
|
column=from_,
|
|
other_table=table_name,
|
|
other_column=to_,
|
|
)
|
|
)
|
|
return fks
|
|
|
|
@property
|
|
def indexes(self):
|
|
sql = 'PRAGMA index_list("{}")'.format(self.name)
|
|
indexes = []
|
|
for row in self.db.execute_returning_dicts(sql):
|
|
index_name = row["name"]
|
|
index_name_quoted = (
|
|
'"{}"'.format(index_name)
|
|
if not index_name.startswith('"')
|
|
else index_name
|
|
)
|
|
column_sql = "PRAGMA index_info({})".format(index_name_quoted)
|
|
columns = []
|
|
for seqno, cid, name in self.db.conn.execute(column_sql).fetchall():
|
|
columns.append(name)
|
|
row["columns"] = columns
|
|
# These columns may be missing on older SQLite versions:
|
|
for key, default in {"origin": "c", "partial": 0}.items():
|
|
if key not in row:
|
|
row[key] = default
|
|
indexes.append(Index(**row))
|
|
return indexes
|
|
|
|
def create(
|
|
self,
|
|
columns,
|
|
pk=None,
|
|
foreign_keys=None,
|
|
column_order=None,
|
|
not_null=None,
|
|
defaults=None,
|
|
hash_id=None,
|
|
extracts=None,
|
|
):
|
|
columns = {name: value for (name, value) in columns.items()}
|
|
with self.db.conn:
|
|
self.db.create_table(
|
|
self.name,
|
|
columns,
|
|
pk=pk,
|
|
foreign_keys=foreign_keys,
|
|
column_order=column_order,
|
|
not_null=not_null,
|
|
defaults=defaults,
|
|
hash_id=hash_id,
|
|
extracts=extracts,
|
|
)
|
|
self.exists = True
|
|
return self
|
|
|
|
def create_index(self, columns, index_name=None, unique=False, if_not_exists=False):
|
|
if index_name is None:
|
|
index_name = "idx_{}_{}".format(
|
|
self.name.replace(" ", "_"), "_".join(columns)
|
|
)
|
|
sql = """
|
|
CREATE {unique}INDEX {if_not_exists}{index_name}
|
|
ON {table_name} ({columns});
|
|
""".format(
|
|
index_name=index_name,
|
|
table_name=self.name,
|
|
columns=", ".join(columns),
|
|
unique="UNIQUE " if unique else "",
|
|
if_not_exists="IF NOT EXISTS " if if_not_exists else "",
|
|
)
|
|
self.db.conn.execute(sql)
|
|
return self
|
|
|
|
def add_column(
|
|
self, col_name, col_type=None, fk=None, fk_col=None, not_null_default=None
|
|
):
|
|
fk_col_type = None
|
|
if fk is not None:
|
|
# fk must be a valid table
|
|
if not fk in self.db.table_names():
|
|
raise AlterError("table '{}' does not exist".format(fk))
|
|
# if fk_col specified, must be a valid column
|
|
if fk_col is not None:
|
|
if fk_col not in self.db[fk].columns_dict:
|
|
raise AlterError("table '{}' has no column {}".format(fk, fk_col))
|
|
else:
|
|
# automatically set fk_col to first primary_key of fk table
|
|
pks = [c for c in self.db[fk].columns if c.is_pk]
|
|
if pks:
|
|
fk_col = pks[0].name
|
|
fk_col_type = pks[0].type
|
|
else:
|
|
fk_col = "rowid"
|
|
fk_col_type = "INTEGER"
|
|
if col_type is None:
|
|
col_type = str
|
|
not_null_sql = None
|
|
if not_null_default is not None:
|
|
not_null_sql = "NOT NULL DEFAULT {}".format(
|
|
self.db.escape(not_null_default)
|
|
)
|
|
sql = "ALTER TABLE [{table}] ADD COLUMN [{col_name}] {col_type}{not_null_default};".format(
|
|
table=self.name,
|
|
col_name=col_name,
|
|
col_type=fk_col_type or COLUMN_TYPE_MAPPING[col_type],
|
|
not_null_default=(" " + not_null_sql) if not_null_sql else "",
|
|
)
|
|
self.db.conn.execute(sql)
|
|
if fk is not None:
|
|
self.add_foreign_key(col_name, fk, fk_col)
|
|
return self
|
|
|
|
def drop(self):
|
|
self.db.conn.execute("DROP TABLE {}".format(self.name))
|
|
|
|
def guess_foreign_table(self, column):
|
|
column = column.lower()
|
|
possibilities = [column]
|
|
if column.endswith("_id"):
|
|
column_without_id = column[:-3]
|
|
possibilities.append(column_without_id)
|
|
if not column_without_id.endswith("s"):
|
|
possibilities.append(column_without_id + "s")
|
|
elif not column.endswith("s"):
|
|
possibilities.append(column + "s")
|
|
existing_tables = {t.lower(): t for t in self.db.table_names()}
|
|
for table in possibilities:
|
|
if table in existing_tables:
|
|
return existing_tables[table]
|
|
# If we get here there's no obvious candidate - raise an error
|
|
raise NoObviousTable(
|
|
"No obvious foreign key table for column '{}' - tried {}".format(
|
|
column, repr(possibilities)
|
|
)
|
|
)
|
|
|
|
def guess_foreign_column(self, other_table):
|
|
pks = [c for c in self.db[other_table].columns if c.is_pk]
|
|
if len(pks) != 1:
|
|
raise BadPrimaryKey(
|
|
"Could not detect single primary key for table '{}'".format(other_table)
|
|
)
|
|
else:
|
|
return pks[0].name
|
|
|
|
def add_foreign_key(self, column, other_table=None, other_column=None):
|
|
# Ensure column exists
|
|
if column not in self.columns_dict:
|
|
raise AlterError("No such column: {}".format(column))
|
|
# If other_table is not specified, attempt to guess it from the column
|
|
if other_table is None:
|
|
other_table = self.guess_foreign_table(column)
|
|
# If other_column is not specified, detect the primary key on other_table
|
|
if other_column is None:
|
|
other_column = self.guess_foreign_column(other_table)
|
|
|
|
# Sanity check that the other column exists
|
|
if (
|
|
not [c for c in self.db[other_table].columns if c.name == other_column]
|
|
and other_column != "rowid"
|
|
):
|
|
raise AlterError("No such column: {}.{}".format(other_table, other_column))
|
|
# Check we do not already have an existing foreign key
|
|
if any(
|
|
fk
|
|
for fk in self.foreign_keys
|
|
if fk.column == column
|
|
and fk.other_table == other_table
|
|
and fk.other_column == other_column
|
|
):
|
|
raise AlterError(
|
|
"Foreign key already exists for {} => {}.{}".format(
|
|
column, other_table, other_column
|
|
)
|
|
)
|
|
self.db.add_foreign_keys([(self.name, column, other_table, other_column)])
|
|
|
|
def enable_fts(self, columns, fts_version="FTS5"):
|
|
"Enables FTS on the specified columns"
|
|
sql = """
|
|
CREATE VIRTUAL TABLE "{table}_fts" USING {fts_version} (
|
|
{columns},
|
|
content="{table}"
|
|
);
|
|
""".format(
|
|
table=self.name,
|
|
columns=", ".join("[{}]".format(c) for c in columns),
|
|
fts_version=fts_version,
|
|
)
|
|
self.db.conn.executescript(sql)
|
|
self.populate_fts(columns)
|
|
return self
|
|
|
|
def populate_fts(self, columns):
|
|
sql = """
|
|
INSERT INTO "{table}_fts" (rowid, {columns})
|
|
SELECT rowid, {columns} FROM {table};
|
|
""".format(
|
|
table=self.name, columns=", ".join(columns)
|
|
)
|
|
self.db.conn.executescript(sql)
|
|
return self
|
|
|
|
def detect_fts(self):
|
|
"Detect if table has a corresponding FTS virtual table and return it"
|
|
rows = self.db.conn.execute(
|
|
"""
|
|
SELECT name FROM sqlite_master
|
|
WHERE rootpage = 0
|
|
AND (
|
|
sql LIKE '%VIRTUAL TABLE%USING FTS%content="{table}"%'
|
|
OR (
|
|
tbl_name = "{table}"
|
|
AND sql LIKE '%VIRTUAL TABLE%USING FTS%'
|
|
)
|
|
)
|
|
""".format(
|
|
table=self.name
|
|
)
|
|
).fetchall()
|
|
if len(rows) == 0:
|
|
return None
|
|
else:
|
|
return rows[0][0]
|
|
|
|
def optimize(self):
|
|
fts_table = self.detect_fts()
|
|
if fts_table is not None:
|
|
self.db.conn.execute(
|
|
"""
|
|
INSERT INTO [{table}] ([{table}]) VALUES ("optimize");
|
|
""".format(
|
|
table=fts_table
|
|
)
|
|
)
|
|
return self
|
|
|
|
def detect_column_types(self, records):
|
|
all_column_types = {}
|
|
for record in records:
|
|
for key, value in record.items():
|
|
all_column_types.setdefault(key, set()).add(type(value))
|
|
column_types = {}
|
|
for key, types in all_column_types.items():
|
|
if len(types) == 1:
|
|
t = list(types)[0]
|
|
# But if it's list / tuple / dict, use str instead as we
|
|
# will be storing it as JSON in the table
|
|
if t in (list, tuple, dict):
|
|
t = str
|
|
elif {int, bool}.issuperset(types):
|
|
t = int
|
|
elif {int, float, bool}.issuperset(types):
|
|
t = float
|
|
elif {bytes, str}.issuperset(types):
|
|
t = bytes
|
|
else:
|
|
t = str
|
|
column_types[key] = t
|
|
return column_types
|
|
|
|
def search(self, q):
|
|
sql = """
|
|
select * from {table} where rowid in (
|
|
select rowid from [{table}_fts]
|
|
where [{table}_fts] match :search
|
|
)
|
|
order by rowid
|
|
""".format(
|
|
table=self.name
|
|
)
|
|
return self.db.conn.execute(sql, (q,)).fetchall()
|
|
|
|
def value_or_default(self, key, value):
|
|
return self._defaults[key] if value is DEFAULT else value
|
|
|
|
def update(self, pk_values, updates=None, alter=False):
|
|
updates = updates or {}
|
|
if not isinstance(pk_values, (list, tuple)):
|
|
pk_values = [pk_values]
|
|
# Sanity check that the record exists (raises error if not):
|
|
self.get(pk_values)
|
|
if not updates:
|
|
return self
|
|
args = []
|
|
sets = []
|
|
wheres = []
|
|
for key, value in updates.items():
|
|
sets.append("[{}] = ?".format(key))
|
|
args.append(value)
|
|
wheres = ["[{}] = ?".format(pk_name) for pk_name in self.pks]
|
|
args.extend(pk_values)
|
|
sql = "update [{table}] set {sets} where {wheres}".format(
|
|
table=self.name, sets=", ".join(sets), wheres=" and ".join(wheres)
|
|
)
|
|
with self.db.conn:
|
|
try:
|
|
rowcount = self.db.conn.execute(sql, args).rowcount
|
|
except OperationalError as e:
|
|
if alter and (" column" in e.args[0]):
|
|
# Attempt to add any missing columns, then try again
|
|
self.add_missing_columns([updates])
|
|
rowcount = self.db.conn.execute(sql, args).rowcount
|
|
else:
|
|
raise
|
|
|
|
# TODO: Test this works (rolls back) - use better exception:
|
|
assert rowcount == 1
|
|
self.last_pk = pk_values[0] if len(self.pks) == 1 else pk_values
|
|
return self
|
|
|
|
def insert(
|
|
self,
|
|
record,
|
|
pk=DEFAULT,
|
|
foreign_keys=DEFAULT,
|
|
column_order=DEFAULT,
|
|
not_null=DEFAULT,
|
|
defaults=DEFAULT,
|
|
upsert=DEFAULT,
|
|
hash_id=DEFAULT,
|
|
alter=DEFAULT,
|
|
ignore=DEFAULT,
|
|
extracts=DEFAULT,
|
|
):
|
|
return self.insert_all(
|
|
[record],
|
|
pk=pk,
|
|
foreign_keys=foreign_keys,
|
|
column_order=column_order,
|
|
not_null=not_null,
|
|
defaults=defaults,
|
|
upsert=upsert,
|
|
hash_id=hash_id,
|
|
alter=alter,
|
|
ignore=ignore,
|
|
extracts=extracts,
|
|
)
|
|
|
|
def insert_all(
|
|
self,
|
|
records,
|
|
pk=DEFAULT,
|
|
foreign_keys=DEFAULT,
|
|
column_order=DEFAULT,
|
|
not_null=DEFAULT,
|
|
defaults=DEFAULT,
|
|
upsert=DEFAULT,
|
|
batch_size=DEFAULT,
|
|
hash_id=DEFAULT,
|
|
alter=DEFAULT,
|
|
ignore=DEFAULT,
|
|
extracts=DEFAULT,
|
|
):
|
|
"""
|
|
Like .insert() but takes a list of records and ensures that the table
|
|
that it creates (if table does not exist) has columns for ALL of that
|
|
data
|
|
"""
|
|
pk = self.value_or_default("pk", pk)
|
|
foreign_keys = self.value_or_default("foreign_keys", foreign_keys)
|
|
column_order = self.value_or_default("column_order", column_order)
|
|
not_null = self.value_or_default("not_null", not_null)
|
|
defaults = self.value_or_default("defaults", defaults)
|
|
upsert = self.value_or_default("upsert", upsert)
|
|
batch_size = self.value_or_default("batch_size", batch_size)
|
|
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 (
|
|
ignore and upsert
|
|
), "Use either ignore=True or upsert=True, not both"
|
|
all_columns = None
|
|
first = True
|
|
# We can only handle a max of 999 variables in a SQL insert, so
|
|
# we need to adjust the batch_size down if we have too many cols
|
|
records = iter(records)
|
|
# Peek at first record to count its columns:
|
|
first_record = next(records)
|
|
num_columns = len(first_record.keys())
|
|
assert (
|
|
num_columns <= SQLITE_MAX_VARS
|
|
), "Rows can have a maximum of {} columns".format(SQLITE_MAX_VARS)
|
|
batch_size = max(1, min(batch_size, SQLITE_MAX_VARS // num_columns))
|
|
for chunk in chunks(itertools.chain([first_record], records), batch_size):
|
|
chunk = list(chunk)
|
|
if first:
|
|
if not self.exists:
|
|
# Use the first batch to derive the table names
|
|
self.create(
|
|
self.detect_column_types(chunk),
|
|
pk,
|
|
foreign_keys,
|
|
column_order=column_order,
|
|
not_null=not_null,
|
|
defaults=defaults,
|
|
hash_id=hash_id,
|
|
extracts=extracts,
|
|
)
|
|
all_columns = set()
|
|
for record in chunk:
|
|
all_columns.update(record.keys())
|
|
all_columns = list(sorted(all_columns))
|
|
if hash_id:
|
|
all_columns.insert(0, hash_id)
|
|
first = False
|
|
or_what = ""
|
|
if upsert:
|
|
or_what = "OR REPLACE "
|
|
elif ignore:
|
|
or_what = "OR IGNORE "
|
|
sql = """
|
|
INSERT {or_what}INTO [{table}] ({columns}) VALUES {rows};
|
|
""".format(
|
|
or_what=or_what,
|
|
table=self.name,
|
|
columns=", ".join("[{}]".format(c) for c in all_columns),
|
|
rows=", ".join(
|
|
"""
|
|
({placeholders})
|
|
""".format(
|
|
placeholders=", ".join(["?"] * len(all_columns))
|
|
)
|
|
for record in chunk
|
|
),
|
|
)
|
|
values = []
|
|
extracts = resolve_extracts(extracts)
|
|
for record in chunk:
|
|
record_values = []
|
|
for key in all_columns:
|
|
value = jsonify_if_needed(
|
|
record.get(key, None if key != hash_id else _hash(record))
|
|
)
|
|
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)
|
|
except OperationalError as e:
|
|
if alter and (" column" in e.args[0]):
|
|
# Attempt to add any missing columns, then try again
|
|
self.add_missing_columns(chunk)
|
|
result = self.db.conn.execute(sql, values)
|
|
else:
|
|
raise
|
|
self.last_rowid = result.lastrowid
|
|
self.last_pk = self.last_rowid
|
|
# self.last_rowid will be 0 if a "INSERT OR IGNORE" happened
|
|
if (hash_id or pk) and self.last_rowid:
|
|
row = list(self.rows_where("rowid = ?", [self.last_rowid]))[0]
|
|
if hash_id:
|
|
self.last_pk = row[hash_id]
|
|
elif isinstance(pk, str):
|
|
self.last_pk = row[pk]
|
|
else:
|
|
self.last_pk = tuple(row[p] for p in pk)
|
|
return self
|
|
|
|
def upsert(
|
|
self,
|
|
record,
|
|
pk=DEFAULT,
|
|
foreign_keys=DEFAULT,
|
|
column_order=DEFAULT,
|
|
not_null=DEFAULT,
|
|
defaults=DEFAULT,
|
|
hash_id=DEFAULT,
|
|
alter=DEFAULT,
|
|
extracts=DEFAULT,
|
|
):
|
|
return self.insert(
|
|
record,
|
|
pk=pk,
|
|
foreign_keys=foreign_keys,
|
|
column_order=column_order,
|
|
not_null=not_null,
|
|
defaults=defaults,
|
|
hash_id=hash_id,
|
|
alter=alter,
|
|
upsert=True,
|
|
extracts=extracts,
|
|
)
|
|
|
|
def upsert_all(
|
|
self,
|
|
records,
|
|
pk=DEFAULT,
|
|
foreign_keys=DEFAULT,
|
|
column_order=DEFAULT,
|
|
not_null=DEFAULT,
|
|
defaults=DEFAULT,
|
|
batch_size=DEFAULT,
|
|
hash_id=DEFAULT,
|
|
alter=DEFAULT,
|
|
extracts=DEFAULT,
|
|
):
|
|
return self.insert_all(
|
|
records,
|
|
pk=pk,
|
|
foreign_keys=foreign_keys,
|
|
column_order=column_order,
|
|
not_null=not_null,
|
|
defaults=defaults,
|
|
batch_size=100,
|
|
hash_id=hash_id,
|
|
alter=alter,
|
|
upsert=True,
|
|
extracts=extracts,
|
|
)
|
|
|
|
def add_missing_columns(self, records):
|
|
needed_columns = self.detect_column_types(records)
|
|
current_columns = self.columns_dict
|
|
for col_name, col_type in needed_columns.items():
|
|
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
|
|
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]
|
|
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 m2m(
|
|
self, other_table, record_or_list=None, pk=DEFAULT, lookup=None, m2m_table=None
|
|
):
|
|
if isinstance(other_table, str):
|
|
other_table = self.db.table(other_table, pk=pk)
|
|
our_id = self.last_pk
|
|
if lookup is not None:
|
|
assert record_or_list is None, "Provide lookup= or record, not both"
|
|
else:
|
|
assert record_or_list is not None, "Provide lookup= or record, not both"
|
|
tables = list(sorted([self.name, other_table.name]))
|
|
columns = ["{}_id".format(t) for t in tables]
|
|
if m2m_table is not None:
|
|
m2m_table_name = m2m_table
|
|
else:
|
|
# Detect if there is a single, unambiguous option
|
|
candidates = self.db.m2m_table_candidates(self.name, other_table.name)
|
|
if len(candidates) == 1:
|
|
m2m_table_name = candidates[0]
|
|
elif len(candidates) > 1:
|
|
raise NoObviousTable(
|
|
"No single obvious m2m table for {}, {} - use m2m_table= parameter".format(
|
|
self.name, other_table.name
|
|
)
|
|
)
|
|
else:
|
|
# If not, create a new table
|
|
m2m_table_name = m2m_table or "{}_{}".format(*tables)
|
|
m2m_table = self.db.table(m2m_table_name, pk=columns, foreign_keys=columns)
|
|
if lookup is None:
|
|
records = (
|
|
[record_or_list]
|
|
if not isinstance(record_or_list, (list, tuple))
|
|
else record_or_list
|
|
)
|
|
# Ensure each record exists in other table
|
|
for record in records:
|
|
id = other_table.upsert(record, pk=pk).last_pk
|
|
m2m_table.upsert(
|
|
{
|
|
"{}_id".format(other_table.name): id,
|
|
"{}_id".format(self.name): our_id,
|
|
}
|
|
)
|
|
else:
|
|
id = other_table.lookup(lookup)
|
|
m2m_table.upsert(
|
|
{
|
|
"{}_id".format(other_table.name): id,
|
|
"{}_id".format(self.name): our_id,
|
|
}
|
|
)
|
|
return self
|
|
|
|
|
|
class View(Queryable):
|
|
exists = True
|
|
|
|
def __repr__(self):
|
|
return "<View {} ({})>".format(
|
|
self.name, ", ".join(c.name for c in self.columns)
|
|
)
|
|
|
|
def drop(self):
|
|
self.db.conn.execute("DROP VIEW {}".format(self.name))
|
|
|
|
|
|
def chunks(sequence, size):
|
|
iterator = iter(sequence)
|
|
for item in iterator:
|
|
yield itertools.chain([item], itertools.islice(iterator, size - 1))
|
|
|
|
|
|
def jsonify_if_needed(value):
|
|
if isinstance(value, (dict, list, tuple)):
|
|
return json.dumps(value)
|
|
elif isinstance(value, (datetime.time, datetime.date, datetime.datetime)):
|
|
return value.isoformat()
|
|
else:
|
|
return value
|
|
|
|
|
|
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
|