database.py

import os
import time
import psycopg
from contextlib import contextmanager
import threading
from aws_database.ssm_tunnel import close_ssm_tunnel, ensure_ssm_tunnel
import logging_setup

logger = logging_setup.setup_logging()

USE_AWS = bool(os.environ.get("AWS_RDS_HOST"))

if USE_AWS:
    from aws_database.generate_token import generate_token

    logger.info("AWS_RDS_HOST detected → using SSM tunnel + IAM authentication")
    host = "localhost"  # Use localhost via SSM tunnel
    port = int(os.environ["LOCAL_PORT"])
else:
    logger.info("AWS_RDS_HOST not set → using direct database connection")
    host = os.environ["POSTGRES_HOST"]
    port = int(os.environ["POSTGRES_PORT"])


# Global excepthook: any error not handled in the THREAD will terminate the app
def _thread_crash_handler(args):
    logger.error(
        f"Unhandled exception in thread '{args.thread.name}': {args.exc_value}"
    )

    shutdown_database_access()

    # Kill the process immediately, in that stage the release lock step won't work anyway
    os._exit(1)


threading.excepthook = _thread_crash_handler


@contextmanager
def get_db_connection():
    if USE_AWS:
        ensure_ssm_tunnel()
        password = generate_token()
    else:
        password = os.environ["POSTGRES_PASSWORD"]

    conn = psycopg.connect(
        dbname=os.environ["POSTGRES_DB"],
        user=os.environ["POSTGRES_USER"],
        password=password,
        host=host,
        port=port,
    )

    try:
        yield conn
    finally:
        conn.close()


def shutdown_database_access():
    if not USE_AWS:
        return

    try:
        close_ssm_tunnel()
    except Exception as e:
        logger.error(f"Failed to close AWS SSM tunnel cleanly: {e}")


_heartbeat_thread = None
_heartbeat_stop_event = None


# Update the timestamp every 2 minutes
def _heartbeat_worker(device_uuid, stop_event, interval=120):
    logger.info(f"Heartbeat thread started for {device_uuid}")
    while not stop_event.is_set():
        with get_db_connection() as conn:
            with conn.cursor() as cursor:
                cursor.execute(
                    "UPDATE devices SET timestamp = NOW() WHERE device_uuid = %s",
                    (device_uuid,),
                )
                conn.commit()

        logger.debug(f"Heartbeat updated for {device_uuid}")

        # Waits for the interval, but exits sooner if stop_event is set
        stop_event.wait(interval)

    logger.info(f"Heartbeat thread stopped for {device_uuid}")


def acquire_lock(device_uuid):
    global _heartbeat_thread, _heartbeat_stop_event

    logger.info(f"Attempting to acquire lock for device {device_uuid}")
    with get_db_connection() as conn:
        with conn.cursor() as cursor:
            # Take the lock if the device is free or its lock is older than
            # 3 minutes (heartbeat updates every 2 minutes, so anything older
            # belongs to a dead process).
            cursor.execute(
                """
                UPDATE devices
                SET is_locked = TRUE, timestamp = NOW()
                WHERE device_uuid = %s
                  AND (is_locked = FALSE
                       OR timestamp < NOW() - INTERVAL '3 minutes')
                RETURNING device_uuid
                """,
                (device_uuid,),
            )
            acquired = cursor.fetchone() is not None
            conn.commit()

            if acquired:
                logger.info(
                    f"Lock acquired successfully for device {device_uuid}"
                )

                _heartbeat_stop_event = threading.Event()
                _heartbeat_thread = threading.Thread(
                    target=_heartbeat_worker,
                    args=(device_uuid, _heartbeat_stop_event),
                    daemon=True,
                )
                _heartbeat_thread.start()

                return True
            logger.info(
                f"Failed to acquire lock for device {device_uuid}. Device is already locked or doesn't exist."
            )
            return False


def release_lock(device_uuid):
    global _heartbeat_thread, _heartbeat_stop_event

    logger.info(f"Attempting to release lock for device {device_uuid}")
    with get_db_connection() as conn:
        with conn.cursor() as cursor:
            cursor.execute(
                "UPDATE devices SET is_locked = FALSE WHERE device_uuid = %s",
                (device_uuid,),
            )
            conn.commit()
            logger.info(f"Lock released for device {device_uuid}")

    if _heartbeat_thread and _heartbeat_stop_event:
        _heartbeat_stop_event.set()
        _heartbeat_thread.join(timeout=5)
        logger.info(f"Heartbeat thread stopped for {device_uuid}")
        _heartbeat_thread = None
        _heartbeat_stop_event = None


def device_exists(device_uuid):
    logger.info(f"Checking if device {device_uuid} exists")
    with get_db_connection() as conn:
        with conn.cursor() as cursor:
            cursor.execute(
                "SELECT 1 FROM devices WHERE device_uuid = %s", (device_uuid,)
            )
            exists = cursor.fetchone() is not None
            logger.info(
                f"Device {device_uuid} {'exists' if exists else 'does not exist'}"
            )
            return exists


def create_device(device_uuid):
    logger.info(f"Attempting to create device {device_uuid}")
    with get_db_connection() as conn:
        with conn.cursor() as cursor:
            cursor.execute(
                "INSERT INTO devices (device_uuid, is_locked) VALUES (%s, FALSE)",
                (device_uuid,),
            )
            conn.commit()
            logger.info(f"Device {device_uuid} created successfully")


def try_until_locked(device_uuid, max_attempts=80, sleep=90, fail_fast=True):
    if fail_fast:
        return acquire_lock(device_uuid)

    attempts = 0
    while attempts < max_attempts:
        if acquire_lock(device_uuid):
            logger.info(f"Device {device_uuid} successfully locked.")
            return True
        else:
            logger.info(
                f"Device {device_uuid} is already locked. Retrying in {sleep} seconds..."
            )
            time.sleep(sleep)
            attempts += 1

    raise Exception(
        f"Wasn't able to lock the device {device_uuid} after {max_attempts} attempts. :("
    )