345 lines
14 KiB
Python
345 lines
14 KiB
Python
|
"""
|
||
|
|
Flexible Database Client for Calejo Control Adapter.
|
||
|
|
|
||
|
|
Supports multiple database backends (PostgreSQL, SQLite) using SQLAlchemy Core.
|
||
|
|
"""
|
||
|
|
|
||
|
|
from typing import Dict, List, Optional, Any
|
||
|
|
from datetime import datetime
|
||
|
|
import structlog
|
||
|
|
from sqlalchemy import create_engine, text, MetaData, Table, Column, String, Float, Integer, Boolean, DateTime
|
||
|
|
from sqlalchemy.engine import Engine
|
||
|
|
from sqlalchemy.exc import SQLAlchemyError
|
||
|
|
|
||
|
|
logger = structlog.get_logger()
|
||
|
|
|
||
|
|
|
||
|
|
class FlexibleDatabaseClient:
|
||
|
|
"""
|
||
|
|
Flexible database client supporting multiple backends.
|
||
|
|
|
||
|
|
Supports:
|
||
|
|
- PostgreSQL: postgresql://user:pass@host:port/dbname
|
||
|
|
- SQLite: sqlite:///path/to/database.db
|
||
|
|
"""
|
||
|
|
|
||
|
|
def __init__(
|
||
|
|
self,
|
||
|
|
database_url: str,
|
||
|
|
pool_size: int = 5,
|
||
|
|
max_overflow: int = 10,
|
||
|
|
pool_timeout: int = 30,
|
||
|
|
pool_recycle: int = 3600
|
||
|
|
):
|
||
|
|
self.database_url = database_url
|
||
|
|
self.pool_size = pool_size
|
||
|
|
self.max_overflow = max_overflow
|
||
|
|
self.pool_timeout = pool_timeout
|
||
|
|
self.pool_recycle = pool_recycle
|
||
|
|
self.engine: Optional[Engine] = None
|
||
|
|
self.metadata = MetaData()
|
||
|
|
|
||
|
|
# Define table schemas
|
||
|
|
self._define_tables()
|
||
|
|
|
||
|
|
def _define_tables(self):
|
||
|
|
"""Define database table schemas."""
|
||
|
|
self.stations = Table(
|
||
|
|
'stations', self.metadata,
|
||
|
|
Column('station_id', String(50), primary_key=True),
|
||
|
|
Column('station_name', String(100)),
|
||
|
|
Column('location', String(200)),
|
||
|
|
Column('created_at', DateTime, default=datetime.now)
|
||
|
|
)
|
||
|
|
|
||
|
|
self.pumps = Table(
|
||
|
|
'pumps', self.metadata,
|
||
|
|
Column('station_id', String(50)),
|
||
|
|
Column('pump_id', String(50)),
|
||
|
|
Column('pump_name', String(100)),
|
||
|
|
Column('control_type', String(50)),
|
||
|
|
Column('min_speed_hz', Float, default=20.0),
|
||
|
|
Column('max_speed_hz', Float, default=60.0),
|
||
|
|
Column('default_setpoint_hz', Float, default=35.0),
|
||
|
|
Column('created_at', DateTime, default=datetime.now)
|
||
|
|
)
|
||
|
|
|
||
|
|
self.pump_plans = Table(
|
||
|
|
'pump_plans', self.metadata,
|
||
|
|
Column('plan_id', Integer, primary_key=True, autoincrement=True),
|
||
|
|
Column('station_id', String(50)),
|
||
|
|
Column('pump_id', String(50)),
|
||
|
|
Column('target_flow_m3h', Float),
|
||
|
|
Column('target_power_kw', Float),
|
||
|
|
Column('target_level_m', Float),
|
||
|
|
Column('suggested_speed_hz', Float),
|
||
|
|
Column('interval_start', DateTime),
|
||
|
|
Column('interval_end', DateTime),
|
||
|
|
Column('plan_version', Integer),
|
||
|
|
Column('plan_status', String(20), default='ACTIVE'),
|
||
|
|
Column('plan_created_at', DateTime, default=datetime.now),
|
||
|
|
Column('plan_updated_at', DateTime, default=datetime.now, onupdate=datetime.now),
|
||
|
|
Column('optimization_run_id', String(100))
|
||
|
|
)
|
||
|
|
|
||
|
|
self.pump_feedback = Table(
|
||
|
|
'pump_feedback', self.metadata,
|
||
|
|
Column('feedback_id', Integer, primary_key=True, autoincrement=True),
|
||
|
|
Column('station_id', String(50)),
|
||
|
|
Column('pump_id', String(50)),
|
||
|
|
Column('timestamp', DateTime, default=datetime.now),
|
||
|
|
Column('actual_speed_hz', Float),
|
||
|
|
Column('actual_power_kw', Float),
|
||
|
|
Column('actual_flow_m3h', Float),
|
||
|
|
Column('wet_well_level_m', Float),
|
||
|
|
Column('pump_running', Boolean),
|
||
|
|
Column('alarm_active', Boolean),
|
||
|
|
Column('alarm_code', String(50))
|
||
|
|
)
|
||
|
|
|
||
|
|
self.emergency_stop_events = Table(
|
||
|
|
'emergency_stop_events', self.metadata,
|
||
|
|
Column('event_id', Integer, primary_key=True, autoincrement=True),
|
||
|
|
Column('triggered_by', String(100)),
|
||
|
|
Column('reason', String),
|
||
|
|
Column('station_id', String(50)),
|
||
|
|
Column('pump_id', String(50)),
|
||
|
|
Column('event_timestamp', DateTime, default=datetime.now),
|
||
|
|
Column('cleared_by', String(100)),
|
||
|
|
Column('cleared_timestamp', DateTime),
|
||
|
|
Column('cleared_notes', String)
|
||
|
|
)
|
||
|
|
|
||
|
|
self.safety_limit_violations = Table(
|
||
|
|
'safety_limit_violations', self.metadata,
|
||
|
|
Column('violation_id', Integer, primary_key=True, autoincrement=True),
|
||
|
|
Column('station_id', String(50)),
|
||
|
|
Column('pump_id', String(50)),
|
||
|
|
Column('requested_setpoint', Float),
|
||
|
|
Column('enforced_setpoint', Float),
|
||
|
|
Column('violations', String),
|
||
|
|
Column('timestamp', DateTime, default=datetime.now)
|
||
|
|
)
|
||
|
|
|
||
|
|
async def connect(self):
|
||
|
|
"""Connect to the database."""
|
||
|
|
try:
|
||
|
|
# Configure engine based on database type
|
||
|
|
if self.database_url.startswith('sqlite://'):
|
||
|
|
# SQLite configuration
|
||
|
|
self.engine = create_engine(
|
||
|
|
self.database_url,
|
||
|
|
poolclass=None, # No connection pooling for SQLite
|
||
|
|
connect_args={"check_same_thread": False}
|
||
|
|
)
|
||
|
|
else:
|
||
|
|
# PostgreSQL configuration
|
||
|
|
self.engine = create_engine(
|
||
|
|
self.database_url,
|
||
|
|
pool_size=self.pool_size,
|
||
|
|
max_overflow=self.max_overflow,
|
||
|
|
pool_timeout=self.pool_timeout,
|
||
|
|
pool_recycle=self.pool_recycle
|
||
|
|
)
|
||
|
|
|
||
|
|
# Test connection
|
||
|
|
with self.engine.connect() as conn:
|
||
|
|
conn.execute(text("SELECT 1"))
|
||
|
|
|
||
|
|
logger.info(
|
||
|
|
"database_connected",
|
||
|
|
database_type=self._get_database_type(),
|
||
|
|
url=self._get_safe_url()
|
||
|
|
)
|
||
|
|
|
||
|
|
except SQLAlchemyError as e:
|
||
|
|
logger.error("database_connection_failed", error=str(e))
|
||
|
|
raise
|
||
|
|
|
||
|
|
async def disconnect(self):
|
||
|
|
"""Disconnect from the database."""
|
||
|
|
if self.engine:
|
||
|
|
self.engine.dispose()
|
||
|
|
logger.info("database_disconnected")
|
||
|
|
|
||
|
|
def _get_database_type(self) -> str:
|
||
|
|
"""Get database type from URL."""
|
||
|
|
if self.database_url.startswith('sqlite://'):
|
||
|
|
return 'SQLite'
|
||
|
|
elif self.database_url.startswith('postgresql://'):
|
||
|
|
return 'PostgreSQL'
|
||
|
|
else:
|
||
|
|
return 'Unknown'
|
||
|
|
|
||
|
|
def _get_safe_url(self) -> str:
|
||
|
|
"""Get safe URL for logging (without credentials)."""
|
||
|
|
if self.database_url.startswith('postgresql://'):
|
||
|
|
# Remove credentials from PostgreSQL URL
|
||
|
|
parts = self.database_url.split('@')
|
||
|
|
if len(parts) > 1:
|
||
|
|
return f"postgresql://...@{parts[1]}"
|
||
|
|
return self.database_url
|
||
|
|
|
||
|
|
def execute_query(self, query: str, params: Optional[Dict[str, Any]] = None) -> List[Dict[str, Any]]:
|
||
|
|
"""Execute a query and return results as dictionaries."""
|
||
|
|
try:
|
||
|
|
with self.engine.connect() as conn:
|
||
|
|
result = conn.execute(text(query), params or {})
|
||
|
|
return [dict(row._mapping) for row in result]
|
||
|
|
except SQLAlchemyError as e:
|
||
|
|
logger.error("query_execution_failed", query=query, error=str(e))
|
||
|
|
raise
|
||
|
|
|
||
|
|
def execute(self, query: str, params: Optional[Dict[str, Any]] = None) -> int:
|
||
|
|
"""Execute a query and return number of affected rows."""
|
||
|
|
try:
|
||
|
|
with self.engine.connect() as conn:
|
||
|
|
result = conn.execute(text(query), params or {})
|
||
|
|
conn.commit()
|
||
|
|
return result.rowcount
|
||
|
|
except SQLAlchemyError as e:
|
||
|
|
logger.error("query_execution_failed", query=query, error=str(e))
|
||
|
|
raise
|
||
|
|
|
||
|
|
def health_check(self) -> bool:
|
||
|
|
"""Check if database is healthy and responsive."""
|
||
|
|
try:
|
||
|
|
with self.engine.connect() as conn:
|
||
|
|
result = conn.execute(text("SELECT 1 as health_check"))
|
||
|
|
row = result.fetchone()
|
||
|
|
return row[0] == 1
|
||
|
|
except SQLAlchemyError as e:
|
||
|
|
logger.error("database_health_check_failed", error=str(e))
|
||
|
|
return False
|
||
|
|
|
||
|
|
def get_connection_stats(self) -> Dict[str, Any]:
|
||
|
|
"""Get connection pool statistics."""
|
||
|
|
if not self.engine:
|
||
|
|
return {"status": "not_connected"}
|
||
|
|
|
||
|
|
return {
|
||
|
|
"database_type": self._get_database_type(),
|
||
|
|
"pool_size": self.pool_size,
|
||
|
|
"max_overflow": self.max_overflow,
|
||
|
|
"status": "connected"
|
||
|
|
}
|
||
|
|
|
||
|
|
# Database-specific methods
|
||
|
|
def get_pump_stations(self) -> List[Dict[str, Any]]:
|
||
|
|
"""Get all pump stations."""
|
||
|
|
query = "SELECT * FROM stations ORDER BY station_id"
|
||
|
|
return self.execute_query(query)
|
||
|
|
|
||
|
|
def get_pumps(self, station_id: Optional[str] = None) -> List[Dict[str, Any]]:
|
||
|
|
"""Get pumps, optionally filtered by station."""
|
||
|
|
if station_id:
|
||
|
|
query = "SELECT * FROM pumps WHERE station_id = :station_id ORDER BY pump_id"
|
||
|
|
return self.execute_query(query, {"station_id": station_id})
|
||
|
|
else:
|
||
|
|
query = "SELECT * FROM pumps ORDER BY station_id, pump_id"
|
||
|
|
return self.execute_query(query)
|
||
|
|
|
||
|
|
def get_pump(self, station_id: str, pump_id: str) -> Optional[Dict[str, Any]]:
|
||
|
|
"""Get specific pump."""
|
||
|
|
query = """
|
||
|
|
SELECT * FROM pumps
|
||
|
|
WHERE station_id = :station_id AND pump_id = :pump_id
|
||
|
|
"""
|
||
|
|
result = self.execute_query(query, {
|
||
|
|
"station_id": station_id,
|
||
|
|
"pump_id": pump_id
|
||
|
|
})
|
||
|
|
return result[0] if result else None
|
||
|
|
|
||
|
|
def get_current_plan(self, station_id: str, pump_id: str) -> Optional[Dict[str, Any]]:
|
||
|
|
"""Get current active plan for a specific pump."""
|
||
|
|
query = """
|
||
|
|
SELECT plan_id, target_flow_m3h, target_power_kw, target_level_m,
|
||
|
|
suggested_speed_hz, interval_start, interval_end, plan_version,
|
||
|
|
plan_status, plan_created_at, plan_updated_at, optimization_run_id
|
||
|
|
FROM pump_plans
|
||
|
|
WHERE station_id = :station_id AND pump_id = :pump_id
|
||
|
|
AND interval_start <= datetime('now') AND interval_end >= datetime('now')
|
||
|
|
AND plan_status = 'ACTIVE'
|
||
|
|
ORDER BY plan_version DESC
|
||
|
|
LIMIT 1
|
||
|
|
"""
|
||
|
|
result = self.execute_query(query, {
|
||
|
|
"station_id": station_id,
|
||
|
|
"pump_id": pump_id
|
||
|
|
})
|
||
|
|
return result[0] if result else None
|
||
|
|
|
||
|
|
def get_latest_feedback(self, station_id: str, pump_id: str) -> Optional[Dict[str, Any]]:
|
||
|
|
"""Get latest feedback for a specific pump."""
|
||
|
|
query = """
|
||
|
|
SELECT timestamp, actual_speed_hz, actual_power_kw, actual_flow_m3h,
|
||
|
|
wet_well_level_m, pump_running, alarm_active, alarm_code
|
||
|
|
FROM pump_feedback
|
||
|
|
WHERE station_id = :station_id AND pump_id = :pump_id
|
||
|
|
ORDER BY timestamp DESC
|
||
|
|
LIMIT 1
|
||
|
|
"""
|
||
|
|
result = self.execute_query(query, {
|
||
|
|
"station_id": station_id,
|
||
|
|
"pump_id": pump_id
|
||
|
|
})
|
||
|
|
return result[0] if result else None
|
||
|
|
|
||
|
|
def get_pump_feedback(self, station_id: str, pump_id: str, limit: int = 10) -> List[Dict[str, Any]]:
|
||
|
|
"""Get recent feedback for a specific pump."""
|
||
|
|
query = """
|
||
|
|
SELECT timestamp, actual_speed_hz, actual_power_kw, actual_flow_m3h,
|
||
|
|
wet_well_level_m, pump_running, alarm_active, alarm_code
|
||
|
|
FROM pump_feedback
|
||
|
|
WHERE station_id = :station_id AND pump_id = :pump_id
|
||
|
|
ORDER BY timestamp DESC
|
||
|
|
LIMIT :limit
|
||
|
|
"""
|
||
|
|
return self.execute_query(query, {
|
||
|
|
"station_id": station_id,
|
||
|
|
"pump_id": pump_id,
|
||
|
|
"limit": limit
|
||
|
|
})
|
||
|
|
|
||
|
|
def get_active_plans(self, resource_type: str = 'pump') -> List[Dict[str, Any]]:
|
||
|
|
"""Get all active plans for a resource type."""
|
||
|
|
query = """
|
||
|
|
SELECT station_id, pump_id, suggested_speed_hz, interval_start, interval_end,
|
||
|
|
plan_version, plan_status, optimization_run_id
|
||
|
|
FROM pump_plans
|
||
|
|
WHERE interval_start <= CURRENT_TIMESTAMP AND interval_end >= CURRENT_TIMESTAMP
|
||
|
|
AND plan_status = 'ACTIVE'
|
||
|
|
AND resource_type = :resource_type
|
||
|
|
ORDER BY station_id, resource_id, plan_version DESC
|
||
|
|
"""
|
||
|
|
return self.execute_query(query, {"resource_type": resource_type})
|
||
|
|
|
||
|
|
def get_safety_limits(self) -> List[Dict[str, Any]]:
|
||
|
|
"""Get safety limits for all pumps."""
|
||
|
|
query = """
|
||
|
|
SELECT station_id, pump_id, min_speed_hz as hard_min_speed_hz,
|
||
|
|
max_speed_hz as hard_max_speed_hz,
|
||
|
|
NULL as hard_min_level_m, NULL as hard_max_level_m,
|
||
|
|
NULL as hard_max_power_kw, 5.0 as max_speed_change_hz_per_min
|
||
|
|
FROM pumps
|
||
|
|
"""
|
||
|
|
return self.execute_query(query)
|
||
|
|
|
||
|
|
def create_tables(self):
|
||
|
|
"""Create all tables if they don't exist."""
|
||
|
|
try:
|
||
|
|
self.metadata.create_all(self.engine)
|
||
|
|
logger.info("database_tables_created")
|
||
|
|
except SQLAlchemyError as e:
|
||
|
|
logger.error("failed_to_create_tables", error=str(e))
|
||
|
|
raise
|
||
|
|
|
||
|
|
def drop_tables(self):
|
||
|
|
"""Drop all tables (for testing)."""
|
||
|
|
try:
|
||
|
|
self.metadata.drop_all(self.engine)
|
||
|
|
logger.info("database_tables_dropped")
|
||
|
|
except SQLAlchemyError as e:
|
||
|
|
logger.error("failed_to_drop_tables", error=str(e))
|
||
|
|
raise
|