sql_databases.py

# (c) Andrew Chen (https://github.com/achen1296)

import csv
import json
import os
import re
import sqlite3
from datetime import date, datetime
from functools import cache
from pathlib import Path
from typing import Any, Iterable, Mapping, Sequence

sqlite3.enable_callback_tracebacks(True)


def adapt_json_serializable(d: dict | list):
 return json.dumps(d)


sqlite3.register_adapter(dict, adapt_json_serializable)
sqlite3.register_adapter(list, adapt_json_serializable)


def convert_json(b: bytes):
 return json.loads(b)


sqlite3.register_converter("json", convert_json)
sqlite3.register_converter("dict", convert_json)
sqlite3.register_converter("object", convert_json)
sqlite3.register_converter("obj", convert_json)
sqlite3.register_converter("list", convert_json)
sqlite3.register_converter("array", convert_json)


def adapt_datetime(d: datetime | date):
 return d.isoformat()


sqlite3.register_adapter(datetime, adapt_datetime)
sqlite3.register_adapter(date, adapt_datetime)


def convert_datetime(b: bytes):
 return datetime.fromisoformat(b.decode())


sqlite3.register_converter("date", convert_datetime)
sqlite3.register_converter("time", convert_datetime)
sqlite3.register_converter("datetime", convert_datetime)


def adapt_bool(b: bool):
 return int(b)


sqlite3.register_adapter(bool, adapt_bool)


def convert_bool(b: bytes):
 """ Numbers: 0 is false and anything else is true.
 Text: keywords "true" or "false", case-insensitive with surrounding whitespace stripped, anything else results in `ValueError`. """
 try:
 return float(b) != 0.
 except ValueError:
 pass
 b = b.lower().strip()
 if b == b"true":
 return True
 if b == b"false":
 return False
 raise ValueError(b)


sqlite3.register_converter("bool", convert_bool)

sqlite3.register_converter("str", lambda b: b.decode())
sqlite3.register_converter("int", int)
sqlite3.register_converter("integer", int)
sqlite3.register_converter("float", float)
sqlite3.register_converter("real", float)


def convert_lenient_int(b: bytes):
 """ Find first instance of text convertbile to int (decimal only) and use that, discarding the rest. """
 m = re.search(b"(\\+|-)?\\d+", b)
 if not m:
 raise ValueError(b)
 return int(m.group(0))


# https://docs.python.org/3/library/functions.html#float
LENIENT_FLOAT_RE = re.compile(
 b"""(\\+|-)? # sign
 ( # value
 inf(inity)?|nan| # special value keywords
 (
 ( \\d*\\.\\d+ | \\d+\\.? ) # digits
 ( e # optional exponent
 (\\+|-)? # exponent sign
 \\d+ # exponent digits
 )?
 )
 )
 """,
 re.VERBOSE | re.I
)


def convert_lenient_float(b: bytes):
 """ Find first instance of text convertbile to float and use that, discarding the rest. """
 m = re.search(LENIENT_FLOAT_RE, b)
 if not m:
 raise ValueError(b)
 return float(m.group(0))


sqlite3.register_converter("lenient float", convert_lenient_float)
sqlite3.register_converter("lenient_float", convert_lenient_float)
sqlite3.register_converter("lenient real", convert_lenient_float)
sqlite3.register_converter("lenient_real", convert_lenient_float)
sqlite3.register_converter("lenient integer", convert_lenient_int)
sqlite3.register_converter("lenient_integer", convert_lenient_int)
sqlite3.register_converter("lenient int", convert_lenient_int)
sqlite3.register_converter("lenient_int", convert_lenient_int)


def col_str(column: str | tuple[str, str], table: str | None = None):
 if isinstance(column, str):
 c = f'"{column.lower()}"'
 else:
 c = f'"{column[0].lower()}" "{column[1]}"'
 if table is not None:
 return f'"{table}".{c}'
 else:
 return c


def col_name_type(column: str | tuple[str, str]) -> tuple[str, str]:
 if isinstance(column, str):
 return column.lower(), ""
 else:
 return (column[0].lower(), column[1])


def cols_strs(columns: Mapping[str, str] | Iterable[str | tuple[str, str]], table: str | None = None) -> list:
 """ `columns`: `Iterable` of either just the column name or `tuple` of the column's name and declared type, or `Mapping` of column name and type. """
 if isinstance(columns, Mapping):
 return [col_str((c, columns[c]), table=table) for c in columns] # type:ignore
 else:
 return [col_str(c, table=table) for c in columns]


def cols_names_types(columns: Mapping[str, str] | Iterable[str | tuple[str, str]]) -> Mapping[str, str]:
 if isinstance(columns, Mapping):
 return columns # type:ignore
 else:
 return {
 c: t
 for c, t in (col_name_type(col) for col in columns)
 }


def cols_joined_str(columns: Mapping[str, str] | Iterable[str | tuple[str, str]], table: str | None = None) -> str:
 return ",".join(cols_strs(columns, table=table))


class Row(sqlite3.Row):
 def __repr__(self):
 return repr(dict(**self))

 def __str__(self):
 return str(dict(**self))

 def __contains__(self, value):
 return value in self.keys()

 def get(self, key, default=None):
 try:
 return self[key]
 except KeyError:
 return default


def _add_connection_features(con: sqlite3.Connection):
 con.row_factory = Row
 con.create_function("regexp", 2, lambda p, s: bool(re.search(p, s or "", re.I)), deterministic=True)


def fetch_rows(db: "Database | Table | sqlite3.Connection", sql: str, parameters=()):
 if isinstance(db, Database) or isinstance(db, Table):
 con = db.con
 else:
 con = db
 # create a new cursor to ensure it's not used for something else, interrupting the query
 cur = con.cursor()
 cur.execute(sql, parameters)
 while (row := cur.fetchone()) is not None:
 yield row


class Database:
 def __init__(self, db_file: str | Path, *, timeout=60., autocommit=False):
 self.db_file = Path(db_file)
 self.con = sqlite3.connect(self.db_file, detect_types=sqlite3.PARSE_DECLTYPES | sqlite3.PARSE_COLNAMES, timeout=timeout, autocommit=autocommit)
 _add_connection_features(self.con)
 self.cur = self.con.cursor()

 def close(self):
 """ This might be needed for long-running programs. """
 self.con.close()

 def __enter__(self):
 return self

 def __exit__(self, exc_type, exc_value, traceback):
 self.close()

 @property
 def tables(self) -> list[str]:
 with self.con:
 return [name[0].lower() for name in self.cur.execute(""" select name from sqlite_schema where type = 'table' """).fetchall()]

 def create_table(self, name: str, columns: Mapping[str, str] | Iterable[str | tuple[str, str]], primary_keys: Iterable[str]):
 cs = cols_joined_str(columns)
 primary_keys_str = cols_joined_str(primary_keys)
 with self.con:
 self.cur.execute(f""" create table \"{name}\" ( {cs}, primary key ({primary_keys_str}) ) """)
 return self.table(name)

 @cache
 def table(self, name: str):
 return Table(self, name)

 def create_sql(self) -> list[str]:
 return [r[0] for r in self.cur.execute(f""" select sql from sqlite_schema where sql is not null """).fetchall()]

 def synchronize_definition_file(self, db_definition_file: Path | str):
 # the database file is not expected to be reconstructed often if at all, this code is mostly to document the intent of matching the saved table definitions in git
 db_definition_file = Path(db_definition_file)
 tables = self.tables
 if db_definition_file.exists():
 with open(db_definition_file) as f:
 for line in f:
 m = re.match("create table \"?(.*?)\"? ?\\(", line, re.I)
 if m:
 # otherwise could be e.g. a trigger or index
 t = m.group(1)
 if t not in tables:
 self.cur.execute(line)
 with open(db_definition_file, "w") as f:
 for sql in self.create_sql():
 print(sql, file=f)


class TableNotFound(Exception):
 pass


class ExtraData(Exception):
 pass


RowType = Mapping[str, Any] | Sequence | sqlite3.Row


class Table:
 # not worrying about SQL injection here
 def __init__(self, db: Database, name: str):
 self.db = db
 self.con = db.con
 self.cur = db.con.cursor()
 self.name = name

 self.altered_table = True
 self.columns # evaluate for existence check

 # _add_connection_features(self.con) # already done by db __init__

 def _cache_columns_and_types(self):
 if self.altered_table:
 with self.con:
 cols_and_types = self.cur.execute(""" select name, type, hidden from pragma_table_xinfo(?) """, (self.name, )).fetchall()
 if not cols_and_types:
 raise TableNotFound(self.name)
 # hidden == 0 are ordinary columns
 # 1 for hidden -- only these do not appear when using "select * from ..." thus they are the only ones excluded from "star columns"
 # 2/3 for dynamic/stored generated columns
 self._columns_and_types = tuple((c[0], c[1]) for c in cols_and_types if c[2] == 0)
 self._star_columns_and_types = tuple((c[0], c[1]) for c in cols_and_types if c[2] != 1)
 self._all_columns_and_types = tuple((c[0], c[1]) for c in cols_and_types)

 self.altered_table = False

 @property
 def columns(self) -> tuple[str, ...]:
 """ Does not include hidden and generated columns """
 self._cache_columns_and_types()
 return tuple(c[0] for c in self._columns_and_types)

 @property
 def column_types(self) -> tuple[str, ...]:
 """ Does not include hidden and generated columns """
 self._cache_columns_and_types()
 return tuple(c[1] for c in self._columns_and_types)

 @property
 def columns_and_types(self) -> tuple[tuple[str, str], ...]:
 """ Does not include hidden and generated columns """
 self._cache_columns_and_types()
 return self._columns_and_types

 @property
 def star_columns(self) -> tuple[str, ...]:
 """ Only those columns which appear when using "select * from ...", which are ordinary and generated columns but not hidden columns"""
 self._cache_columns_and_types()
 return tuple(c[0] for c in self._star_columns_and_types)

 @property
 def star_column_types(self) -> tuple[str, ...]:
 """ Only those columns which appear when using "select * from ...", which are ordinary and generated columns but not hidden columns"""
 self._cache_columns_and_types()
 return tuple(c[1] for c in self._star_columns_and_types)

 @property
 def star_columns_and_types(self) -> tuple[tuple[str, str], ...]:
 """ Only those columns which appear when using "select * from ...", which are ordinary and generated columns but not hidden columns"""
 self._cache_columns_and_types()
 return self._star_columns_and_types

 @property
 def all_columns(self) -> tuple[str, ...]:
 self._cache_columns_and_types()
 return tuple(c[0] for c in self._all_columns_and_types)

 @property
 def all_column_types(self) -> tuple[str, ...]:
 self._cache_columns_and_types()
 return tuple(c[1] for c in self._all_columns_and_types)

 @property
 def all_columns_and_types(self) -> tuple[tuple[str, str], ...]:
 self._cache_columns_and_types()
 return self._all_columns_and_types

 def add_columns(self, columns: Mapping[str, str] | Iterable[str | tuple[str, str]]):
 """ Adds columns, unless they are already in the table. Returns `True` if any new columns were added, `False` otherwise. """
 added_any = False
 existing_cols = [c.lower() for c in self.columns]
 for c, t in cols_names_types(columns).items():
 if c.lower() not in existing_cols:
 added_any = True
 with self.con:
 if t:
 self.cur.execute(f""" alter table "{self.name}" add column {col_str((c, t))} """)
 else:
 self.cur.execute(f""" alter table "{self.name}" add column {col_str(c)} """)
 if added_any:
 self.altered_table = True
 return added_any

 @property
 @cache # primary keys cannot be changed except by recreating the table
 def primary_keys(self) -> tuple[str]:
 with self.con:
 cols = tuple(name[0].lower() for name in self.cur.execute(""" select name from pragma_table_info(?) where pk > 0 """, (self.name, )).fetchall())
 if not cols:
 raise TableNotFound(self.name)
 return cols

 def _parse_row(self, row: RowType, *, add_missing_columns: bool, add_column_types: bool, ignore_extra_data: bool):
 if isinstance(row, Mapping) or isinstance(row, sqlite3.Row):
 lower_columns: list[str] | None = None
 keys = row.keys()
 if add_missing_columns:
 if add_column_types:
 cols = {
 c: type(row[c]).__name__
 for c in row.keys()
 }
 else:
 cols = [c for c in keys]
 self.altered_table = self.add_columns(cols)
 elif not ignore_extra_data:
 lower_columns = [c.lower() for c in self.columns]
 extra_keys = [c for c in keys if c.lower() not in lower_columns]
 if extra_keys:
 raise ExtraData(extra_keys)

 if lower_columns is None:
 lower_columns = [c.lower() for c in self.columns]
 # `in row` is keys for `Mapping` but values for `sqlite3.Row`
 # need to use the case of the keys as they are in `row` for retrieving them below
 operation_cols = [c for c in keys if c.lower() in lower_columns]
 params = [row[c] for c in operation_cols]
 else:
 lr = len(row)
 lc = len(self.columns)
 if not ignore_extra_data and lr > lc:
 raise ExtraData(row[lc:])
 operation_cols = self.columns[:lr]
 params = list(row[:lc])
 return operation_cols, params

 def insert(self, row: RowType, *, add_missing_columns=False, add_column_types=True, ignore_extra_data=False, upsert=False):
 """ If `add_missing_columns`, will add keys of a `row` that is a `Mapping` as new columns if one with the same name doesn't exist (SQLite columns are case-insensitive), and if `add_column_types`, will add declared column types using `type(v).__name__`. Else, if `ignore_extra_data`, ignores the additional keys, otherwise raise an exception.

 If `row` is a `Sequence` with length at most the number of columns, always succeeds. Otherwise, either ignores or raises an exception based on `ignore_extra_data`. Cannot add new columns this way because a name is not provided.

 Note: `sqlite3.Row` is treated as a `Mapping`, not a `Sequence`. It is designed such that it could be treated as either in many ways. """

 operation_cols, params = self._parse_row(row, add_missing_columns=add_missing_columns, add_column_types=add_column_types, ignore_extra_data=ignore_extra_data)
 self._insert(operation_cols, params, upsert=upsert)

 def _insert(self, operation_cols: Iterable[str], params: list, *, upsert: bool):
 sql = f""" insert into {self.name} ({cols_joined_str(operation_cols)}) values({",".join("?"*len(params))}) """
 if upsert:
 sql += f""" on conflict do update set ({cols_joined_str(operation_cols)}) = ({",".join("?"*len(params))}) """
 params = params + params

 with self.con:
 self.cur.execute(sql, params)

 def upsert(self, row: RowType, *, upsert=True, **kwargs,):
 """ See `insert`. `upsert` argument is just to absorb accidentally including this argument, always passed as `True` to `insert`. """
 return self.insert(row, upsert=True, **kwargs)

 def bulk_insert(self, rows: Iterable[RowType], *, add_missing_columns=False, add_column_types=True, ignore_extra_data=False, upsert=False):
 """ Translating from Python data to the database is slow. This method uses a temporary table to do that part, before performing the transfer to the target table inside of SQLite which is much faster, reducing time spent with the database locked.

 NOTE: It is possible specify different sets of columns for the rows. However, if one row specifies one or more columns that a second row does *not* specify, and the second row results in an upsert, then the missing values in the second row *will be updated to null instead of being ignored*. (Presumably, most of the time, all rows will have the same set of columns and this won't be an issue.) """
 self.cur.execute(""" drop table if exists bulk_insert_temp_table """)
 # precompute all the necessary columns
 parse_results = [self._parse_row(r, add_missing_columns=add_missing_columns, add_column_types=add_column_types, ignore_extra_data=ignore_extra_data) for r in rows]

 if len(parse_results) == 0:
 return 0 # otherwise syntax error with no columns

 all_operation_columns: set[str] = set()
 for operation_cols, _ in parse_results:
 for c in operation_cols:
 all_operation_columns.add(c)

 self.cur.execute(f""" create temp table bulk_insert_temp_table({cols_joined_str(all_operation_columns)}) """)
 temp_table = Table(self.db, "bulk_insert_temp_table")

 count = 0
 for operation_cols, params in parse_results:
 # upsert is not applicable with no primary key/uniqueness constraints on the temp table, and it will be applied later
 temp_table._insert(operation_cols, params, upsert=False)
 count += 1

 operation_cols_str = cols_joined_str(all_operation_columns)

 sql = f"""\
 insert into {self.name}
 ({operation_cols_str})

 select {operation_cols_str}
 from {temp_table.name}
 """

 if upsert:
 sql += f"""
 where true

 on conflict do update
 set ({operation_cols_str}) = ({",".join("excluded.\"" + c + "\"" for c in temp_table.columns)})
 """

 with self.con:
 self.cur.execute(sql)

 return count

 def bulk_upsert(self, rows: Iterable[RowType], *, upsert=True, **kwargs):
 return self.bulk_insert(rows, upsert=True, **kwargs)

 def import_csv(self, csv_file: Path | str, *, add_missing_columns: bool = False, ignore_extra_data=False, upsert=False):
 """ Cannot add types to columns this way, as CSV reader would of course always produce string values. Returns count of entries added. """
 with open(csv_file, newline="", encoding="utf-8-sig") as f: # encoding handles byte order mark
 reader = csv.DictReader(f)
 return self.bulk_insert(reader, add_missing_columns=add_missing_columns, ignore_extra_data=ignore_extra_data, upsert=upsert, add_column_types=False)

 def update(self, row: RowType, where: str, where_params=[], *, add_missing_columns: bool = False, add_column_types=True, ignore_extra_data=False):
 """ See `insert`. """
 operation_cols, params = self._parse_row(row, add_missing_columns=add_missing_columns, add_column_types=add_column_types, ignore_extra_data=ignore_extra_data)
 sql = f""" update {self.name} set ({cols_joined_str(operation_cols)}) = ({",".join("?"*len(params))}) where {where} """
 params += where_params
 with self.con:
 self.cur.execute(sql, params)

 def select(self, columns: Iterable[str] | None = None, where: str = "true", where_params=[], *, as_types: Mapping[str, str] = {}) -> list[Row]:
 """ Don't forget to `sqlite3.register_converter` if you use `as_types`! Some converters have already been registered for common Python built-in types. """
 if columns is None:
 columns = self.columns
 else:
 columns = [c.lower() for c in columns]
 as_types = {
 k.lower(): v
 for k, v in as_types.items()
 }
 col_str = ",".join(f"\"{c}\" as '{c} [{as_types[c]}]'" if c in as_types else f'"{c}"' for c in columns)

 with self.con:
 sql = f""" select {col_str} from {self.name} where {where} """
 return self.cur.execute(sql, where_params).fetchall()

 def delete(self, where: str = "true", where_params=[], ):
 with self.con:
 sql = f""" delete from "{self.name}" where {where} """
 self.cur.execute(sql, where_params)

 def __iter__(self):
 return iter(self.select())


DATABASE_IN_MEMORY = ":memory:"

if __name__ == "__main__":
 test_db = Database(DATABASE_IN_MEMORY)
 t = test_db.create_table("t", [("a", "int")], ["a"])
 assert t.columns == ("a",), t.columns
 assert t.primary_keys == ("a",), t.primary_keys

 try:
 t.insert({"a": 1, "b": 2})
 except ExtraData:
 pass
 else:
 assert False
 t.insert({"a": 1, "b": 2}, ignore_extra_data=True)

 try:
 t.insert([3, 4])
 except ExtraData:
 pass
 else:
 assert False
 t.insert([3, 4], ignore_extra_data=True)

 rows = list(dict(**r) for r in t)
 assert rows == [{"a": 1}, {"a": 3}], rows

 t.insert({"a": 5, "b": 6}, add_missing_columns=True, add_column_types=False)

 assert t.columns == ("a", "b"), t.columns

 rows = list(dict(**r) for r in t)
 assert rows == [{"a": 1, "b": None}, {"a": 3, "b": None}, {"a": 5, "b": 6}], rows

 t.upsert({"a": 1, "b": {"hello": "world"}})
 rows = list(dict(**r) for r in t)
 assert rows == [{"a": 1, "b": '{"hello": "world"}'}, {"a": 3, "b": None}, {"a": 5, "b": 6}], rows

 t.update((7, ["asdf"]), "a=1")
 rows = list(dict(**r) for r in t)
 assert rows == [{"a": 7, "b": '["asdf"]'}, {"a": 3, "b": None}, {"a": 5, "b": 6}], rows

 # new column should get list type
 t.insert({"a": 8, "c": []}, add_missing_columns=True)
 selected = [dict(**r) for r in t.select(where="a=8")]
 assert selected == [{"a": 8, "b": None, "c": []}], selected

 selected = [dict(**r) for r in t.select(where="a=8", as_types={"c": "str"})]
 assert selected == [{"a": 8, "b": None, "c": "[]"}], selected

 with open("test.csv", "w", newline="") as f:
 w = csv.writer(f)
 w.writerow(["a", "c", "d"])
 w.writerow([9, '["hello", "world"]', "hello"])
 w.writerow([7, '["bye", "world"]'])
 t.import_csv("test.csv", upsert=True, add_missing_columns=True)
 os.remove("test.csv")

 selected = [dict(**r) for r in t.select(where="a=9", as_types={"c": "list"})]
 assert selected == [{"a": 9, "b": None, "c": ["hello", "world"], "d": "hello"}], selected
 selected = [dict(**r) for r in t.select(where="a=7", as_types={"c": "list"})]
 # column b should be retained through upsert
 assert selected == [{"a": 7, "b": '["asdf"]', "c": ["bye", "world"], "d": None}], selected

 test_db.con.close()

 print("tests passed")