import os import sqlite3 import time import json from datetime import datetime from app.config import ( DB_PATH, CACHE_EXPIRY_HOURS, USE_POSTGRES, POSTGRES_HOST, POSTGRES_PORT, POSTGRES_DB, POSTGRES_USER, POSTGRES_PASSWORD ) 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, 'database': POSTGRES_DB, 'user': POSTGRES_USER, 'password': POSTGRES_PASSWORD } else: self._init_sqlite_path() def _init_sqlite_path(self): """Initialize SQLite database path with fallback logic""" # 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}") self.db_path = db_path def _wait_for_postgres(self, max_retries=30, retry_delay=2): """Wait for PostgreSQL to be ready""" print(f"Waiting for PostgreSQL at {POSTGRES_HOST}:{POSTGRES_PORT}...") for attempt in range(max_retries): try: # First try to connect to the default 'postgres' database default_params = self.connection_params.copy() default_params['database'] = 'postgres' conn = self.psycopg2.connect(**default_params) conn.close() print(f"PostgreSQL is ready after {attempt + 1} attempts") return True except Exception as e: if attempt < max_retries - 1: print(f"PostgreSQL not ready (attempt {attempt + 1}/{max_retries}): {e}") time.sleep(retry_delay) else: print(f"PostgreSQL connection failed after {max_retries} attempts: {e}") return False return False def _create_database_if_not_exists(self): """Create the target database if it doesn't exist""" try: # Connect to the default 'postgres' database default_params = self.connection_params.copy() default_params['database'] = 'postgres' conn = self.psycopg2.connect(**default_params) conn.autocommit = True # Required for CREATE DATABASE cursor = conn.cursor() # Check if the target database exists cursor.execute("SELECT 1 FROM pg_database WHERE datname = %s", (POSTGRES_DB,)) exists = cursor.fetchone() if not exists: print(f"Creating database '{POSTGRES_DB}'...") cursor.execute(f"CREATE DATABASE {POSTGRES_DB}") print(f"Database '{POSTGRES_DB}' created successfully") else: print(f"Database '{POSTGRES_DB}' already exists") cursor.close() conn.close() return True except Exception as e: print(f"Error creating database: {e}") return False def get_connection(self): """Get database connection based on configured database type""" if self.db_type == "postgresql": return self.psycopg2.connect(**self.connection_params) else: return sqlite3.connect(self.db_path) def init_db(self): """Initialize database and create tables""" if self.db_type == "postgresql": self._init_postgres_db() else: self._init_sqlite_db() def _init_postgres_db(self): """Initialize PostgreSQL database""" try: # Wait for PostgreSQL to be ready if not self._wait_for_postgres(): print("PostgreSQL connection failed, falling back to SQLite") self.db_type = "sqlite" self._init_sqlite_path() self._init_sqlite_db() return # Create the target database if it doesn't exist if not self._create_database_if_not_exists(): print("Failed to create or check database, falling back to SQLite") self.db_type = "sqlite" self._init_sqlite_path() self._init_sqlite_db() return 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}") 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""" try: conn = self.get_connection() 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"SQLite database initialized at {self.db_path}") except sqlite3.OperationalError as 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 self.db_path = 'cache.db' print(f"Falling back to local database file: {self.db_path}") try: conn = sqlite3.connect(self.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 {self.db_path}") except sqlite3.OperationalError as e2: print(f"Error initializing local database: {e2}") raise 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 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 = %s AND route = %s AND timestamp > %s", (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 def save_to_cache(self, url, route, data): """Save data to cache""" conn = self.get_connection() cursor = conn.cursor() timestamp = int(time.time()) # Convert data to JSON string data_json = json.dumps(data) 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) ) 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}") 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()