""" Database utility functions for the OpenEventDatabase. """ import os import sys import psycopg2 import psycopg2.extras from oedb.utils.logging import logger def load_env_from_file(): """ Load environment variables from .env file. This ensures that database connection parameters are properly set. """ if os.path.exists('.env'): logger.info("Loading environment variables from .env file...") with open('.env', 'r') as f: for line in f: line = line.strip() if line and not line.startswith('#'): key, value = line.split('=', 1) os.environ[key] = value def db_connect(): """ Connect to the PostgreSQL database using environment variables. Returns: psycopg2.connection: A connection to the database. """ # Load environment variables from .env file load_env_from_file() # Get connection parameters from environment variables dbname = os.getenv("DB_NAME", "oedb") host = os.getenv("DB_HOST", "") password = os.getenv("POSTGRES_PASSWORD", None) user = os.getenv("DB_USER", "") # If host is empty, set it to localhost to force TCP/IP connection # instead of Unix socket (which uses peer authentication) if not host: host = "localhost" logger.debug(f"Connecting to database: {dbname} on {host} as user {user}") # For localhost connections, add additional parameters to ensure proper connection if host in ('localhost', '127.0.0.1'): logger.debug("Using TCP/IP connection with additional parameters") return psycopg2.connect( dbname=dbname, host=host, password=password, user=user, options="-c client_encoding=utf8 -c statement_timeout=3000", connect_timeout=3, application_name="oedb", # Disable SSL for local connections sslmode='disable') else: # For remote connections, use standard parameters logger.debug("Using standard connection parameters") return psycopg2.connect( dbname=dbname, host=host, password=password, user=user) def check_db_connection(): """ Check if the database is accessible. Returns: bool: True if the database is accessible, False otherwise. """ try: conn = db_connect() conn.close() logger.success("Successfully connected to PostgreSQL database") return True except psycopg2.OperationalError as e: logger.error(f"Failed to connect to PostgreSQL database: {e}") return False