84 lines
No EOL
2.6 KiB
Python
84 lines
No EOL
2.6 KiB
Python
"""
|
|
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 |