dumb ai

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)