ske087
150 lines
5.4 KiB
Python
150 lines
5.4 KiB
Python
"""
|
|
Database configuration and connection management
|
|
"""
|
|
import os
|
|
from sqlalchemy import create_engine, MetaData
|
|
from sqlalchemy.orm import sessionmaker
|
|
from contextlib import contextmanager
|
|
from app.models import Base
|
|
import logging
|
|
|
|
class DatabaseConfig:
|
|
"""Database configuration and connection management"""
|
|
|
|
def __init__(self, database_url=None):
|
|
if database_url is None:
|
|
# Default to SQLite with improved path
|
|
self.database_url = f"sqlite:///data/enhanced_monitoring.db"
|
|
else:
|
|
self.database_url = database_url
|
|
|
|
self.engine = None
|
|
self.Session = None
|
|
self._setup_database()
|
|
|
|
def _setup_database(self):
|
|
"""Initialize database connection and session factory"""
|
|
# Create data directory if it doesn't exist
|
|
os.makedirs('data', exist_ok=True)
|
|
|
|
# Create engine with connection pooling for SQLite
|
|
self.engine = create_engine(
|
|
self.database_url,
|
|
echo=False, # Set to True for SQL debugging
|
|
pool_pre_ping=True,
|
|
connect_args={"check_same_thread": False} # For SQLite
|
|
)
|
|
|
|
# Create session factory
|
|
self.Session = sessionmaker(bind=self.engine)
|
|
|
|
def create_tables(self):
|
|
"""Create all database tables"""
|
|
try:
|
|
Base.metadata.create_all(self.engine)
|
|
self.ensure_schema()
|
|
logging.info("Database tables created successfully")
|
|
return True
|
|
except Exception as e:
|
|
logging.error(f"Error creating database tables: {e}")
|
|
return False
|
|
|
|
def ensure_schema(self):
|
|
"""Idempotently add columns that were introduced after the table was first
|
|
created. SQLAlchemy's create_all() never ALTERs existing tables, so new
|
|
columns on an existing SQLite database must be added manually.
|
|
"""
|
|
from sqlalchemy import text
|
|
# Columns added over time: (table, column, SQL type definition)
|
|
required_columns = [
|
|
('devices', 'wmt_last_seen', 'DATETIME'),
|
|
('devices', 'config_synced_at', 'DATETIME'),
|
|
('devices', 'custom_chrome_url', 'VARCHAR(500)'),
|
|
]
|
|
try:
|
|
with self.engine.connect() as conn:
|
|
for table, column, col_type in required_columns:
|
|
exists = conn.execute(
|
|
text("SELECT name FROM sqlite_master WHERE type='table' AND name=:t"),
|
|
{'t': table},
|
|
).fetchone()
|
|
if not exists:
|
|
continue
|
|
cols = [row[1] for row in conn.execute(text(f"PRAGMA table_info({table})")).fetchall()]
|
|
if column not in cols:
|
|
conn.execute(text(f"ALTER TABLE {table} ADD COLUMN {column} {col_type}"))
|
|
conn.commit()
|
|
logging.info(f"Schema: added column {table}.{column}")
|
|
except Exception as e:
|
|
logging.error(f"Error ensuring schema columns: {e}")
|
|
|
|
def drop_tables(self):
|
|
"""Drop all database tables (use with caution!)"""
|
|
try:
|
|
Base.metadata.drop_all(self.engine)
|
|
logging.info("Database tables dropped successfully")
|
|
return True
|
|
except Exception as e:
|
|
logging.error(f"Error dropping database tables: {e}")
|
|
return False
|
|
|
|
@contextmanager
|
|
def get_session(self):
|
|
"""Context manager for database sessions"""
|
|
session = self.Session()
|
|
try:
|
|
yield session
|
|
session.commit()
|
|
except Exception as e:
|
|
session.rollback()
|
|
logging.error(f"Database session error: {e}")
|
|
raise
|
|
finally:
|
|
session.close()
|
|
|
|
def get_session_direct(self):
|
|
"""Get session directly (remember to close it)"""
|
|
return self.Session()
|
|
|
|
def backup_database(self, backup_path=None):
|
|
"""Create database backup"""
|
|
if backup_path is None:
|
|
from datetime import datetime
|
|
timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
|
|
backup_path = f"data/backups/backup_{timestamp}.db"
|
|
|
|
try:
|
|
# Create backups directory
|
|
os.makedirs('data/backups', exist_ok=True)
|
|
|
|
# For SQLite, simple file copy
|
|
if self.database_url.startswith('sqlite'):
|
|
import shutil
|
|
db_file = self.database_url.replace('sqlite:///', '')
|
|
shutil.copy2(db_file, backup_path)
|
|
logging.info(f"Database backup created: {backup_path}")
|
|
return backup_path
|
|
else:
|
|
# For other databases, implement proper backup
|
|
logging.warning("Backup not implemented for non-SQLite databases")
|
|
return None
|
|
|
|
except Exception as e:
|
|
logging.error(f"Database backup failed: {e}")
|
|
return None
|
|
|
|
# Global database instance
|
|
db_config = None
|
|
|
|
def init_database(database_url=None):
|
|
"""Initialize global database configuration"""
|
|
global db_config
|
|
db_config = DatabaseConfig(database_url)
|
|
return db_config.create_tables()
|
|
|
|
def get_db():
|
|
"""Get database configuration instance"""
|
|
global db_config
|
|
if db_config is None:
|
|
init_database()
|
|
return db_config |