232 lines
9.5 KiB
Python
232 lines
9.5 KiB
Python
# Licensed under the Apache License: http://www.apache.org/licenses/LICENSE-2.0
|
|
# For details: https://github.com/nedbat/coveragepy/blob/master/NOTICE.txt
|
|
|
|
"""SQLite abstraction for coverage.py"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import contextlib
|
|
import re
|
|
import sqlite3
|
|
|
|
from typing import cast, Any
|
|
from collections.abc import Iterable, Iterator
|
|
|
|
from coverage.debug import auto_repr, clipped_repr, exc_one_line
|
|
from coverage.exceptions import DataError
|
|
from coverage.types import TDebugCtl
|
|
|
|
|
|
class SqliteDb:
|
|
"""A simple abstraction over a SQLite database.
|
|
|
|
Use as a context manager, then you can use it like a
|
|
:class:`python:sqlite3.Connection` object::
|
|
|
|
with SqliteDb(filename, debug_control) as db:
|
|
with db.execute("select a, b from some_table") as cur:
|
|
for a, b in cur:
|
|
etc(a, b)
|
|
|
|
"""
|
|
def __init__(self, filename: str, debug: TDebugCtl) -> None:
|
|
self.debug = debug
|
|
self.filename = filename
|
|
self.nest = 0
|
|
self.con: sqlite3.Connection | None = None
|
|
|
|
__repr__ = auto_repr
|
|
|
|
def _connect(self) -> None:
|
|
"""Connect to the db and do universal initialization."""
|
|
if self.con is not None:
|
|
return
|
|
|
|
# It can happen that Python switches threads while the tracer writes
|
|
# data. The second thread will also try to write to the data,
|
|
# effectively causing a nested context. However, given the idempotent
|
|
# nature of the tracer operations, sharing a connection among threads
|
|
# is not a problem.
|
|
if self.debug.should("sql"):
|
|
self.debug.write(f"Connecting to {self.filename!r}")
|
|
try:
|
|
self.con = sqlite3.connect(self.filename, check_same_thread=False)
|
|
except sqlite3.Error as exc:
|
|
raise DataError(f"Couldn't use data file {self.filename!r}: {exc}") from exc
|
|
|
|
if self.debug.should("sql"):
|
|
self.debug.write(f"Connected to {self.filename!r} as {self.con!r}")
|
|
|
|
self.con.create_function("REGEXP", 2, lambda txt, pat: re.search(txt, pat) is not None)
|
|
|
|
# Turning off journal_mode can speed up writing. It can't always be
|
|
# disabled, so we have to be prepared for *-journal files elsewhere.
|
|
# In Python 3.12+, we can change the config to allow journal_mode=off.
|
|
if hasattr(sqlite3, "SQLITE_DBCONFIG_DEFENSIVE"):
|
|
# Turn off defensive mode, so that journal_mode=off can succeed.
|
|
self.con.setconfig( # type: ignore[attr-defined, unused-ignore]
|
|
sqlite3.SQLITE_DBCONFIG_DEFENSIVE, False,
|
|
)
|
|
|
|
# This pragma makes writing faster. It disables rollbacks, but we never need them.
|
|
self.execute_void("pragma journal_mode=off")
|
|
|
|
# This pragma makes writing faster. It can fail in unusual situations
|
|
# (https://github.com/nedbat/coveragepy/issues/1646), so use fail_ok=True
|
|
# to keep things going.
|
|
self.execute_void("pragma synchronous=off", fail_ok=True)
|
|
|
|
def close(self) -> None:
|
|
"""If needed, close the connection."""
|
|
if self.con is not None and self.filename != ":memory:":
|
|
if self.debug.should("sql"):
|
|
self.debug.write(f"Closing {self.con!r} on {self.filename!r}")
|
|
self.con.close()
|
|
self.con = None
|
|
|
|
def __enter__(self) -> SqliteDb:
|
|
if self.nest == 0:
|
|
self._connect()
|
|
assert self.con is not None
|
|
self.con.__enter__()
|
|
self.nest += 1
|
|
return self
|
|
|
|
def __exit__(self, exc_type, exc_value, traceback) -> None: # type: ignore[no-untyped-def]
|
|
self.nest -= 1
|
|
if self.nest == 0:
|
|
try:
|
|
assert self.con is not None
|
|
self.con.__exit__(exc_type, exc_value, traceback)
|
|
self.close()
|
|
except Exception as exc:
|
|
if self.debug.should("sql"):
|
|
self.debug.write(f"EXCEPTION from __exit__: {exc_one_line(exc)}")
|
|
raise DataError(f"Couldn't end data file {self.filename!r}: {exc}") from exc
|
|
|
|
def _execute(self, sql: str, parameters: Iterable[Any]) -> sqlite3.Cursor:
|
|
"""Same as :meth:`python:sqlite3.Connection.execute`."""
|
|
if self.debug.should("sql"):
|
|
tail = f" with {parameters!r}" if parameters else ""
|
|
self.debug.write(f"Executing {sql!r}{tail}")
|
|
try:
|
|
assert self.con is not None
|
|
try:
|
|
return self.con.execute(sql, parameters) # type: ignore[arg-type]
|
|
except Exception:
|
|
# In some cases, an error might happen that isn't really an
|
|
# error. Try again immediately.
|
|
# https://github.com/nedbat/coveragepy/issues/1010
|
|
return self.con.execute(sql, parameters) # type: ignore[arg-type]
|
|
except sqlite3.Error as exc:
|
|
msg = str(exc)
|
|
if self.filename != ":memory:":
|
|
try:
|
|
# `execute` is the first thing we do with the database, so try
|
|
# hard to provide useful hints if something goes wrong now.
|
|
with open(self.filename, "rb") as bad_file:
|
|
cov4_sig = b"!coverage.py: This is a private format"
|
|
if bad_file.read(len(cov4_sig)) == cov4_sig:
|
|
msg = (
|
|
"Looks like a coverage 4.x data file. " +
|
|
"Are you mixing versions of coverage?"
|
|
)
|
|
except Exception:
|
|
pass
|
|
if self.debug.should("sql"):
|
|
self.debug.write(f"EXCEPTION from execute: {exc_one_line(exc)}")
|
|
raise DataError(f"Couldn't use data file {self.filename!r}: {msg}") from exc
|
|
|
|
@contextlib.contextmanager
|
|
def execute(
|
|
self,
|
|
sql: str,
|
|
parameters: Iterable[Any] = (),
|
|
) -> Iterator[sqlite3.Cursor]:
|
|
"""Context managed :meth:`python:sqlite3.Connection.execute`.
|
|
|
|
Use with a ``with`` statement to auto-close the returned cursor.
|
|
"""
|
|
cur = self._execute(sql, parameters)
|
|
try:
|
|
yield cur
|
|
finally:
|
|
cur.close()
|
|
|
|
def execute_void(self, sql: str, parameters: Iterable[Any] = (), fail_ok: bool = False) -> None:
|
|
"""Same as :meth:`python:sqlite3.Connection.execute` when you don't need the cursor.
|
|
|
|
If `fail_ok` is True, then SQLite errors are ignored.
|
|
"""
|
|
try:
|
|
# PyPy needs the .close() calls here, or sqlite gets twisted up:
|
|
# https://bitbucket.org/pypy/pypy/issues/2872/default-isolation-mode-is-different-on
|
|
self._execute(sql, parameters).close()
|
|
except DataError:
|
|
if not fail_ok:
|
|
raise
|
|
|
|
def execute_for_rowid(self, sql: str, parameters: Iterable[Any] = ()) -> int:
|
|
"""Like execute, but returns the lastrowid."""
|
|
with self.execute(sql, parameters) as cur:
|
|
assert cur.lastrowid is not None
|
|
rowid: int = cur.lastrowid
|
|
if self.debug.should("sqldata"):
|
|
self.debug.write(f"Row id result: {rowid!r}")
|
|
return rowid
|
|
|
|
def execute_one(self, sql: str, parameters: Iterable[Any] = ()) -> tuple[Any, ...] | None:
|
|
"""Execute a statement and return the one row that results.
|
|
|
|
This is like execute(sql, parameters).fetchone(), except it is
|
|
correct in reading the entire result set. This will raise an
|
|
exception if more than one row results.
|
|
|
|
Returns a row, or None if there were no rows.
|
|
"""
|
|
with self.execute(sql, parameters) as cur:
|
|
rows = list(cur)
|
|
if len(rows) == 0:
|
|
return None
|
|
elif len(rows) == 1:
|
|
return cast(tuple[Any, ...], rows[0])
|
|
else:
|
|
raise AssertionError(f"SQL {sql!r} shouldn't return {len(rows)} rows")
|
|
|
|
def _executemany(self, sql: str, data: list[Any]) -> sqlite3.Cursor:
|
|
"""Same as :meth:`python:sqlite3.Connection.executemany`."""
|
|
if self.debug.should("sql"):
|
|
final = ":" if self.debug.should("sqldata") else ""
|
|
self.debug.write(f"Executing many {sql!r} with {len(data)} rows{final}")
|
|
if self.debug.should("sqldata"):
|
|
for i, row in enumerate(data):
|
|
self.debug.write(f"{i:4d}: {row!r}")
|
|
assert self.con is not None
|
|
try:
|
|
return self.con.executemany(sql, data)
|
|
except Exception:
|
|
# In some cases, an error might happen that isn't really an
|
|
# error. Try again immediately.
|
|
# https://github.com/nedbat/coveragepy/issues/1010
|
|
return self.con.executemany(sql, data)
|
|
|
|
def executemany_void(self, sql: str, data: Iterable[Any]) -> None:
|
|
"""Same as :meth:`python:sqlite3.Connection.executemany` when you don't need the cursor."""
|
|
data = list(data)
|
|
if data:
|
|
self._executemany(sql, data).close()
|
|
|
|
def executescript(self, script: str) -> None:
|
|
"""Same as :meth:`python:sqlite3.Connection.executescript`."""
|
|
if self.debug.should("sql"):
|
|
self.debug.write("Executing script with {} chars: {}".format(
|
|
len(script), clipped_repr(script, 100),
|
|
))
|
|
assert self.con is not None
|
|
self.con.executescript(script).close()
|
|
|
|
def dump(self) -> str:
|
|
"""Return a multi-line string, the SQL dump of the database."""
|
|
assert self.con is not None
|
|
return "\n".join(self.con.iterdump())
|