347 lines
14 KiB
Python
347 lines
14 KiB
Python
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
|
|
|
|
if self.db_type == "postgresql":
|
|
cursor.execute(
|
|
"SELECT data FROM cache WHERE url = %s AND route = %s AND timestamp > %s",
|
|
(url, route, cache_expiry)
|
|
)
|
|
else:
|
|
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
|
|
|
|
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)
|
|
|
|
if self.db_type == "postgresql":
|
|
# 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))
|
|
else:
|
|
# Get count of entries to be deleted (non-pagination)
|
|
cursor.execute("SELECT COUNT(*) FROM cache WHERE route != ? AND timestamp < ?", ('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 = ? AND timestamp < ?", ('pagination', pagination_expiry_timestamp))
|
|
count_pagination = cursor.fetchone()[0]
|
|
|
|
# Delete old non-pagination entries
|
|
cursor.execute("DELETE FROM cache WHERE route != ? AND timestamp < ?", ('pagination', expiry_timestamp))
|
|
|
|
# Delete old pagination entries
|
|
cursor.execute("DELETE FROM cache WHERE route = ? AND timestamp < ?", ('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() |