diff --git a/Dockers/puppeteer-api/app/database.py b/Dockers/puppeteer-api/app/database.py index fb5db02..b097cae 100644 --- a/Dockers/puppeteer-api/app/database.py +++ b/Dockers/puppeteer-api/app/database.py @@ -8,17 +8,26 @@ from app.config import ( POSTGRES_HOST, POSTGRES_PORT, POSTGRES_DB, POSTGRES_USER, POSTGRES_PASSWORD ) -# Import PostgreSQL dependencies only if needed -if USE_POSTGRES: - import psycopg2 - from psycopg2.extras import RealDictCursor - class DatabaseManager: def __init__(self): self.db_type = "postgresql" if USE_POSTGRES else "sqlite" self.connection_params = None + self.db_path = None if USE_POSTGRES: + # Import PostgreSQL dependencies only when needed + try: + import psycopg2 + from psycopg2.extras import RealDictCursor + self.psycopg2 = psycopg2 + self.RealDictCursor = RealDictCursor + except ImportError as e: + print(f"PostgreSQL dependencies not available: {e}") + print("Falling back to SQLite") + self.db_type = "sqlite" + self._init_sqlite_path() + return + self.connection_params = { 'host': POSTGRES_HOST, 'port': POSTGRES_PORT, @@ -31,8 +40,6 @@ class DatabaseManager: def _init_sqlite_path(self): """Initialize SQLite database path with fallback logic""" - global DB_PATH - # Try to use the mounted volume first db_path = '/db/cache.db' db_dir = os.path.dirname(db_path) @@ -67,18 +74,18 @@ class DatabaseManager: db_path = 'cache.db' print(f"Using local database file: {db_path}") - DB_PATH = db_path + self.db_path = db_path def get_connection(self): """Get database connection based on configured database type""" - if USE_POSTGRES: - return psycopg2.connect(**self.connection_params) + if self.db_type == "postgresql": + return self.psycopg2.connect(**self.connection_params) else: - return sqlite3.connect(DB_PATH) + return sqlite3.connect(self.db_path) def init_db(self): """Initialize database and create tables""" - if USE_POSTGRES: + if self.db_type == "postgresql": self._init_postgres_db() else: self._init_sqlite_db() @@ -105,7 +112,10 @@ class DatabaseManager: print(f"PostgreSQL database initialized at {POSTGRES_HOST}:{POSTGRES_PORT}/{POSTGRES_DB}") except Exception as e: print(f"Error initializing PostgreSQL database: {e}") - raise + print("Falling back to SQLite") + self.db_type = "sqlite" + self._init_sqlite_path() + self._init_sqlite_db() def _init_sqlite_db(self): """Initialize SQLite database""" @@ -123,15 +133,14 @@ class DatabaseManager: ''') conn.commit() conn.close() - print(f"SQLite database initialized at {DB_PATH}") + print(f"SQLite database initialized at {self.db_path}") except sqlite3.OperationalError as e: - print(f"Error initializing database at {DB_PATH}: {e}") + print(f"Error initializing database at {self.db_path}: {e}") # Fallback to using a local database file if the mounted volume has permission issues - global DB_PATH - DB_PATH = 'cache.db' - print(f"Falling back to local database file: {DB_PATH}") + self.db_path = 'cache.db' + print(f"Falling back to local database file: {self.db_path}") try: - conn = sqlite3.connect(DB_PATH) + conn = sqlite3.connect(self.db_path) cursor = conn.cursor() cursor.execute(''' CREATE TABLE IF NOT EXISTS cache ( @@ -144,7 +153,7 @@ class DatabaseManager: ''') conn.commit() conn.close() - print(f"Local database initialized at {DB_PATH}") + print(f"Local database initialized at {self.db_path}") except sqlite3.OperationalError as e2: print(f"Error initializing local database: {e2}") raise @@ -182,7 +191,7 @@ class DatabaseManager: # Convert data to JSON string data_json = json.dumps(data) - if USE_POSTGRES: + if self.db_type == "postgresql": cursor.execute( "INSERT INTO cache (url, route, data, timestamp) VALUES (%s, %s, %s, %s) ON CONFLICT (url, route) DO UPDATE SET data = %s, timestamp = %s", (url, route, data_json, timestamp, data_json, timestamp)