Changeset 527c097


Ignore:
Timestamp:
04/05/26 20:02:04 (3 months ago)
Author:
teo <teodorbogoeski@…>
Branches:
main
Parents:
95cf703
Message:

Add connection pooling and transaction management

Files:
5 edited

Legend:

Unmodified
Added
Removed

  • db.py

    r95cf703 r527c097  
    11import datetime
    22import psycopg2
     3from psycopg2 import pool
     4from contextlib import contextmanager
    35
    46DB_CONFIG = {
    … …  
    1113}
    1214
     15# -- Connection Pool ----------------------------------------------------------
     16# Reuses database connections instead of opening a new one per query.
     17# SimpleConnectionPool maintains between minconn and maxconn connections.
    1318
     19_pool = None
     20
     21
     22def _get_pool():
     23    global _pool
     24    if _pool is None:
     25        _pool = pool.SimpleConnectionPool(minconn=2, maxconn=10, **DB_CONFIG)
     26    return _pool
     27
     28
     29@contextmanager
    1430def get_connection():
    15     return psycopg2.connect(**DB_CONFIG)
     31    """Get a read-only connection from the pool.
     32    Autocommit is enabled so SELECTs don't need explicit commit.
     33    Connection is returned to the pool on exit."""
     34    p = _get_pool()
     35    conn = p.getconn()
     36    try:
     37        conn.autocommit = True
     38        yield conn
     39    finally:
     40        conn.autocommit = False
     41        p.putconn(conn)
     42
     43
     44@contextmanager
     45def get_transaction():
     46    """Get a connection with explicit transaction management.
     47    Commits on successful exit, rolls back on exception.
     48    Connection is returned to the pool in both cases."""
     49    p = _get_pool()
     50    conn = p.getconn()
     51    try:
     52        conn.autocommit = False
     53        yield conn
     54        conn.commit()
     55    except Exception:
     56        conn.rollback()
     57        raise
     58    finally:
     59        p.putconn(conn)
     60
     61
     62def close_pool():
     63    """Close all connections in the pool."""
     64    global _pool
     65    if _pool is not None:
     66        _pool.closeall()
     67        _pool = None
     68
     69
     70# -- UI Helpers ---------------------------------------------------------------
    1671
    1772

  • main.py

    r95cf703 r527c097  
    44import bcrypt
    55
    6 from db import get_connection, pick_from_list
     6from db import get_connection, close_pool, pick_from_list
    77from uc_browse import browse_resources
    88from uc_reserve import make_reservation
    … …  
    9595    except KeyboardInterrupt:
    9696        print("\n\n  Goodbye.")
     97    finally:
     98        close_pool()
    9799
    98100

  • uc_approve.py

    r95cf703 r527c097  
    1 from db import get_connection, pick_from_list, print_table
     1from db import get_connection, get_transaction, pick_from_list, print_table
    22
    33
    … …  
    120120
    121121    new_status = "approved" if action == 0 else "rejected"
    122     conn = get_connection()
    123122    try:
    124         with conn.cursor() as cur:
    125             cur.execute(
    126                 """
    127                 UPDATE reservations
    128                 SET status = %s, approved_by = %s
    129                 WHERE reservation_id = %s AND status = 'pending'
    130                 RETURNING reservation_id, status
    131                 """,
    132                 (new_status, admin_user_id, res["reservation_id"]),
    133             )
    134             result = cur.fetchone()
    135         conn.commit()
     123        with get_transaction() as conn:
     124            with conn.cursor() as cur:
     125                # For approvals, re-check conflicts within the transaction
     126                # to prevent approving overlapping reservations
     127                if new_status == "approved":
     128                    cur.execute(
     129                        """
     130                        SELECT COUNT(*) FROM reservations
     131                        WHERE resource_id = %s
     132                          AND reservation_id != %s
     133                          AND status = 'approved'
     134                          AND start_time < %s AND end_time > %s
     135                        """,
     136                        (res["resource_id"], res["reservation_id"],
     137                         res["end_time"], res["start_time"]),
     138                    )
     139                    if cur.fetchone()[0] > 0:
     140                        raise RuntimeError("Conflict: another reservation was approved for this time slot")
     141
     142                cur.execute(
     143                    """
     144                    UPDATE reservations
     145                    SET status = %s, approved_by = %s
     146                    WHERE reservation_id = %s AND status = 'pending'
     147                    RETURNING reservation_id, status
     148                    """,
     149                    (new_status, admin_user_id, res["reservation_id"]),
     150                )
     151                result = cur.fetchone()
     152        # Transaction committed automatically
    136153        if result:
    137154            print(f"\n  Reservation #{result[0]} has been {result[1]}.")
    … …  
    139156            print("\n  Reservation was already processed by another administrator.")
    140157    except Exception as e:
    141         conn.rollback()
     158        # Transaction rolled back automatically
    142159        print(f"\n  Error: {e}")
    143     finally:
    144         conn.close()

  • uc_reserve.py

    r95cf703 r527c097  
    11import datetime
    22
    3 from db import get_connection, pick_from_list, print_table, input_date, input_time
     3from db import get_connection, get_transaction, pick_from_list, print_table, input_date, input_time
    44
    55
    … …  
    9696            return
    9797
    98         # Step 7: INSERT
    99         conn = get_connection()
     98        # Step 7: INSERT (transaction: re-check conflicts + insert atomically)
    10099        try:
    101             with conn.cursor() as cur:
    102                 cur.execute(
    103                     """
    104                     INSERT INTO reservations
    105                         (start_time, end_time, status, purpose, created_at, user_id, resource_id)
    106                     VALUES (%s, %s, 'pending', %s, CURRENT_TIMESTAMP, %s, %s)
    107                     RETURNING reservation_id, status, created_at
    108                     """,
    109                     (start_dt, end_dt, purpose, user_id, rid),
    110                 )
    111                 result = cur.fetchone()
    112             conn.commit()
     100            with get_transaction() as conn:
     101                with conn.cursor() as cur:
     102                    # Re-check conflicts inside the transaction to prevent
     103                    # race conditions (another user inserting between our
     104                    # check and insert)
     105                    cur.execute(
     106                        """
     107                        SELECT COUNT(*) FROM reservations
     108                        WHERE resource_id = %s
     109                          AND status IN ('approved', 'pending')
     110                          AND start_time < %s AND end_time > %s
     111                        """,
     112                        (rid, end_dt, start_dt),
     113                    )
     114                    if cur.fetchone()[0] > 0:
     115                        raise RuntimeError("Conflict: another reservation was just created for this time slot")
     116
     117                    cur.execute(
     118                        """
     119                        INSERT INTO reservations
     120                            (start_time, end_time, status, purpose, created_at, user_id, resource_id)
     121                        VALUES (%s, %s, 'pending', %s, CURRENT_TIMESTAMP, %s, %s)
     122                        RETURNING reservation_id, status, created_at
     123                        """,
     124                        (start_dt, end_dt, purpose, user_id, rid),
     125                    )
     126                    result = cur.fetchone()
     127            # Transaction committed automatically
    113128            print(f"\n  Reservation created successfully!")
    114129            print(f"  Reservation ID: {result[0]}")
    … …  
    117132            return
    118133        except Exception as e:
    119             conn.rollback()
     134            # Transaction rolled back automatically
    120135            print(f"\n  Error creating reservation: {e}")
    121136            return
    122         finally:
    123             conn.close()
    124137
    125138

  • uc_users.py

    r95cf703 r527c097  
    33import bcrypt
    44
    5 from db import get_connection, pick_from_list
     5from db import get_connection, get_transaction, pick_from_list
    66
    77
    … …  
    4040        return
    4141
    42     # Step 3: Check email uniqueness
    43     with get_connection() as conn:
    44         with conn.cursor() as cur:
    45             cur.execute("SELECT COUNT(*) FROM users WHERE email = %s", (email,))
    46             exists = cur.fetchone()[0]
    47     if exists > 0:
    48         print(f"\n  ERROR: A user with email '{email}' already exists.")
    49         return
    50 
    51     # Step 4: Enter password
     42    # Step 3: Enter password
    5243    password = getpass.getpass("  Initial password: ").strip()
    5344    if not password:
    … …  
    5950        return
    6051
    61     # Step 5: Hash and insert
     52    # Step 4: Hash password
    6253    hashed = bcrypt.hashpw(password.encode(), bcrypt.gensalt()).decode()
    6354
    64     conn = get_connection()
     55    # Step 5: Transaction — check email uniqueness + insert atomically
     56    # This prevents a race condition where two users register with the
     57    # same email simultaneously (one check passes, both insert).
    6558    try:
    66         with conn.cursor() as cur:
    67             cur.execute(
    68                 """
    69                 INSERT INTO users (first_name, last_name, email, password, type_id)
    70                 VALUES (%s, %s, %s, %s, %s)
    71                 RETURNING user_id, first_name, last_name, email
    72                 """,
    73                 (first_name, last_name, email, hashed, type_id),
    74             )
    75             result = cur.fetchone()
    76         conn.commit()
     59        with get_transaction() as conn:
     60            with conn.cursor() as cur:
     61                cur.execute("SELECT COUNT(*) FROM users WHERE email = %s", (email,))
     62                if cur.fetchone()[0] > 0:
     63                    raise ValueError(f"A user with email '{email}' already exists.")
    7764
    78         # Step 6: Confirmation
     65                cur.execute(
     66                    """
     67                    INSERT INTO users (first_name, last_name, email, password, type_id)
     68                    VALUES (%s, %s, %s, %s, %s)
     69                    RETURNING user_id, first_name, last_name, email
     70                    """,
     71                    (first_name, last_name, email, hashed, type_id),
     72                )
     73                result = cur.fetchone()
     74
     75        # Transaction committed automatically — Step 6: Confirmation
    7976        print(f"\n  User created successfully!")
    8077        print(f"  User ID:  {result[0]}")
    … …  
    8279        print(f"  Email:    {result[3]}")
    8380        print(f"  Role:     {role_name}")
     81    except ValueError as e:
     82        print(f"\n  ERROR: {e}")
    8483    except Exception as e:
    85         conn.rollback()
    8684        print(f"\n  Error creating user: {e}")
    87     finally:
    88         conn.close()
Note: See TracChangeset for help on using the changeset viewer.