"""InfluxDB 2.x API client."""
import csv
import io
import logging
from typing import Any
from overwatch_mcp.clients.base import BaseHTTPClient
from overwatch_mcp.models.config import InfluxDBConfig
from overwatch_mcp.models.errors import ErrorCode, OverwatchError
logger = logging.getLogger(__name__)
class InfluxDBClient(BaseHTTPClient):
"""Client for InfluxDB 2.x API."""
def __init__(self, config: InfluxDBConfig):
"""
Initialize InfluxDB client.
Args:
config: InfluxDB configuration
"""
# Normalize URL - strip /api/v2 suffix if present (paths already include it)
url = config.url.rstrip("/")
if url.endswith("/api/v2"):
url = url[:-7]
elif url.endswith("/api"):
url = url[:-4]
headers = {
"Authorization": f"Token {config.token}",
"Accept": "application/csv",
"Content-Type": "application/vnd.flux",
}
super().__init__(
base_url=url,
timeout_seconds=config.timeout_seconds,
headers=headers,
verify_ssl=config.verify_ssl,
)
self.config = config
def _validate_bucket(self, bucket: str) -> None:
"""
Validate bucket against allowlist.
Args:
bucket: Bucket name to validate
Raises:
OverwatchError: If bucket is not in allowlist
"""
if bucket not in self.config.allowed_buckets:
raise OverwatchError(
code=ErrorCode.BUCKET_NOT_ALLOWED,
message=f"Bucket '{bucket}' is not in allowed list",
details={
"requested_bucket": bucket,
"allowed_buckets": self.config.allowed_buckets,
},
)
def _validate_query_bucket(self, query: str, bucket: str) -> None:
"""
Validate that query references the specified bucket.
Args:
query: Flux query
bucket: Expected bucket name
Raises:
OverwatchError: If query doesn't reference the correct bucket
"""
# Simple validation - check if 'from(bucket: "name")' exists
if f'from(bucket: "{bucket}")' not in query and f"from(bucket:\"{bucket}\")" not in query:
raise OverwatchError(
code=ErrorCode.INVALID_QUERY,
message=f"Query must reference bucket '{bucket}'",
details={"bucket": bucket, "query": query[:200]},
)
def _parse_csv_response(self, csv_text: str) -> list[dict[str, Any]]:
"""
Parse InfluxDB annotated CSV response to list of records.
InfluxDB returns CSV with annotation rows (starting with #) followed by data.
Args:
csv_text: CSV response text
Returns:
List of record dictionaries
Raises:
OverwatchError: If CSV parsing fails
"""
try:
lines = csv_text.strip().split("\n")
if not lines:
return []
# Find where annotation rows end and data begins
# Annotations start with #, then comes header row, then data rows
has_annotations = any(line.startswith("#") for line in lines)
non_annotation_lines = []
for line in lines:
if not line.startswith("#"):
non_annotation_lines.append(line)
# If we have no annotations and multiple lines, this isn't InfluxDB format
if not has_annotations and len(lines) > 1:
raise ValueError("Invalid InfluxDB CSV format: missing annotation rows")
if len(non_annotation_lines) < 1: # Need at least header
return []
# First non-annotation line is the header
header_line = non_annotation_lines[0]
data_lines = non_annotation_lines[1:] if len(non_annotation_lines) > 1 else []
# If there are no data lines (only empty lines), return empty
if not any(line.strip() for line in data_lines):
return []
# Parse CSV
reader = csv.DictReader(
io.StringIO("\n".join([header_line] + data_lines))
)
return list(reader)
except Exception as e:
raise OverwatchError(
code=ErrorCode.UPSTREAM_CLIENT_ERROR,
message=f"Failed to parse InfluxDB CSV response: {str(e)}",
details={"error": str(e), "csv_preview": csv_text[:500]},
)
async def query(self, query: str, bucket: str) -> list[dict[str, Any]]:
"""
Execute Flux query against InfluxDB.
Args:
query: Flux query
bucket: Target bucket
Returns:
List of records
Raises:
OverwatchError: On API errors, bucket not allowed, or parsing failures
"""
# Validate bucket
self._validate_bucket(bucket)
self._validate_query_bucket(query, bucket)
# Add org parameter
params = {"org": self.config.org}
# Make request with raw Flux query body
client = await self._get_client()
try:
response = await client.post(
"/api/v2/query",
params=params,
content=query,
headers={
**self.default_headers,
"Content-Type": "application/vnd.flux",
},
)
if response.status_code != 200:
raise OverwatchError(
code=ErrorCode.UPSTREAM_CLIENT_ERROR,
message=f"InfluxDB query failed: {response.status_code}",
details={
"status_code": response.status_code,
"response": response.text,
},
)
# Parse CSV response
return self._parse_csv_response(response.text)
except OverwatchError:
raise
except Exception as e:
raise OverwatchError(
code=ErrorCode.UPSTREAM_SERVER_ERROR,
message=f"InfluxDB request failed: {str(e)}",
details={"error": str(e)},
)
async def health_check(self) -> bool:
"""
Check if InfluxDB is reachable.
Returns:
True if healthy, False otherwise
"""
endpoint = "/health"
full_url = f"{self.base_url}{endpoint}"
logger.debug(f"InfluxDB health check: {full_url}")
try:
response = await self.get(endpoint)
logger.debug(f"InfluxDB health response: {response.status_code} - {response.text[:200]}")
return response.status_code == 200
except OverwatchError as e:
logger.debug(f"InfluxDB health check failed: {e}")
return False
