database.py

import sqlite3
from contextlib import contextmanager
from typing import Optional


class Database:
    """
    Handles all raw database connection and query execution for the bank system.
    No business logic lives here - just connection management and query execution.
    """

    def __init__(self, db_path: str = "bank.db"):
        self.db_path = db_path
        self.con = sqlite3.connect(self.db_path, check_same_thread=False)
        self.con.row_factory = sqlite3.Row   # lets us access columns by name, not just index
        self.cur = self.con.cursor()
        self._create_tables_if_missing()

    def _create_tables_if_missing(self):
        self.cur.execute("""
            CREATE TABLE IF NOT EXISTS bank (
                account_no      INTEGER PRIMARY KEY,
                full_name       TEXT NOT NULL,
                date_of_birth   TEXT NOT NULL,
                gender          TEXT NOT NULL,
                phone_no        TEXT NOT NULL,
                current_balance REAL NOT NULL DEFAULT 0,
                opening_date    TEXT NOT NULL,
                card_issued     INTEGER NOT NULL DEFAULT 0,
                pin_hash        TEXT NOT NULL,
                email           TEXT,
                address         TEXT,
                branch          TEXT
            )
        """)
        self.cur.execute("""
            CREATE TABLE IF NOT EXISTS staff (
                staff_id      TEXT PRIMARY KEY,
                staff_name    TEXT NOT NULL,
                password_hash TEXT NOT NULL
            )
        """)
        self.con.commit()

    def execute(self, query: str, params: tuple = (), fetch: Optional[str] = None):
        """
        Runs a single query safely.
        fetch: None (default, for INSERT/UPDATE/DELETE), "one", or "all"
        Returns: None, a single row, or a list of rows - depending on fetch.
        Raises sqlite3.Error on failure (caller decides how to handle it).
        """
        try:
            self.cur.execute(query, params)
            if fetch == "one":
                return self.cur.fetchone()
            elif fetch == "all":
                return self.cur.fetchall()
            return None
        except sqlite3.Error as e:
            raise RuntimeError(f"Database error: {e}") from e

    def executemany(self, query: str, param_list: list):
        """For bulk inserts, e.g. seeding dummy data."""
        try:
            self.cur.executemany(query, param_list)
        except sqlite3.Error as e:
            raise RuntimeError(f"Database error: {e}") from e

    def commit(self):
        self.con.commit()

    def rollback(self):
        self.con.rollback()

    def close(self):
        self.con.close()

    @contextmanager
    def transaction(self):
        """
        Use for multi-step operations that must succeed or fail together,
        e.g. transfer() which debits one account and credits another.
        Usage:
            with db.transaction():
                db.execute(...)
                db.execute(...)
        """
        try:
            yield
            self.commit()
        except Exception:
            self.rollback()
            raise

    def __enter__(self):
        return self

    def __exit__(self, exc_type, exc_val, exc_tb):
        self.close()