grabowski
6e78225d00
- MeasurementResponse.discharge is now Optional[float]; measurements with a null discharge no longer raise a Pydantic ValidationError (HTTP 500) on the /measurements/latest and /measurements/station endpoints. - InfluxDB save_measurements guards float(discharge) against None instead of crashing with TypeError. - Extract the duplicated measurement->response mapping into a single _to_measurement_response helper used by both measurement endpoints.
796 lines
32 KiB
Python
796 lines
32 KiB
Python
#!/usr/bin/env python3
|
|
"""
|
|
Database adapters for different storage backends
|
|
"""
|
|
|
|
import datetime
|
|
import logging
|
|
from abc import ABC, abstractmethod
|
|
from typing import Dict, List, Optional
|
|
|
|
|
|
# Base adapter interface
|
|
class DatabaseAdapter(ABC):
|
|
@abstractmethod
|
|
def connect(self):
|
|
pass
|
|
|
|
@abstractmethod
|
|
def save_measurements(self, measurements: List[Dict]) -> bool:
|
|
pass
|
|
|
|
@abstractmethod
|
|
def get_latest_measurements(self, limit: int = 100) -> List[Dict]:
|
|
pass
|
|
|
|
@abstractmethod
|
|
def get_measurements_by_timerange(
|
|
self,
|
|
start_time: datetime.datetime,
|
|
end_time: datetime.datetime,
|
|
station_codes: Optional[List[str]] = None,
|
|
) -> List[Dict]:
|
|
pass
|
|
|
|
@abstractmethod
|
|
def get_measurements_for_date(self, target_date: datetime.datetime) -> List[Dict]:
|
|
pass
|
|
|
|
|
|
# InfluxDB Adapter
|
|
class InfluxDBAdapter(DatabaseAdapter):
|
|
def __init__(
|
|
self,
|
|
host: str = "localhost",
|
|
port: int = 8086,
|
|
database: str = "water_monitoring",
|
|
username: str = None,
|
|
password: str = None,
|
|
):
|
|
self.host = host
|
|
self.port = port
|
|
self.database = database
|
|
self.username = username
|
|
self.password = password
|
|
self.client = None
|
|
|
|
def connect(self):
|
|
try:
|
|
from influxdb import InfluxDBClient
|
|
|
|
self.client = InfluxDBClient(
|
|
host=self.host,
|
|
port=self.port,
|
|
username=self.username,
|
|
password=self.password,
|
|
database=self.database,
|
|
)
|
|
|
|
# Create database if it doesn't exist
|
|
databases = self.client.get_list_database()
|
|
if not any(db["name"] == self.database for db in databases):
|
|
self.client.create_database(self.database)
|
|
logging.info(f"Created InfluxDB database: {self.database}")
|
|
|
|
# Create retention policy (keep data for 2 years, downsample after 30 days)
|
|
retention_policies = self.client.get_list_retention_policies(self.database)
|
|
if not any(rp["name"] == "water_data_policy" for rp in retention_policies):
|
|
self.client.create_retention_policy(
|
|
"water_data_policy",
|
|
"730d", # 2 years
|
|
"1", # replication factor
|
|
database=self.database,
|
|
default=True,
|
|
)
|
|
|
|
logging.info("Connected to InfluxDB successfully")
|
|
return True
|
|
|
|
except ImportError:
|
|
logging.error("InfluxDB client not installed. Run: pip install influxdb")
|
|
return False
|
|
except Exception as e:
|
|
logging.error(f"Failed to connect to InfluxDB: {e}")
|
|
return False
|
|
|
|
def save_measurements(self, measurements: List[Dict]) -> bool:
|
|
if not self.client:
|
|
logging.error("InfluxDB client not connected")
|
|
return False
|
|
|
|
try:
|
|
points = []
|
|
for measurement in measurements:
|
|
point = {
|
|
"measurement": "water_data",
|
|
"tags": {
|
|
"station_code": measurement["station_code"],
|
|
"station_name_en": measurement["station_name_en"],
|
|
"station_name_th": measurement["station_name_th"],
|
|
},
|
|
"time": measurement["timestamp"].isoformat(),
|
|
"fields": {
|
|
"water_level": float(measurement["water_level"]),
|
|
"discharge": float(measurement["discharge"])
|
|
if measurement.get("discharge") is not None
|
|
else None,
|
|
"discharge_percent": float(measurement["discharge_percent"])
|
|
if measurement.get("discharge_percent")
|
|
else None,
|
|
},
|
|
}
|
|
points.append(point)
|
|
|
|
success = self.client.write_points(points)
|
|
if success:
|
|
logging.info(f"Successfully wrote {len(points)} points to InfluxDB")
|
|
return success
|
|
|
|
except Exception as e:
|
|
logging.error(f"Error writing to InfluxDB: {e}")
|
|
return False
|
|
|
|
def get_latest_measurements(self, limit: int = 100) -> List[Dict]:
|
|
if not self.client:
|
|
return []
|
|
|
|
try:
|
|
# Cast limit to int so it can never carry an injection payload.
|
|
limit = int(limit)
|
|
query = f"""
|
|
SELECT last("water_level") as water_level,
|
|
last("discharge") as discharge,
|
|
last("discharge_percent") as discharge_percent
|
|
FROM "water_data"
|
|
GROUP BY "station_code", "station_name_en", "station_name_th"
|
|
LIMIT {limit}
|
|
"""
|
|
|
|
result = self.client.query(query)
|
|
measurements = []
|
|
|
|
for point in result.get_points():
|
|
measurements.append(
|
|
{
|
|
"timestamp": point["time"],
|
|
"station_code": point.get("station_code"),
|
|
"station_name_en": point.get("station_name_en"),
|
|
"station_name_th": point.get("station_name_th"),
|
|
"water_level": point.get("water_level"),
|
|
"discharge": point.get("discharge"),
|
|
"discharge_percent": point.get("discharge_percent"),
|
|
}
|
|
)
|
|
|
|
return measurements
|
|
|
|
except Exception as e:
|
|
logging.error(f"Error querying InfluxDB: {e}")
|
|
return []
|
|
|
|
def get_measurements_by_timerange(
|
|
self,
|
|
start_time: datetime.datetime,
|
|
end_time: datetime.datetime,
|
|
station_codes: Optional[List[str]] = None,
|
|
) -> List[Dict]:
|
|
if not self.client:
|
|
return []
|
|
|
|
try:
|
|
# start_time/end_time are datetime objects (fixed isoformat, injection-safe).
|
|
# station_codes are untrusted strings -> bind them as parameters.
|
|
bind_params = {}
|
|
where_clause = f"time >= '{start_time.isoformat()}' AND time <= '{end_time.isoformat()}'"
|
|
if station_codes:
|
|
placeholders = []
|
|
for i, code in enumerate(station_codes):
|
|
key = f"sc{i}"
|
|
bind_params[key] = code
|
|
placeholders.append(f"station_code = ${key}")
|
|
where_clause += " AND (" + " OR ".join(placeholders) + ")"
|
|
|
|
query = f"""
|
|
SELECT "water_level", "discharge", "discharge_percent", "station_code", "station_name_en", "station_name_th"
|
|
FROM "water_data"
|
|
WHERE {where_clause}
|
|
ORDER BY time DESC
|
|
"""
|
|
|
|
result = self.client.query(query, bind_params=bind_params)
|
|
measurements = []
|
|
|
|
for point in result.get_points():
|
|
measurements.append(
|
|
{
|
|
"timestamp": point["time"],
|
|
"station_code": point.get("station_code"),
|
|
"station_name_en": point.get("station_name_en"),
|
|
"station_name_th": point.get("station_name_th"),
|
|
"water_level": point.get("water_level"),
|
|
"discharge": point.get("discharge"),
|
|
"discharge_percent": point.get("discharge_percent"),
|
|
}
|
|
)
|
|
|
|
return measurements
|
|
|
|
except Exception as e:
|
|
logging.error(f"Error querying InfluxDB: {e}")
|
|
return []
|
|
|
|
|
|
# MySQL/PostgreSQL Adapter
|
|
class SQLAdapter(DatabaseAdapter):
|
|
def __init__(self, connection_string: str, db_type: str = "mysql"):
|
|
self.connection_string = connection_string
|
|
self.db_type = db_type.lower()
|
|
self.engine = None
|
|
|
|
# Add SQLite-specific connection parameters for better concurrency
|
|
if self.db_type == "sqlite":
|
|
if "?" not in connection_string:
|
|
self.connection_string += "?timeout=30&check_same_thread=False"
|
|
else:
|
|
self.connection_string += "&timeout=30&check_same_thread=False"
|
|
|
|
def connect(self):
|
|
try:
|
|
from sqlalchemy import create_engine
|
|
|
|
self.engine = create_engine(self.connection_string, pool_pre_ping=True)
|
|
|
|
# Create tables
|
|
self._create_tables()
|
|
|
|
logging.info(f"Connected to {self.db_type.upper()} successfully")
|
|
return True
|
|
|
|
except ImportError:
|
|
logging.error("SQLAlchemy not installed. Run: pip install sqlalchemy pymysql")
|
|
return False
|
|
except Exception as e:
|
|
logging.error(f"Failed to connect to {self.db_type.upper()}: {e}")
|
|
return False
|
|
|
|
def _create_tables(self):
|
|
from sqlalchemy import text
|
|
|
|
# Stations table - adjust for different databases
|
|
if self.db_type == "sqlite":
|
|
stations_sql = """
|
|
CREATE TABLE IF NOT EXISTS stations (
|
|
id INTEGER PRIMARY KEY,
|
|
station_code TEXT UNIQUE NOT NULL,
|
|
thai_name TEXT NOT NULL,
|
|
english_name TEXT NOT NULL,
|
|
latitude REAL,
|
|
longitude REAL,
|
|
geohash TEXT,
|
|
created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
|
|
updated_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP
|
|
)
|
|
"""
|
|
|
|
measurements_sql = """
|
|
CREATE TABLE IF NOT EXISTS water_measurements (
|
|
id INTEGER PRIMARY KEY AUTOINCREMENT,
|
|
timestamp DATETIME NOT NULL,
|
|
station_id INTEGER NOT NULL,
|
|
water_level REAL,
|
|
discharge REAL,
|
|
discharge_percent REAL,
|
|
status TEXT DEFAULT 'active',
|
|
created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
|
|
FOREIGN KEY (station_id) REFERENCES stations(id),
|
|
UNIQUE(timestamp, station_id)
|
|
)
|
|
"""
|
|
|
|
# Create indexes separately for SQLite
|
|
index_sql = [
|
|
"CREATE INDEX IF NOT EXISTS idx_timestamp ON water_measurements(timestamp)",
|
|
"CREATE INDEX IF NOT EXISTS idx_station_timestamp ON water_measurements(station_id, timestamp)",
|
|
]
|
|
|
|
elif self.db_type == "postgresql":
|
|
stations_sql = """
|
|
CREATE TABLE IF NOT EXISTS stations (
|
|
id SERIAL PRIMARY KEY,
|
|
station_code VARCHAR(10) UNIQUE NOT NULL,
|
|
thai_name VARCHAR(255) NOT NULL,
|
|
english_name VARCHAR(255) NOT NULL,
|
|
latitude DECIMAL(10,8),
|
|
longitude DECIMAL(11,8),
|
|
geohash VARCHAR(20),
|
|
created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
|
|
updated_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP
|
|
)
|
|
"""
|
|
|
|
measurements_sql = """
|
|
CREATE TABLE IF NOT EXISTS water_measurements (
|
|
id BIGSERIAL PRIMARY KEY,
|
|
timestamp TIMESTAMP NOT NULL,
|
|
station_id INTEGER NOT NULL,
|
|
water_level NUMERIC(10,3),
|
|
discharge NUMERIC(10,2),
|
|
discharge_percent NUMERIC(5,2),
|
|
status VARCHAR(20) DEFAULT 'active',
|
|
created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
|
|
FOREIGN KEY (station_id) REFERENCES stations(id),
|
|
UNIQUE(timestamp, station_id)
|
|
)
|
|
"""
|
|
|
|
index_sql = [
|
|
"CREATE INDEX IF NOT EXISTS idx_timestamp ON water_measurements(timestamp)",
|
|
"CREATE INDEX IF NOT EXISTS idx_station_timestamp ON water_measurements(station_id, timestamp DESC)",
|
|
]
|
|
|
|
else: # MySQL
|
|
stations_sql = """
|
|
CREATE TABLE IF NOT EXISTS stations (
|
|
id INT PRIMARY KEY,
|
|
station_code VARCHAR(10) UNIQUE NOT NULL,
|
|
thai_name VARCHAR(255) NOT NULL,
|
|
english_name VARCHAR(255) NOT NULL,
|
|
latitude DECIMAL(10,8),
|
|
longitude DECIMAL(11,8),
|
|
geohash VARCHAR(20),
|
|
created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
|
|
updated_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP ON UPDATE CURRENT_TIMESTAMP
|
|
)
|
|
"""
|
|
|
|
measurements_sql = """
|
|
CREATE TABLE IF NOT EXISTS water_measurements (
|
|
id BIGINT AUTO_INCREMENT PRIMARY KEY,
|
|
timestamp DATETIME NOT NULL,
|
|
station_id INT NOT NULL,
|
|
water_level DECIMAL(10,3),
|
|
discharge DECIMAL(10,2),
|
|
discharge_percent DECIMAL(5,2),
|
|
status VARCHAR(20) DEFAULT 'active',
|
|
created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
|
|
FOREIGN KEY (station_id) REFERENCES stations(id),
|
|
UNIQUE KEY unique_measurement (timestamp, station_id),
|
|
INDEX idx_timestamp (timestamp),
|
|
INDEX idx_station_timestamp (station_id, timestamp)
|
|
)
|
|
"""
|
|
index_sql = []
|
|
|
|
with self.engine.begin() as conn:
|
|
conn.execute(text(stations_sql))
|
|
conn.execute(text(measurements_sql))
|
|
|
|
# Create indexes for SQLite and PostgreSQL
|
|
for index in index_sql:
|
|
conn.execute(text(index))
|
|
|
|
# Transaction is automatically committed when context manager exits
|
|
|
|
def save_measurements(self, measurements: List[Dict]) -> bool:
|
|
if not self.engine:
|
|
return False
|
|
|
|
try:
|
|
from sqlalchemy import text
|
|
|
|
with self.engine.begin() as conn:
|
|
# Insert/update stations
|
|
for measurement in measurements:
|
|
if self.db_type == "sqlite":
|
|
station_sql = """
|
|
INSERT OR REPLACE INTO stations
|
|
(id, station_code, thai_name, english_name,
|
|
latitude, longitude, geohash, updated_at)
|
|
VALUES
|
|
(:station_id, :station_code, :thai_name, :english_name,
|
|
:latitude, :longitude, :geohash, CURRENT_TIMESTAMP)
|
|
"""
|
|
elif self.db_type == "postgresql":
|
|
station_sql = """
|
|
INSERT INTO stations
|
|
(id, station_code, thai_name, english_name,
|
|
latitude, longitude, geohash, updated_at)
|
|
VALUES
|
|
(:station_id, :station_code, :thai_name, :english_name,
|
|
:latitude, :longitude, :geohash, NOW())
|
|
ON CONFLICT (id) DO UPDATE SET
|
|
thai_name = EXCLUDED.thai_name,
|
|
english_name = EXCLUDED.english_name,
|
|
latitude = EXCLUDED.latitude,
|
|
longitude = EXCLUDED.longitude,
|
|
geohash = EXCLUDED.geohash,
|
|
updated_at = NOW()
|
|
"""
|
|
else: # MySQL
|
|
station_sql = """
|
|
INSERT INTO stations
|
|
(id, station_code, thai_name, english_name,
|
|
latitude, longitude, geohash, updated_at)
|
|
VALUES
|
|
(:station_id, :station_code, :thai_name, :english_name,
|
|
:latitude, :longitude, :geohash, NOW())
|
|
ON DUPLICATE KEY UPDATE
|
|
thai_name = VALUES(thai_name),
|
|
english_name = VALUES(english_name),
|
|
latitude = VALUES(latitude),
|
|
longitude = VALUES(longitude),
|
|
geohash = VALUES(geohash),
|
|
updated_at = NOW()
|
|
"""
|
|
|
|
conn.execute(
|
|
text(station_sql),
|
|
{
|
|
"station_id": measurement["station_id"],
|
|
"station_code": measurement["station_code"],
|
|
"thai_name": measurement["station_name_th"],
|
|
"english_name": measurement["station_name_en"],
|
|
"latitude": measurement.get("latitude"),
|
|
"longitude": measurement.get("longitude"),
|
|
"geohash": measurement.get("geohash"),
|
|
},
|
|
)
|
|
|
|
# Insert measurements
|
|
for measurement in measurements:
|
|
if self.db_type == "sqlite":
|
|
measurement_sql = """
|
|
INSERT OR REPLACE INTO water_measurements
|
|
(timestamp, station_id, water_level, discharge, discharge_percent, status)
|
|
VALUES (:timestamp, :station_id, :water_level, :discharge, :discharge_percent, :status)
|
|
"""
|
|
elif self.db_type == "postgresql":
|
|
measurement_sql = """
|
|
INSERT INTO water_measurements
|
|
(timestamp, station_id, water_level, discharge, discharge_percent, status)
|
|
VALUES (:timestamp, :station_id, :water_level, :discharge, :discharge_percent, :status)
|
|
ON CONFLICT (timestamp, station_id) DO UPDATE SET
|
|
water_level = EXCLUDED.water_level,
|
|
discharge = EXCLUDED.discharge,
|
|
discharge_percent = EXCLUDED.discharge_percent,
|
|
status = EXCLUDED.status
|
|
"""
|
|
else: # MySQL
|
|
measurement_sql = """
|
|
INSERT INTO water_measurements
|
|
(timestamp, station_id, water_level, discharge, discharge_percent, status)
|
|
VALUES (:timestamp, :station_id, :water_level, :discharge, :discharge_percent, :status)
|
|
ON DUPLICATE KEY UPDATE
|
|
water_level = VALUES(water_level),
|
|
discharge = VALUES(discharge),
|
|
discharge_percent = VALUES(discharge_percent),
|
|
status = VALUES(status)
|
|
"""
|
|
|
|
conn.execute(
|
|
text(measurement_sql),
|
|
{
|
|
"timestamp": measurement["timestamp"],
|
|
"station_id": measurement["station_id"],
|
|
"water_level": measurement["water_level"],
|
|
"discharge": measurement["discharge"],
|
|
"discharge_percent": measurement["discharge_percent"],
|
|
"status": measurement["status"],
|
|
},
|
|
)
|
|
|
|
# Transaction is automatically committed when context manager exits
|
|
logging.info(f"Successfully saved {len(measurements)} measurements to {self.db_type.upper()}")
|
|
return True
|
|
|
|
except Exception as e:
|
|
logging.error(f"Error saving to {self.db_type.upper()}: {e}")
|
|
return False
|
|
|
|
def get_latest_measurements(self, limit: int = 100) -> List[Dict]:
|
|
if not self.engine:
|
|
return []
|
|
|
|
try:
|
|
from sqlalchemy import text
|
|
|
|
query = """
|
|
SELECT m.timestamp, s.station_code, s.english_name, s.thai_name,
|
|
m.water_level, m.discharge, m.discharge_percent, m.status
|
|
FROM water_measurements m
|
|
JOIN stations s ON m.station_id = s.id
|
|
INNER JOIN (
|
|
SELECT station_id, MAX(timestamp) as max_timestamp
|
|
FROM water_measurements
|
|
GROUP BY station_id
|
|
) latest ON m.station_id = latest.station_id AND m.timestamp = latest.max_timestamp
|
|
ORDER BY s.station_code
|
|
LIMIT :limit
|
|
"""
|
|
|
|
with self.engine.connect() as conn:
|
|
result = conn.execute(text(query), {"limit": limit})
|
|
measurements = []
|
|
|
|
for row in result:
|
|
measurements.append(
|
|
{
|
|
"timestamp": row[0],
|
|
"station_code": row[1],
|
|
"station_name_en": row[2],
|
|
"station_name_th": row[3],
|
|
"water_level": float(row[4]) if row[4] else None,
|
|
"discharge": float(row[5]) if row[5] else None,
|
|
"discharge_percent": float(row[6]) if row[6] else None,
|
|
"status": row[7],
|
|
}
|
|
)
|
|
|
|
return measurements
|
|
|
|
except Exception as e:
|
|
logging.error(f"Error querying {self.db_type.upper()}: {e}")
|
|
return []
|
|
|
|
def get_measurements_by_timerange(
|
|
self,
|
|
start_time: datetime.datetime,
|
|
end_time: datetime.datetime,
|
|
station_codes: Optional[List[str]] = None,
|
|
) -> List[Dict]:
|
|
if not self.engine:
|
|
return []
|
|
|
|
try:
|
|
from sqlalchemy import text
|
|
|
|
where_clause = "m.timestamp BETWEEN :start_time AND :end_time"
|
|
params = {"start_time": start_time, "end_time": end_time}
|
|
|
|
if station_codes:
|
|
placeholders = ",".join([f":station_{i}" for i in range(len(station_codes))])
|
|
where_clause += f" AND s.station_code IN ({placeholders})"
|
|
for i, code in enumerate(station_codes):
|
|
params[f"station_{i}"] = code
|
|
|
|
query = f"""
|
|
SELECT m.timestamp, s.station_code, s.english_name, s.thai_name,
|
|
m.water_level, m.discharge, m.discharge_percent, m.status
|
|
FROM water_measurements m
|
|
JOIN stations s ON m.station_id = s.id
|
|
WHERE {where_clause}
|
|
ORDER BY m.timestamp DESC, s.station_code
|
|
"""
|
|
|
|
with self.engine.connect() as conn:
|
|
result = conn.execute(text(query), params)
|
|
measurements = []
|
|
|
|
for row in result:
|
|
measurements.append(
|
|
{
|
|
"timestamp": row[0],
|
|
"station_code": row[1],
|
|
"station_name_en": row[2],
|
|
"station_name_th": row[3],
|
|
"water_level": float(row[4]) if row[4] else None,
|
|
"discharge": float(row[5]) if row[5] else None,
|
|
"discharge_percent": float(row[6]) if row[6] else None,
|
|
"status": row[7],
|
|
}
|
|
)
|
|
|
|
return measurements
|
|
|
|
except Exception as e:
|
|
logging.error(f"Error querying {self.db_type.upper()}: {e}")
|
|
return []
|
|
|
|
def get_measurements_for_date(self, target_date: datetime.datetime) -> List[Dict]:
|
|
"""Get all measurements for a specific date"""
|
|
if not self.engine:
|
|
return []
|
|
|
|
try:
|
|
from sqlalchemy import text
|
|
|
|
# Get start and end of the target date
|
|
start_of_day = target_date.replace(hour=0, minute=0, second=0, microsecond=0)
|
|
end_of_day = target_date.replace(hour=23, minute=59, second=59, microsecond=999999)
|
|
|
|
query = """
|
|
SELECT m.timestamp, m.station_id, s.station_code, s.thai_name,
|
|
m.water_level, m.discharge, m.discharge_percent, m.status
|
|
FROM water_measurements m
|
|
LEFT JOIN stations s ON m.station_id = s.id
|
|
WHERE m.timestamp >= :start_time AND m.timestamp <= :end_time
|
|
ORDER BY m.timestamp DESC
|
|
"""
|
|
|
|
with self.engine.connect() as conn:
|
|
result = conn.execute(text(query), {"start_time": start_of_day, "end_time": end_of_day})
|
|
|
|
measurements = []
|
|
for row in result:
|
|
measurements.append(
|
|
{
|
|
"timestamp": row[0],
|
|
"station_id": row[1],
|
|
"station_code": row[2] or f"Station_{row[1]}",
|
|
"station_name_th": row[3] or f"Station {row[1]}",
|
|
"water_level": float(row[4]) if row[4] else None,
|
|
"discharge": float(row[5]) if row[5] else None,
|
|
"discharge_percent": float(row[6]) if row[6] else None,
|
|
"status": row[7],
|
|
}
|
|
)
|
|
|
|
return measurements
|
|
|
|
except Exception as e:
|
|
logging.error(f"Error querying {self.db_type.upper()} for date {target_date.date()}: {e}")
|
|
return []
|
|
|
|
|
|
# VictoriaMetrics Adapter (using Prometheus format)
|
|
class VictoriaMetricsAdapter(DatabaseAdapter):
|
|
def __init__(self, host: str = "localhost", port: int = 8428):
|
|
self.host = host
|
|
self.port = port
|
|
|
|
# Handle HTTPS URLs and reverse proxy configurations
|
|
if host.startswith(("http://", "https://")):
|
|
self.base_url = host
|
|
if port != 80 and port != 443 and not host.endswith(f":{port}"):
|
|
# Only add port if it's not standard and not already in URL
|
|
if "://" in host and ":" not in host.split("://")[1]:
|
|
self.base_url = f"{host}:{port}"
|
|
else:
|
|
# Default to HTTP for localhost, HTTPS for remote hosts
|
|
protocol = "https" if host != "localhost" and not host.startswith("127.") else "http"
|
|
if (protocol == "https" and port == 443) or (protocol == "http" and port == 80):
|
|
self.base_url = f"{protocol}://{host}"
|
|
else:
|
|
self.base_url = f"{protocol}://{host}:{port}"
|
|
|
|
@staticmethod
|
|
def _escape_label(value) -> str:
|
|
"""Escape a Prometheus label value per the exposition format spec.
|
|
|
|
Station names include arbitrary Thai text (and could be set via the API),
|
|
so backslashes, double-quotes and newlines must be escaped to avoid
|
|
producing malformed or injected exposition lines.
|
|
"""
|
|
return str(value).replace("\\", "\\\\").replace('"', '\\"').replace("\n", "\\n")
|
|
|
|
@staticmethod
|
|
def _metric_value(value) -> Optional[float]:
|
|
"""Coerce a numeric field to float, or None if it isn't a valid number."""
|
|
try:
|
|
return float(value)
|
|
except (TypeError, ValueError):
|
|
return None
|
|
|
|
def connect(self):
|
|
try:
|
|
import requests
|
|
|
|
# Test connection with SSL verification and timeout
|
|
response = requests.get(
|
|
f"{self.base_url}/api/v1/status/config",
|
|
timeout=10,
|
|
verify=True, # Enable SSL verification for HTTPS
|
|
)
|
|
if response.status_code == 200:
|
|
logging.info(f"Connected to VictoriaMetrics successfully at {self.base_url}")
|
|
return True
|
|
else:
|
|
logging.error(f"VictoriaMetrics connection failed: {response.status_code}")
|
|
return False
|
|
except requests.exceptions.SSLError as e:
|
|
logging.error(f"SSL error connecting to VictoriaMetrics: {e}")
|
|
return False
|
|
except requests.exceptions.ConnectionError as e:
|
|
logging.error(f"Connection error to VictoriaMetrics: {e}")
|
|
return False
|
|
except Exception as e:
|
|
logging.error(f"Failed to connect to VictoriaMetrics: {e}")
|
|
return False
|
|
|
|
def save_measurements(self, measurements: List[Dict]) -> bool:
|
|
try:
|
|
import requests
|
|
|
|
# Convert to Prometheus format
|
|
metrics_data = []
|
|
timestamp_ms = int(datetime.datetime.now().timestamp() * 1000)
|
|
|
|
for measurement in measurements:
|
|
# Escape label values once per measurement (untrusted Thai/English names).
|
|
labels = (
|
|
f'station_code="{self._escape_label(measurement["station_code"])}",'
|
|
f'station_name_en="{self._escape_label(measurement["station_name_en"])}",'
|
|
f'station_name_th="{self._escape_label(measurement["station_name_th"])}"'
|
|
)
|
|
|
|
# Water level metric
|
|
water_level = self._metric_value(measurement.get("water_level"))
|
|
if water_level is not None:
|
|
metrics_data.append(f"water_level{{{labels}}} {water_level} {timestamp_ms}")
|
|
|
|
# Discharge metric
|
|
discharge = self._metric_value(measurement.get("discharge"))
|
|
if discharge is not None:
|
|
metrics_data.append(f"water_discharge{{{labels}}} {discharge} {timestamp_ms}")
|
|
|
|
# Discharge percentage metric
|
|
discharge_percent = self._metric_value(measurement.get("discharge_percent"))
|
|
if discharge_percent is not None:
|
|
metrics_data.append(f"water_discharge_percent{{{labels}}} {discharge_percent} {timestamp_ms}")
|
|
|
|
# Send to VictoriaMetrics
|
|
data = "\n".join(metrics_data)
|
|
response = requests.post(
|
|
f"{self.base_url}/api/v1/import/prometheus",
|
|
data=data,
|
|
headers={"Content-Type": "text/plain"},
|
|
timeout=30,
|
|
)
|
|
|
|
if response.status_code == 204:
|
|
logging.info(f"Successfully sent {len(measurements)} measurements to VictoriaMetrics")
|
|
return True
|
|
else:
|
|
logging.error(f"VictoriaMetrics import failed: {response.status_code} - {response.text}")
|
|
return False
|
|
|
|
except Exception as e:
|
|
logging.error(f"Error sending to VictoriaMetrics: {e}")
|
|
return False
|
|
|
|
def get_latest_measurements(self, limit: int = 100) -> List[Dict]:
|
|
# VictoriaMetrics queries would be implemented here
|
|
# This is a simplified version
|
|
logging.warning("get_latest_measurements not fully implemented for VictoriaMetrics")
|
|
return []
|
|
|
|
def get_measurements_by_timerange(
|
|
self,
|
|
start_time: datetime.datetime,
|
|
end_time: datetime.datetime,
|
|
station_codes: Optional[List[str]] = None,
|
|
) -> List[Dict]:
|
|
# VictoriaMetrics range queries would be implemented here
|
|
logging.warning("get_measurements_by_timerange not fully implemented for VictoriaMetrics")
|
|
return []
|
|
|
|
def get_measurements_for_date(self, target_date: datetime.datetime) -> List[Dict]:
|
|
"""Get all measurements for a specific date"""
|
|
logging.warning("get_measurements_for_date not fully implemented for VictoriaMetrics")
|
|
return []
|
|
|
|
|
|
# Factory function to create appropriate adapter
|
|
def create_database_adapter(db_type: str, **kwargs) -> DatabaseAdapter:
|
|
"""
|
|
Factory function to create database adapter
|
|
|
|
Args:
|
|
db_type: 'influxdb', 'mysql', 'postgresql', 'sqlite', or 'victoriametrics'
|
|
**kwargs: Database-specific connection parameters
|
|
"""
|
|
db_type = db_type.lower()
|
|
|
|
if db_type == "influxdb":
|
|
return InfluxDBAdapter(**kwargs)
|
|
elif db_type == "mysql":
|
|
return SQLAdapter(db_type="mysql", **kwargs)
|
|
elif db_type == "postgresql":
|
|
return SQLAdapter(db_type="postgresql", **kwargs)
|
|
elif db_type == "sqlite":
|
|
return SQLAdapter(db_type="sqlite", **kwargs)
|
|
elif db_type == "victoriametrics":
|
|
return VictoriaMetricsAdapter(**kwargs)
|
|
else:
|
|
raise ValueError(f"Unsupported database type: {db_type}")
|