diff --git a/sqlite_utils/db.py b/sqlite_utils/db.py index f5b1b12..fd7fd4b 100644 --- a/sqlite_utils/db.py +++ b/sqlite_utils/db.py @@ -1,5 +1,6 @@ from .utils import sqlite3, OperationalError, suggest_column_types, column_affinity from collections import namedtuple, OrderedDict +import contextlib import datetime import decimal import hashlib @@ -109,6 +110,7 @@ class Database: memory=False, recreate=False, recursive_triggers=True, + tracer=None, ): assert (filename_or_conn is not None and not memory) or ( filename_or_conn is None and memory @@ -124,6 +126,16 @@ class Database: self.conn = filename_or_conn if recursive_triggers: self.conn.execute("PRAGMA recursive_triggers=on;") + self._tracer = tracer + + @contextlib.contextmanager + def tracer(self, tracer=None): + prev_tracer = self.tracer + self._tracer = tracer or print + try: + yield self + finally: + self._tracer = prev_tracer def __getitem__(self, table_name): return self.table(table_name) @@ -131,6 +143,19 @@ class Database: def __repr__(self): return "".format(self.conn) + def execute(self, sql, parameters=None): + if self._tracer: + self._tracer(sql, parameters) + if parameters is not None: + return self.conn.execute(sql, parameters) + else: + return self.conn.execute(sql) + + def executescript(self, sql): + if self._tracer: + self._tracer(sql) + return self.conn.executescript(sql) + def table(self, table_name, **kwargs): klass = View if table_name in self.view_names() else Table return klass(self, table_name, **kwargs) @@ -139,7 +164,7 @@ class Database: # 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( + return self.execute( # Use SQLite itself to correctly escape this string: "SELECT quote(:value)", {"value": value}, @@ -152,12 +177,12 @@ class Database: 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()] + return [r[0] for r in self.execute(sql).fetchall()] def view_names(self): return [ r[0] - for r in self.conn.execute( + for r in self.execute( "select name from sqlite_master where type = 'view'" ).fetchall() ] @@ -174,25 +199,25 @@ class Database: def triggers(self): return [ Trigger(*r) - for r in self.conn.execute( + for r in self.execute( "select name, tbl_name, sql from sqlite_master where type = 'trigger'" ).fetchall() ] @property def journal_mode(self): - return self.conn.execute("PRAGMA journal_mode;").fetchone()[0] + return self.execute("PRAGMA journal_mode;").fetchone()[0] def enable_wal(self): if self.journal_mode != "wal": - self.conn.execute("PRAGMA journal_mode=wal;") + self.execute("PRAGMA journal_mode=wal;") def disable_wal(self): if self.journal_mode != "delete": - self.conn.execute("PRAGMA journal_mode=delete;") + self.execute("PRAGMA journal_mode=delete;") def execute_returning_dicts(self, sql, params=None): - cursor = self.conn.execute(sql, params or tuple()) + cursor = self.execute(sql, params or tuple()) keys = [d[0] for d in cursor.description] return [dict(zip(keys, row)) for row in cursor.fetchall()] @@ -337,7 +362,7 @@ class Database: """.format( table=name, columns_sql=columns_sql, extra_pk=extra_pk ) - self.conn.execute(sql) + self.execute(sql) return self.table( name, pk=pk, @@ -363,7 +388,7 @@ class Database: if create_sql == self[name].schema: return self self[name].drop() - self.conn.execute(create_sql) + self.execute(create_sql) return self def m2m_table_candidates(self, table, other_table): @@ -451,7 +476,7 @@ class Database: table.create_index([fk.column]) def vacuum(self): - self.conn.execute("VACUUM;") + self.execute("VACUUM;") class Queryable: @@ -464,7 +489,7 @@ class Queryable: @property def count(self): - return self.db.conn.execute( + return self.db.execute( "select count(*) from [{}]".format(self.name) ).fetchone()[0] @@ -480,7 +505,7 @@ class Queryable: sql += " where " + where if order_by is not None: sql += " order by " + order_by - cursor = self.db.conn.execute(sql, where_args or []) + cursor = self.db.execute(sql, where_args or []) columns = [c[0] for c in cursor.description] for row in cursor: yield dict(zip(columns, row)) @@ -489,9 +514,7 @@ class Queryable: def columns(self): if not self.exists(): return [] - rows = self.db.conn.execute( - "PRAGMA table_info([{}])".format(self.name) - ).fetchall() + rows = self.db.execute("PRAGMA table_info([{}])".format(self.name)).fetchall() return [Column(*row) for row in rows] @property @@ -501,7 +524,7 @@ class Queryable: @property def schema(self): - return self.db.conn.execute( + return self.db.execute( "select sql from sqlite_master where name = ?", (self.name,) ).fetchone()[0] @@ -587,7 +610,7 @@ class Table(Queryable): @property def foreign_keys(self): fks = [] - for row in self.db.conn.execute( + for row in self.db.execute( "PRAGMA foreign_key_list([{}])".format(self.name) ).fetchall(): if row is not None: @@ -615,7 +638,7 @@ class Table(Queryable): ) column_sql = "PRAGMA index_info({})".format(index_name_quoted) columns = [] - for seqno, cid, name in self.db.conn.execute(column_sql).fetchall(): + for seqno, cid, name in self.db.execute(column_sql).fetchall(): columns.append(name) row["columns"] = columns # These columns may be missing on older SQLite versions: @@ -629,7 +652,7 @@ class Table(Queryable): def triggers(self): return [ Trigger(*r) - for r in self.db.conn.execute( + for r in self.db.execute( "select name, tbl_name, sql from sqlite_master where type = 'trigger'" " and tbl_name = ?", (self.name,), @@ -683,7 +706,7 @@ class Table(Queryable): if_not_exists="IF NOT EXISTS " if if_not_exists else "", ) ) - self.db.conn.execute(sql) + self.db.execute(sql) return self def add_column( @@ -720,13 +743,13 @@ class Table(Queryable): 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) + self.db.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)) + self.db.execute("DROP TABLE [{}]".format(self.name)) def guess_foreign_table(self, column): column = column.lower() @@ -811,7 +834,7 @@ class Table(Queryable): tokenize="\n tokenize='{}',".format(tokenize) if tokenize else "", ) ) - self.db.conn.executescript(sql) + self.db.executescript(sql) self.populate_fts(columns) if create_triggers: @@ -840,7 +863,7 @@ class Table(Queryable): new_cols=new_cols, ) ) - self.db.conn.executescript(triggers) + self.db.executescript(triggers) return self def populate_fts(self, columns): @@ -850,7 +873,7 @@ class Table(Queryable): """.format( table=self.name, columns=", ".join("[{}]".format(c) for c in columns) ) - self.db.conn.executescript(sql) + self.db.executescript(sql) return self def disable_fts(self): @@ -866,11 +889,11 @@ class Table(Queryable): fts_table ) trigger_names = [] - for row in self.db.conn.execute(sql).fetchall(): + for row in self.db.execute(sql).fetchall(): trigger_names.append(row[0]) with self.db.conn: for trigger_name in trigger_names: - self.db.conn.execute("DROP TRIGGER IF EXISTS [{}]".format(trigger_name)) + self.db.execute("DROP TRIGGER IF EXISTS [{}]".format(trigger_name)) def detect_fts(self): "Detect if table has a corresponding FTS virtual table and return it" @@ -887,7 +910,7 @@ class Table(Queryable): """.format( table=self.name ) - rows = self.db.conn.execute(sql).fetchall() + rows = self.db.execute(sql).fetchall() if len(rows) == 0: return None else: @@ -896,7 +919,7 @@ class Table(Queryable): def optimize(self): fts_table = self.detect_fts() if fts_table is not None: - self.db.conn.execute( + self.db.execute( """ INSERT INTO [{table}] ([{table}]) VALUES ("optimize"); """.format( @@ -925,7 +948,7 @@ class Table(Queryable): """.format( table=self.name ) - return self.db.conn.execute(sql, (q,)).fetchall() + return self.db.execute(sql, (q,)).fetchall() def value_or_default(self, key, value): return self._defaults[key] if value is DEFAULT else value @@ -939,7 +962,7 @@ class Table(Queryable): table=self.name, wheres=" and ".join(wheres) ) with self.db.conn: - self.db.conn.execute(sql, pk_values) + self.db.execute(sql, pk_values) def delete_where(self, where=None, where_args=None): if not self.exists(): @@ -947,7 +970,7 @@ class Table(Queryable): sql = "delete from [{}]".format(self.name) if where is not None: sql += " where " + where - self.db.conn.execute(sql, where_args or []) + self.db.execute(sql, where_args or []) def update(self, pk_values, updates=None, alter=False, conversions=None): updates = updates or {} @@ -972,12 +995,12 @@ class Table(Queryable): ) with self.db.conn: try: - rowcount = self.db.conn.execute(sql, args).rowcount + rowcount = self.db.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 + rowcount = self.db.execute(sql, args).rowcount else: raise @@ -1084,7 +1107,7 @@ class Table(Queryable): self.last_rowid = None self.last_pk = None if truncate and self.exists(): - self.db.conn.execute("DELETE FROM [{}];".format(self.name)) + self.db.execute("DELETE FROM [{}];".format(self.name)) for chunk in chunks(itertools.chain([first_record], records), batch_size): chunk = list(chunk) num_records_processed += len(chunk) @@ -1205,12 +1228,12 @@ class Table(Queryable): with self.db.conn: for query, params in queries_and_params: try: - result = self.db.conn.execute(query, params) + result = self.db.execute(query, params) 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(query, params) + result = self.db.execute(query, params) else: raise if num_records_processed == 1 and not upsert: @@ -1383,7 +1406,7 @@ class View(Queryable): ) def drop(self): - self.db.conn.execute("DROP VIEW [{}]".format(self.name)) + self.db.execute("DROP VIEW [{}]".format(self.name)) def chunks(sequence, size):