"""Implements ``Database`` class."""
from __future__ import annotations
import contextlib
import pathlib
import sqlite3
from collections.abc import Iterator
from typing import Any, Optional
from fractal.typing import Row
__all__ = ['Database']
[docs]
class Database:
"""Thin SQLite wrapper for the tree's central database.
Provides generic table operations -- ``init``, ``read``,
``write``, ``update``, ``merge``, ``delete``, ``exists``,
and ``count`` -- with no domain-specific logic. Each operation runs
on its own one-shot handle unless passed a :meth:`transaction`
connection via ``connection=``, which makes a multi-statement block
one atomic unit.
"""
_timeout = 30
def __init__(
self: Database,
path: pathlib.Path,
schema: pathlib.Path,
) -> None:
"""Initialize ``Database``.
Args:
path: Path to the ``.db`` file.
schema: Path to the ``.sql`` schema executed by ``init``.
"""
self._path = pathlib.Path(path).expanduser().resolve()
self._schema = pathlib.Path(schema).expanduser().resolve()
@property
def path(self: Database) -> pathlib.Path:
"""Return the database file path (read-only)."""
return self._path
[docs]
def init(self: Database) -> None:
"""Create the database and tables from the schema.
Idempotent -- safe to call on an existing database -- but purely
additive: the schema's ``IF NOT EXISTS`` DDL never alters an
existing table or refreshes a changed view, so a database created
under an older schema is not upgraded and must be rebuilt.
"""
sql = self._schema.read_text(encoding='utf-8')
connection = self.connect()
try:
# stamp WAL once at creation: the mode persists in the file
# header, and the delete-to-WAL transition takes an exclusive
# lock, so it must never run per-connect under a live fleet
connection.execute('PRAGMA journal_mode = WAL')
connection.executescript(sql)
# stamp the schema version on an unstamped database only, so a
# future migration mechanism can key on the stored version
version, *_ = connection.execute('PRAGMA user_version').fetchone()
if version == 0:
connection.execute('PRAGMA user_version = 1')
finally:
connection.close()
# backfill the schema into the main file: a fresh tree's .db
# must be self-contained at rest, not a bare header page whose
# tables live only in the -wal
self.checkpoint()
[docs]
def checkpoint(self: Database) -> None:
"""Backfill the WAL into the main file and truncate it.
Writable handles never checkpoint on close (the sidecars must stay
put for write-denied readers), and one-shot connections never
observe the full backfill that lets SQLite restart the log -- so
without an explicit consolidation the WAL only grows and the bare
``.db`` is incomplete at rest. TRUNCATE backfills and zeroes the
log while leaving both sidecars in place. Best-effort: a busy
fleet defers the consolidation to the next caller.
"""
connection = self.connect()
try:
connection.execute('PRAGMA wal_checkpoint(TRUNCATE)')
except sqlite3.OperationalError:
pass
finally:
connection.close()
[docs]
def connect(
self: Database,
*,
read_only: bool = False,
timeout: Optional[float] = None,
) -> sqlite3.Connection:
"""Open a database connection.
Enables foreign keys (the WAL journal mode is stamped once by
``init`` and persists in the file header) and gives every handle a
30-second busy timeout (overriding ``sqlite3.connect``'s 5-second
default) so a contending writer waits for the lock under wide node
fan-out instead of failing fast and aborting the loop. ``timeout``
overrides the default -- a UI-thread reader passes a short one so a
busy writer never stalls a refresh. Writable handles keep the
close-time checkpoint from unlinking the WAL sidecars. Each caller
is responsible for closing the returned connection.
Args:
read_only: Open in read-only mode.
timeout: Busy timeout in seconds (default 30).
Returns:
Database connection.
"""
# default the busy timeout to 30s (the fleet contends on one DB under
# fan-out); a caller's explicit timeout still wins
if timeout is None:
timeout = self._timeout
# create connection
if read_only:
# a read-only open of a missing file yields a cryptic "unable to open
# database file" -- raise an actionable error pointing at fractal init
if not self._path.exists():
raise FileNotFoundError(
f'No database at {self._path}; run `fractal init` at the repo root.'
)
uri = f'file:{self._path}?mode=ro'
connection = sqlite3.connect(uri, uri=True, timeout=timeout)
else:
connection = sqlite3.connect(f'{self._path}', timeout=timeout)
# configure connection -- fractal's own handles never unlink the
# -wal/-shm sidecars (a cooperative guarantee: a foreign writable
# client closing last can still unlink them once): a last close that
# checkpoints and removes them strands write-denied (sandboxed)
# readers on SQLITE_CANTOPEN until a writer rebuilds the wal-index,
# so writable handles disable the close-time checkpoint; backfill
# comes from wal_autocheckpoint (never zero it) plus the explicit
# loop-exit checkpoint()
if not read_only:
connection.setconfig(sqlite3.SQLITE_DBCONFIG_NO_CKPT_ON_CLOSE, True)
connection.execute('PRAGMA foreign_keys = ON')
connection.row_factory = sqlite3.Row
return connection
[docs]
@contextlib.contextmanager
def transaction(self: Database) -> Iterator[sqlite3.Connection]:
"""Open one ``BEGIN IMMEDIATE`` transaction over a dedicated connection.
Yields a connection holding the database write lock for the whole
block, so a read-decide-write transition commits or rolls back as
one atomic unit; the table operations join it via ``connection=``.
Every operation inside the block must pass it: a write without
``connection=`` deadlocks against the block's own lock until the
busy timeout, and a read without it silently sees the
pre-transaction snapshot. ``BEGIN IMMEDIATE`` takes the lock up
front -- two concurrent transactions serialize (bounded by the
busy timeout) instead of failing mid-block. Commits on clean
exit; rolls back on any exception (re-raised).
Yields:
The transaction's connection.
"""
connection = self.connect()
# explicit BEGIN/COMMIT -- disable the sqlite3 module's implicit
# transaction management for this handle
connection.isolation_level = None
try:
connection.execute('BEGIN IMMEDIATE')
yield connection
connection.commit()
except BaseException:
# BEGIN itself may be what failed (busy timeout), leaving no
# transaction to roll back
if connection.in_transaction:
connection.rollback()
raise
finally:
connection.close()
@contextlib.contextmanager
def _handle(
self: Database,
connection: Optional[sqlite3.Connection],
*,
read_only: bool = False,
) -> Iterator[sqlite3.Connection]:
"""Yield ``connection`` if given, else an owned one-shot handle.
An owned handle commits on clean exit and always closes; a caller's
connection is yielded untouched -- its transaction owns commit and
close, and reads on it see the transaction's uncommitted writes.
Args:
connection: Transaction to run inside (from
:meth:`transaction`), or ``None`` for an owned handle.
read_only: Open an owned handle read-only (no commit).
Yields:
The connection the operation runs on.
"""
if connection is not None:
yield connection
return
owned = self.connect(read_only=read_only)
try:
yield owned
if not read_only:
owned.commit()
finally:
owned.close()
[docs]
def read(
self: Database,
table: Optional[str] = None,
*,
query: Optional[str] = None,
params: tuple[Any, ...] = (),
where: Optional[dict[str, Any]] = None,
limit: Optional[int] = None,
connection: Optional[sqlite3.Connection] = None,
) -> list[Row]:
"""Read rows from a table.
Either ``query`` or ``table`` must be provided.
``query`` is mutually exclusive with ``where``
and ``limit``. When ``query`` is given, it is
executed as raw read-only SQL with ``params`` bound
positionally to its ``?`` placeholders. Otherwise, a
``SELECT`` is built from ``table``, ``where``, and
``limit``.
Built ``SELECT``s are ordered by ``rowid DESC`` -- the
true write order (a monotonic alias of the
``INTEGER PRIMARY KEY``) -- so the most recently written row is
first, without assuming any particular timestamp
column (tables name their start ``started_at`` or
``created_at``). A raw ``query`` keeps its own ordering
(none is imposed); ``limit=0`` is valid and returns no rows.
Args:
table: Table name.
query: Raw SQL query (mutually exclusive with
``where`` and ``limit``).
params: Parameters bound to the ``query``'s ``?``
placeholders.
where: Column filters (AND-joined equality).
limit: Maximum rows to return.
connection: Transaction to run inside (from
:meth:`transaction`).
Returns:
List of row dicts.
"""
# validate arguments
if query is not None and (where is not None or limit is not None):
raise ValueError('Cannot combine query with where or limit.')
if query is None and table is None:
raise ValueError('Either table or query is required.')
if limit is not None and limit < 0:
raise ValueError(f'limit must be non-negative, got {limit}.')
# build query -- the table/where path binds its own params
# (query and where are mutually exclusive, so the caller's
# params are untouched there)
if query is None:
query = f'SELECT * FROM {table}'
if where:
where_clause, where_params = self._where_clause(where)
query += ' WHERE ' + where_clause
params += where_params
query += ' ORDER BY rowid DESC'
if limit is not None:
query += ' LIMIT ?'
params += (limit,)
# execute query
with self._handle(connection, read_only=True) as handle:
cursor = handle.execute(query, params)
result = [dict(row) for row in cursor.fetchall()]
return result
[docs]
def write(
self: Database,
data: Row,
table: str,
*,
connection: Optional[sqlite3.Connection] = None,
) -> int:
"""Insert a row into a table.
Args:
data: Column values.
table: Table name.
connection: Transaction to run inside (from
:meth:`transaction`).
Returns:
Row ID of the inserted row.
"""
# build statement
columns = ', '.join(data.keys())
placeholders = ', '.join('?' for _ in data)
statement = f'INSERT INTO {table} ({columns}) VALUES ({placeholders})'
# execute statement
with self._handle(connection) as handle:
cursor = handle.execute(statement, tuple(data.values()))
row_id = cursor.lastrowid
return row_id
[docs]
def update(
self: Database,
data: Row,
table: str,
*,
where: dict[str, Any],
connection: Optional[sqlite3.Connection] = None,
) -> int:
"""Update rows in a table.
Args:
data: Column values to set.
table: Table name.
where: Column filters (AND-joined equality).
connection: Transaction to run inside (from
:meth:`transaction`).
Returns:
Number of rows updated -- the observable outcome of a
compare-and-swap ``where`` (0 means another writer won).
"""
# build statement
set_clause = ', '.join(f'{key} = ?' for key in data)
where_clause, params = self._where_clause(where)
statement = f'UPDATE {table} SET {set_clause} WHERE {where_clause}'
params = tuple(data.values()) + params
# execute statement
with self._handle(connection) as handle:
cursor = handle.execute(statement, params)
count = cursor.rowcount
return count
[docs]
def merge(
self: Database,
data: Row,
table: str,
*,
conflict: Optional[list[str]] = None,
connection: Optional[sqlite3.Connection] = None,
) -> int:
"""Insert a row, or upsert it on a unique conflict.
With ``conflict`` (the unique column(s) identifying an existing row), an
existing row is updated in place for **only** the columns in ``data`` --
every other column keeps its stored value. This is a real upsert: a
partial merge (e.g. just ``status``) preserves the row's other columns
rather than wiping them the way a whole-row ``INSERT OR REPLACE`` would.
A partial merge must still supply every ``NOT NULL`` column that lacks a
SQL default: the base ``INSERT``'s constraints are checked before the
conflict routes to the update, so omitting one fails (it is not
backfilled). Without ``conflict`` it falls back to ``INSERT OR REPLACE``
(a full-row insert).
Args:
data: Column values.
table: Table name.
conflict: Unique column(s) to upsert on. ``None`` does
a whole-row insert-or-replace.
connection: Transaction to run inside (from
:meth:`transaction`).
Returns:
Row ID of the inserted, updated, or (on a no-op) existing row.
"""
# build statement
columns = ', '.join(data.keys())
placeholders = ', '.join('?' for _ in data)
if conflict:
targets = ', '.join(conflict)
# update only the provided non-conflict columns so
# unprovided columns keep their stored values
updates = ', '.join(
f'{key} = excluded.{key}' for key in data if key not in conflict
)
if updates:
statement = (
f'INSERT INTO {table} ({columns}) VALUES ({placeholders})'
f' ON CONFLICT({targets}) DO UPDATE SET {updates}'
' RETURNING rowid'
)
else:
# nothing to update -- a pure "insert if new, else leave alone";
# INSERT OR IGNORE skips the conflicting row rather than failing;
# ON CONFLICT DO NOTHING would still evaluate the base INSERT's
# NOT NULL checks first and raise on a conflict-only row
statement = (
f'INSERT OR IGNORE INTO {table} ({columns})'
f' VALUES ({placeholders}) RETURNING rowid'
)
else:
statement = (
f'INSERT OR REPLACE INTO {table} ({columns}) VALUES ({placeholders})'
)
# execute statement
with self._handle(connection) as handle:
cursor = handle.execute(statement, tuple(data.values()))
if conflict:
# RETURNING yields the affected row on insert/update, but nothing
# on a no-op (the row already exists and was left alone) -- so fall
# back to resolving the existing row's id by the conflict key(s)
returned = cursor.fetchone()
if returned is not None:
row_id, *_ = returned
else:
where = ' AND '.join(f'{key} = ?' for key in conflict)
lookup = handle.execute(
f'SELECT rowid FROM {table} WHERE {where}',
tuple(data[key] for key in conflict),
)
found = lookup.fetchone()
row_id = found[0] if found else 0
else:
row_id = cursor.lastrowid
return row_id
[docs]
def delete(
self: Database,
table: str,
*,
where: dict[str, Any],
connection: Optional[sqlite3.Connection] = None,
) -> int:
"""Delete rows from a table.
Args:
table: Table name.
where: Column filters (AND-joined equality).
connection: Transaction to run inside (from
:meth:`transaction`).
Returns:
Number of rows deleted -- so callers can report a zero-match
delete honestly instead of implying one landed.
"""
# build statement
where_clause, params = self._where_clause(where)
statement = f'DELETE FROM {table} WHERE {where_clause}'
# execute statement
with self._handle(connection) as handle:
cursor = handle.execute(statement, params)
count = cursor.rowcount
return count
[docs]
def exists(
self: Database,
table: str,
*,
where: dict[str, Any],
connection: Optional[sqlite3.Connection] = None,
) -> bool:
"""Check whether a matching row exists.
Args:
table: Table name.
where: Column filters (AND-joined equality).
connection: Transaction to run inside (from
:meth:`transaction`).
Returns:
Whether a matching row exists.
"""
# build query
where_clause, params = self._where_clause(where)
query = f'SELECT 1 FROM {table} WHERE {where_clause} LIMIT 1'
# execute query
with self._handle(connection, read_only=True) as handle:
cursor = handle.execute(query, params)
result = cursor.fetchone() is not None
return result
[docs]
def count(
self: Database,
table: str,
*,
where: Optional[dict[str, Any]] = None,
connection: Optional[sqlite3.Connection] = None,
) -> int:
"""Count rows in a table.
Args:
table: Table name.
where: Column filters (AND-joined equality).
connection: Transaction to run inside (from
:meth:`transaction`).
Returns:
Row count.
"""
# build query
query = f'SELECT COUNT(*) FROM {table}'
params = ()
if where:
where_clause, where_params = self._where_clause(where)
query += ' WHERE ' + where_clause
params += where_params
# execute query
with self._handle(connection, read_only=True) as handle:
cursor = handle.execute(query, params)
result, *_ = cursor.fetchone()
return result
@staticmethod
def _where_clause(where: dict[str, Any]) -> tuple[str, tuple[Any, ...]]:
"""Build an AND-joined WHERE clause.
Renders a ``None`` filter value as ``col IS NULL``,
binding no parameter. SQL never matches ``col = NULL``
(the comparison is always unknown), so equality cannot
stand in for a null check.
Args:
where: Column filters (AND-joined equality).
Returns:
The clause text (without the ``WHERE`` keyword) and
the parameter tuple for the non-null comparisons,
positionally aligned with it.
"""
clauses = []
params = []
for key, value in where.items():
if value is None:
clauses.append(f'{key} IS NULL')
else:
clauses.append(f'{key} = ?')
params.append(value)
return ' AND '.join(clauses), tuple(params)