tracer-cloud--opensre/integrations/coralogix/client.py

"""Coralogix DataPrime client for RCA log queries."""

from __future__ import annotations

import json
import logging
from datetime import UTC, datetime, timedelta
from typing import Any

import httpx

from integrations.config_models import CoralogixIntegrationConfig
from integrations.probes import ProbeResult
from platform.observability.errors.service import capture_service_error
from platform.observability.streaming import StreamingParseStats

logger = logging.getLogger(__name__)

_DEFAULT_TIMEOUT_SECONDS = 30.0


def _escape_query_value(value: str) -> str:
    return value.replace("\\", "\\\\").replace("'", "\\'")


def _ensure_limit_clause(query: str, limit: int) -> str:
    normalized = query.strip()
    if " limit " in normalized.lower() or normalized.lower().endswith("| limit"):
        return normalized
    return f"{normalized} | limit {max(int(limit), 1)}"


def build_coralogix_logs_query(
    *,
    raw_query: str = "",
    application_name: str = "",
    subsystem_name: str = "",
    text_query: str = "",
    trace_id: str = "",
    limit: int = 50,
) -> str:
    """Build a Coralogix DataPrime log query with optional RCA-friendly filters."""
    if raw_query.strip():
        return _ensure_limit_clause(raw_query.strip(), limit)

    filters: list[str] = []
    if application_name:
        filters.append(f"$l.applicationname == '{_escape_query_value(application_name)}'")
    if subsystem_name:
        filters.append(f"$l.subsystemname == '{_escape_query_value(subsystem_name)}'")
    if trace_id:
        filters.append(f"$d.trace_id == '{_escape_query_value(trace_id)}'")
    if text_query:
        filters.append(f"$d.message.contains('{_escape_query_value(text_query)}')")

    query_parts = ["source logs"]
    query_parts.extend(f"filter {clause}" for clause in filters)
    query_parts.append(f"limit {max(int(limit), 1)}")
    return " | ".join(query_parts)


class CoralogixClient:
    """Synchronous Coralogix client backed by the DataPrime HTTP API."""

    def __init__(self, config: CoralogixIntegrationConfig) -> None:
        self.config = config

    @property
    def is_configured(self) -> bool:
        return bool(self.config.api_key and self.config.base_url)

    @property
    def query_url(self) -> str:
        return f"{self.config.base_url}/api/v1/dataprime/query"

    def _request_headers(self) -> dict[str, str]:
        return {
            "Authorization": f"Bearer {self.config.api_key}",
            "Content-Type": "application/json",
            "Accept": "application/json",
        }

    def _parse_ndjson(
        self, response_text: str, stats: StreamingParseStats | None = None
    ) -> dict[str, Any]:
        query_ids: list[str] = []
        warnings: list[str] = []
        rows: list[dict[str, Any]] = []

        for line in response_text.splitlines():
            stripped = line.strip()
            if not stripped:
                continue
            try:
                payload = json.loads(stripped)
            except json.JSONDecodeError as exc:
                if stats is not None:
                    stats.record_error(exc)
                continue
            if stats is not None:
                stats.record_parsed()

            if isinstance(payload, dict):
                query_id = payload.get("queryId")
                if isinstance(query_id, dict):
                    maybe_id = str(query_id.get("queryId", "")).strip()
                    if maybe_id:
                        query_ids.append(maybe_id)
                elif isinstance(query_id, str) and query_id.strip():
                    query_ids.append(query_id.strip())

                warning = payload.get("warning")
                if isinstance(warning, str) and warning.strip():
                    warnings.append(warning.strip())

                direct_results = payload.get("result", {})
                if isinstance(direct_results, dict):
                    result_rows = direct_results.get("results", [])
                    if isinstance(result_rows, list):
                        rows.extend(item for item in result_rows if isinstance(item, dict))

                background_results = payload.get("response", {})
                if isinstance(background_results, dict):
                    nested_results = background_results.get("results", {})
                    if isinstance(nested_results, dict):
                        result_rows = nested_results.get("results", [])
                        if isinstance(result_rows, list):
                            rows.extend(item for item in result_rows if isinstance(item, dict))

        return {
            "query_ids": query_ids,
            "warnings": warnings,
            "rows": rows,
        }

    @staticmethod
    def _items_to_dict(items: object) -> dict[str, Any]:
        if not isinstance(items, list):
            return {}
        result: dict[str, Any] = {}
        for item in items:
            if not isinstance(item, dict):
                continue
            key = str(item.get("key", "")).strip()
            if not key:
                continue
            result[key] = item.get("value")
        return result

    @staticmethod
    def _parse_user_data(value: object, stats: StreamingParseStats | None = None) -> dict[str, Any]:
        if isinstance(value, dict):
            return value
        if not isinstance(value, str) or not value.strip():
            return {}
        # Only the string branch is an actual parse attempt, so both
        # record_parsed and record_error live inside it — otherwise the
        # denominator (parsed + skipped) would only count failures and skew
        # the skip ratio.
        try:
            payload = json.loads(value)
        except json.JSONDecodeError as exc:
            if stats is not None:
                stats.record_error(exc)
            return {}
        if stats is not None:
            stats.record_parsed()
        return payload if isinstance(payload, dict) else {}

    def _normalize_row(
        self, row: dict[str, Any], stats: StreamingParseStats | None = None
    ) -> dict[str, Any]:
        metadata = self._items_to_dict(row.get("metadata"))
        labels = self._items_to_dict(row.get("labels"))
        user_data = self._parse_user_data(row.get("userData"), stats=stats)
        log_obj = user_data.get("log_obj", {}) if isinstance(user_data, dict) else {}
        if not isinstance(log_obj, dict):
            log_obj = {}

        message = str(
            log_obj.get("message") or user_data.get("message") or metadata.get("message") or ""
        ).strip()
        timestamp = str(
            log_obj.get("timestamp")
            or user_data.get("timestamp")
            or metadata.get("timestamp")
            or ""
        ).strip()
        level = str(
            log_obj.get("level") or user_data.get("level") or metadata.get("severity") or ""
        ).strip()
        trace_id = str(
            user_data.get("trace_id") or log_obj.get("trace_id") or metadata.get("trace_id") or ""
        ).strip()

        return {
            "timestamp": timestamp,
            "message": message,
            "level": level,
            "trace_id": trace_id,
            "application_name": str(labels.get("applicationname", "")).strip(),
            "subsystem_name": str(labels.get("subsystemname", "")).strip(),
            "labels": labels,
            "metadata": metadata,
            "user_data": user_data,
        }

    def query_logs(
        self,
        query: str,
        *,
        time_range_minutes: int = 60,
        limit: int = 50,
    ) -> dict[str, Any]:
        """Run a Coralogix DataPrime direct query and normalize log rows."""
        query_with_limit = _ensure_limit_clause(query, limit)
        now = datetime.now(UTC)
        start = now - timedelta(minutes=max(int(time_range_minutes), 1))
        payload = {
            "query": query_with_limit,
            "metadata": {
                "startDate": start.isoformat(),
                "endDate": now.isoformat(),
                "defaultSource": "logs",
            },
        }

        try:
            response = httpx.post(
                self.query_url,
                headers=self._request_headers(),
                json=payload,
                timeout=_DEFAULT_TIMEOUT_SECONDS,
            )
            response.raise_for_status()
        except httpx.HTTPStatusError as exc:
            capture_service_error(exc, logger=logger, integration="coralogix", method="query_logs")
            return {
                "success": False,
                "error": f"HTTP {exc.response.status_code}: {exc.response.text[:200]}",
            }
        except Exception as exc:
            capture_service_error(exc, logger=logger, integration="coralogix", method="query_logs")
            return {"success": False, "error": str(exc)}

        stats = StreamingParseStats()
        parsed = self._parse_ndjson(response.text, stats=stats)
        # Aggregate userData parse drops into the same stats so the threshold
        # fires once per response no matter where the drops happened.
        logs = [self._normalize_row(row, stats=stats) for row in parsed["rows"]]
        stats.report_if_unhealthy(logger=logger, integration="coralogix", source="dataprime/query")
        return {
            "success": True,
            "logs": logs,
            "total": len(logs),
            "query": query_with_limit,
            "query_ids": parsed["query_ids"],
            "warnings": parsed["warnings"],
        }

    def validate_access(self) -> dict[str, Any]:
        """Validate Coralogix access with a lightweight direct query."""
        result = self.query_logs("source logs | limit 1", time_range_minutes=15, limit=1)
        if not result.get("success"):
            return result
        return {
            "success": True,
            "total": result.get("total", 0),
            "warnings": result.get("warnings", []),
        }

    def probe_access(self) -> ProbeResult:
        """Validate Coralogix access using the lightweight direct query probe."""
        if not self.is_configured:
            return ProbeResult.missing("Missing Coralogix API key or API URL.")

        result = self.validate_access()
        if not result.get("success"):
            return ProbeResult.failed(
                f"DataPrime check failed: {result.get('error', 'unknown error')}"
            )

        scope: list[str] = []
        if self.config.application_name:
            scope.append(f"application {self.config.application_name}")
        if self.config.subsystem_name:
            scope.append(f"subsystem {self.config.subsystem_name}")
        scope_detail = f" ({', '.join(scope)})" if scope else ""
        total = int(result.get("total", 0) or 0)
        return ProbeResult.passed(
            (
                f"Connected to {self.config.base_url}{scope_detail}; "
                f"DataPrime returned {total} row(s)."
            ),
            total=total,
        )