From 8b7c598407a5fea9a3d78473e72df87606998cd4 Mon Sep 17 00:00:00 2001 From: Frederick Muriuki Muriithi Date: Fri, 4 Aug 2023 10:10:28 +0300 Subject: Copy over files from GN3 repository. --- gn_auth/auth/db.py | 78 ++++++++++++++++++++++++++++++++++++++++++++++++++++++ 1 file changed, 78 insertions(+) create mode 100644 gn_auth/auth/db.py (limited to 'gn_auth/auth/db.py') diff --git a/gn_auth/auth/db.py b/gn_auth/auth/db.py new file mode 100644 index 0000000..2ba6619 --- /dev/null +++ b/gn_auth/auth/db.py @@ -0,0 +1,78 @@ +"""Handle connection to auth database.""" +import sqlite3 +import logging +import contextlib +from typing import Any, Callable, Iterator, Protocol + +import traceback + +class DbConnection(Protocol): + """Type annotation for a generic database connection object.""" + def cursor(self) -> Any: + """A cursor object""" + ... + + def commit(self) -> Any: + """Commit the transaction.""" + ... + + def rollback(self) -> Any: + """Rollback the transaction.""" + ... + +class DbCursor(Protocol): + """Type annotation for a generic database cursor object.""" + def execute(self, *args, **kwargs) -> Any: + """Execute a single query""" + ... + + def executemany(self, *args, **kwargs) -> Any: + """ + Execute parameterized SQL statement sql against all parameter sequences + or mappings found in the sequence parameters. + """ + ... + + def fetchone(self, *args, **kwargs): + """Fetch single result if present, or `None`.""" + ... + + def fetchmany(self, *args, **kwargs): + """Fetch many results if present or `None`.""" + ... + + def fetchall(self, *args, **kwargs): + """Fetch all results if present or `None`.""" + ... + +@contextlib.contextmanager +def connection(db_path: str, row_factory: Callable = sqlite3.Row) -> Iterator[DbConnection]: + """Create the connection to the auth database.""" + logging.debug("SQLite3 DB Path: '%s'.", db_path) + conn = sqlite3.connect(db_path) + conn.row_factory = row_factory + conn.set_trace_callback(logging.debug) + conn.execute("PRAGMA foreign_keys = ON") + try: + yield conn + except sqlite3.Error as exc: + conn.rollback() + logging.debug(traceback.format_exc()) + raise exc + finally: + conn.commit() + conn.close() + +@contextlib.contextmanager +def cursor(conn: DbConnection) -> Iterator[DbCursor]: + """Get a cursor from the given connection to the auth database.""" + cur = conn.cursor() + try: + yield cur + except sqlite3.Error as exc: + conn.rollback() + logging.debug(traceback.format_exc()) + raise exc + finally: + conn.commit() + cur.close() -- cgit v1.2.3