allow postgres db

Dockers/puppeteer-api/app/database.py

@@ -3,70 +3,114 @@ import sqlite3
 import time
 import json
 from datetime import datetime
-from app.config import DB_PATH, CACHE_EXPIRY_HOURS
+from app.config import (
+    DB_PATH, CACHE_EXPIRY_HOURS, USE_POSTGRES,
+    POSTGRES_HOST, POSTGRES_PORT, POSTGRES_DB, POSTGRES_USER, POSTGRES_PASSWORD
+)
 
-# Initialize SQLite database
-def init_db():
-    global DB_PATH
+# Import PostgreSQL dependencies only if needed
+if USE_POSTGRES:
+    import psycopg2
+    from psycopg2.extras import RealDictCursor
 
-    # Try to use the mounted volume first
-    db_path = '/db/cache.db'
-    db_dir = os.path.dirname(db_path)
+class DatabaseManager:
+    def __init__(self):
+        self.db_type = "postgresql" if USE_POSTGRES else "sqlite"
+        self.connection_params = None
 
-    # Check if directory exists and is writable
-    dir_writable = False
-    if os.path.exists(db_dir):
-        try:
-            test_file = os.path.join(db_dir, '.write_test')
-            with open(test_file, 'w') as f:
-                f.write('test')
-            os.remove(test_file)
-            dir_writable = True
-        except (IOError, PermissionError):
-            print(f"Directory {db_dir} exists but is not writable")
-            dir_writable = False
+        if USE_POSTGRES:
+            self.connection_params = {
+                'host': POSTGRES_HOST,
+                'port': POSTGRES_PORT,
+                'database': POSTGRES_DB,
+                'user': POSTGRES_USER,
+                'password': POSTGRES_PASSWORD
+            }
+        else:
+            self._init_sqlite_path()
 
-    # If directory doesn't exist or isn't writable, try to create it
-    if not os.path.exists(db_dir) or not dir_writable:
-        try:
-            os.makedirs(db_dir, exist_ok=True)
-            # Test if we can write to the directory
-            test_file = os.path.join(db_dir, '.write_test')
-            with open(test_file, 'w') as f:
-                f.write('test')
-            os.remove(test_file)
-            print(f"Created directory: {db_dir}")
-            dir_writable = True
-        except Exception as e:
-            print(f"Warning: Could not create or write to directory {db_dir}: {e}")
-            # Fallback to using a local database file
-            db_path = 'cache.db'
-            print(f"Using local database file: {db_path}")
+    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)
+
+        # Check if directory exists and is writable
+        dir_writable = False
+        if os.path.exists(db_dir):
+            try:
+                test_file = os.path.join(db_dir, '.write_test')
+                with open(test_file, 'w') as f:
+                    f.write('test')
+                os.remove(test_file)
+                dir_writable = True
+            except (IOError, PermissionError):
+                print(f"Directory {db_dir} exists but is not writable")
+                dir_writable = False
+
+        # If directory doesn't exist or isn't writable, try to create it
+        if not os.path.exists(db_dir) or not dir_writable:
+            try:
+                os.makedirs(db_dir, exist_ok=True)
+                # Test if we can write to the directory
+                test_file = os.path.join(db_dir, '.write_test')
+                with open(test_file, 'w') as f:
+                    f.write('test')
+                os.remove(test_file)
+                print(f"Created directory: {db_dir}")
+                dir_writable = True
+            except Exception as e:
+                print(f"Warning: Could not create or write to directory {db_dir}: {e}")
+                # Fallback to using a local database file
+                db_path = 'cache.db'
+                print(f"Using local database file: {db_path}")
 
-    try:
-        conn = sqlite3.connect(db_path)
-        cursor = conn.cursor()
-        cursor.execute('''
-        CREATE TABLE IF NOT EXISTS cache (
-            url TEXT,
-            route TEXT,
-            data TEXT,
-            timestamp INTEGER,
-            PRIMARY KEY (url, route)
-        )
-        ''')
-        conn.commit()
-        conn.close()
-        print(f"Database initialized at {db_path}")
-        # Update the global DB_PATH
         DB_PATH = db_path
-    except sqlite3.OperationalError as e:
-        print(f"Error initializing database at {db_path}: {e}")
-        # Fallback to using a local database file if the mounted volume has permission issues
-        db_path = 'cache.db'
-        print(f"Falling back to local database file: {db_path}")
+
+    def get_connection(self):
+        """Get database connection based on configured database type"""
+        if USE_POSTGRES:
+            return psycopg2.connect(**self.connection_params)
+        else:
+            return sqlite3.connect(DB_PATH)
+
+    def init_db(self):
+        """Initialize database and create tables"""
+        if USE_POSTGRES:
+            self._init_postgres_db()
+        else:
+            self._init_sqlite_db()
+
+    def _init_postgres_db(self):
+        """Initialize PostgreSQL database"""
         try:
-            conn = sqlite3.connect(db_path)
+            conn = self.get_connection()
+            cursor = conn.cursor()
+
+            # Create cache table
+            cursor.execute('''
+            CREATE TABLE IF NOT EXISTS cache (
+                url TEXT,
+                route TEXT,
+                data TEXT,
+                timestamp BIGINT,
+                PRIMARY KEY (url, route)
+            )
+            ''')
+
+            conn.commit()
+            conn.close()
+            print(f"PostgreSQL database initialized at {POSTGRES_HOST}:{POSTGRES_PORT}/{POSTGRES_DB}")
+        except Exception as e:
+            print(f"Error initializing PostgreSQL database: {e}")
+            raise
+
+    def _init_sqlite_db(self):
+        """Initialize SQLite database"""
+        try:
+            conn = self.get_connection()
             cursor = conn.cursor()
             cursor.execute('''
             CREATE TABLE IF NOT EXISTS cache (
@@ -79,82 +123,128 @@ def init_db():
             ''')
             conn.commit()
             conn.close()
-            print(f"Local database initialized at {db_path}")
-            # Update the global DB_PATH
-            DB_PATH = db_path
-        except sqlite3.OperationalError as e2:
-            print(f"Error initializing local database: {e2}")
-            raise
+            print(f"SQLite database initialized at {DB_PATH}")
+        except sqlite3.OperationalError as e:
+            print(f"Error initializing database at {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}")
+            try:
+                conn = sqlite3.connect(DB_PATH)
+                cursor = conn.cursor()
+                cursor.execute('''
+                CREATE TABLE IF NOT EXISTS cache (
+                    url TEXT,
+                    route TEXT,
+                    data TEXT,
+                    timestamp INTEGER,
+                    PRIMARY KEY (url, route)
+                )
+                ''')
+                conn.commit()
+                conn.close()
+                print(f"Local database initialized at {DB_PATH}")
+            except sqlite3.OperationalError as e2:
+                print(f"Error initializing local database: {e2}")
+                raise
 
-# Get cached data if it exists and is not older than the expiry time
-def get_cached_data(url, route):
-    conn = sqlite3.connect(DB_PATH)
-    cursor = conn.cursor()
-
-    # Calculate cache expiry time
-    if route == 'pagination':
-        # Pagination cache expires after 31 times the normal cache expiry
-        cache_expiry = int(time.time()) - (CACHE_EXPIRY_HOURS * 31 * 60 * 60)
-    else:
-        cache_expiry = int(time.time()) - (CACHE_EXPIRY_HOURS * 60 * 60)  # Convert hours to seconds
-
-    cursor.execute(
-        "SELECT data FROM cache WHERE url = ? AND route = ? AND timestamp > ?",
-        (url, route, cache_expiry)
-    )
-    result = cursor.fetchone()
-    conn.close()
-
-    if result:
-        print(f"Cache hit for {url} on route {route}")
-        return json.loads(result[0])
-    return None
-
-# Save data to cache
-def save_to_cache(url, route, data):
-    conn = sqlite3.connect(DB_PATH)
-    cursor = conn.cursor()
-    timestamp = int(time.time())
-
-    # Convert data to JSON string
-    data_json = json.dumps(data)
-
-    cursor.execute(
-        "INSERT OR REPLACE INTO cache (url, route, data, timestamp) VALUES (?, ?, ?, ?)",
-        (url, route, data_json, timestamp)
-    )
-    conn.commit()
-    conn.close()
-    print(f"Saved to cache: {url} on route {route}")
-
-# Function to clean up old cache entries
-def cleanup_old_cache_entries():
-    try:
-        print(f"Running scheduled cache cleanup (entries older than {CACHE_EXPIRY_HOURS} hours, pagination: {CACHE_EXPIRY_HOURS * 31} hours)")
-        conn = sqlite3.connect(DB_PATH)
+    def get_cached_data(self, url, route):
+        """Get cached data if it exists and is not older than the expiry time"""
+        conn = self.get_connection()
         cursor = conn.cursor()
 
-        # Calculate the timestamp for entries older than the expiry time
-        expiry_timestamp = int(time.time()) - (CACHE_EXPIRY_HOURS * 60 * 60)
-        pagination_expiry_timestamp = int(time.time()) - (CACHE_EXPIRY_HOURS * 31 * 60 * 60)
+        # Calculate cache expiry time
+        if route == 'pagination':
+            # Pagination cache expires after 31 times the normal cache expiry
+            cache_expiry = int(time.time()) - (CACHE_EXPIRY_HOURS * 31 * 60 * 60)
+        else:
+            cache_expiry = int(time.time()) - (CACHE_EXPIRY_HOURS * 60 * 60)  # Convert hours to seconds
 
-        # Get count of entries to be deleted (non-pagination)
-        cursor.execute("SELECT COUNT(*) FROM cache WHERE route != 'pagination' AND timestamp < ?", (expiry_timestamp,))
-        count_non_pagination = cursor.fetchone()[0]
+        cursor.execute(
+            "SELECT data FROM cache WHERE url = %s AND route = %s AND timestamp > %s",
+            (url, route, cache_expiry)
+        )
+        result = cursor.fetchone()
+        conn.close()
 
-        # Get count of pagination entries to be deleted
-        cursor.execute("SELECT COUNT(*) FROM cache WHERE route = 'pagination' AND timestamp < ?", (pagination_expiry_timestamp,))
-        count_pagination = cursor.fetchone()[0]
+        if result:
+            print(f"Cache hit for {url} on route {route}")
+            return json.loads(result[0])
+        return None
 
-        # Delete old non-pagination entries
-        cursor.execute("DELETE FROM cache WHERE route != 'pagination' AND timestamp < ?", (expiry_timestamp,))
+    def save_to_cache(self, url, route, data):
+        """Save data to cache"""
+        conn = self.get_connection()
+        cursor = conn.cursor()
+        timestamp = int(time.time())
 
-        # Delete old pagination entries
-        cursor.execute("DELETE FROM cache WHERE route = 'pagination' AND timestamp < ?", (pagination_expiry_timestamp,))
+        # Convert data to JSON string
+        data_json = json.dumps(data)
+
+        if USE_POSTGRES:
+            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)
+            )
+        else:
+            cursor.execute(
+                "INSERT OR REPLACE INTO cache (url, route, data, timestamp) VALUES (?, ?, ?, ?)",
+                (url, route, data_json, timestamp)
+            )
 
         conn.commit()
         conn.close()
+        print(f"Saved to cache: {url} on route {route}")
 
-        print(f"Cache cleanup completed: {count_non_pagination} non-pagination entries and {count_pagination} pagination entries removed")
-    except Exception as e:
-        print(f"Error during cache cleanup: {e}")
+    def cleanup_old_cache_entries(self):
+        """Clean up old cache entries"""
+        try:
+            print(f"Running scheduled cache cleanup (entries older than {CACHE_EXPIRY_HOURS} hours, pagination: {CACHE_EXPIRY_HOURS * 31} hours)")
+            conn = self.get_connection()
+            cursor = conn.cursor()
+
+            # Calculate the timestamp for entries older than the expiry time
+            expiry_timestamp = int(time.time()) - (CACHE_EXPIRY_HOURS * 60 * 60)
+            pagination_expiry_timestamp = int(time.time()) - (CACHE_EXPIRY_HOURS * 31 * 60 * 60)
+
+            # Get count of entries to be deleted (non-pagination)
+            cursor.execute("SELECT COUNT(*) FROM cache WHERE route != %s AND timestamp < %s", ('pagination', expiry_timestamp))
+            count_non_pagination = cursor.fetchone()[0]
+
+            # Get count of pagination entries to be deleted
+            cursor.execute("SELECT COUNT(*) FROM cache WHERE route = %s AND timestamp < %s", ('pagination', pagination_expiry_timestamp))
+            count_pagination = cursor.fetchone()[0]
+
+            # Delete old non-pagination entries
+            cursor.execute("DELETE FROM cache WHERE route != %s AND timestamp < %s", ('pagination', expiry_timestamp))
+
+            # Delete old pagination entries
+            cursor.execute("DELETE FROM cache WHERE route = %s AND timestamp < %s", ('pagination', pagination_expiry_timestamp))
+
+            conn.commit()
+            conn.close()
+
+            print(f"Cache cleanup completed: {count_non_pagination} non-pagination entries and {count_pagination} pagination entries removed")
+        except Exception as e:
+            print(f"Error during cache cleanup: {e}")
+
+# Create global database manager instance
+db_manager = DatabaseManager()
+
+# Backward compatibility functions
+def init_db():
+    """Initialize database (backward compatibility)"""
+    db_manager.init_db()
+
+def get_cached_data(url, route):
+    """Get cached data (backward compatibility)"""
+    return db_manager.get_cached_data(url, route)
+
+def save_to_cache(url, route, data):
+    """Save data to cache (backward compatibility)"""
+    db_manager.save_to_cache(url, route, data)
+
+def cleanup_old_cache_entries():
+    """Clean up old cache entries (backward compatibility)"""
+    db_manager.cleanup_old_cache_entries()