tracer-cloud--opensre/integrations/betterstack/__init__.py

"""Shared Better Stack Telemetry integration helpers.

Provides configuration, connectivity validation, and read-only log queries
against Better Stack's ClickHouse SQL over HTTP endpoint (Basic auth). Recent
and historical log rows are UNIONed across ``remote(<source>_logs)`` and
``s3Cluster(primary, <source>_s3) WHERE _row_type = 1`` per Better Stack's
own sample query, so a single call returns both live-stream and archived data.

Credentials are generated in the dashboard under ``Integrations → Connect
ClickHouse HTTP client`` (username + password + a region-specific endpoint
like ``https://eu-nbg-2-connect.betterstackdata.com``). The SQL endpoint does
not expose a documented source-listing query, so the optional
``BETTERSTACK_SOURCES`` env var surfaces user-supplied source IDs to the
planner via ``extract_params``.
"""

from __future__ import annotations

import json
import logging
import os
import re
from dataclasses import dataclass
from datetime import datetime
from typing import Any

import httpx
from pydantic import Field, field_validator

from config.strict_config import StrictConfigModel
from integrations._validation_helpers import report_classify_failure, report_validation_failure

logger = logging.getLogger(__name__)

DEFAULT_BETTERSTACK_TIMEOUT_S = 15
DEFAULT_BETTERSTACK_MAX_ROWS = 500
_MAX_ALLOWED_ROWS = 10_000
_REQUIRED_CONTENT_TYPE = "text/plain"
_REQUIRED_QUERY_PARAMS = {"output_format_pretty_row_numbers": "0"}
_VALIDATION_PROBE_SQL = "SELECT 1 FORMAT JSONEachRow"
# Source identifiers are ClickHouse bare identifiers (``t{team_id}_{source}``).
# Anything else would land in the FROM clause as SQL injection, so reject
# anything not matching this whitelist.
_SOURCE_NAME_RE = re.compile(r"^[A-Za-z0-9_]+$")


class BetterStackConfig(StrictConfigModel):
    """Normalized Better Stack SQL Query API connection settings."""

    query_endpoint: str = ""
    username: str = ""
    password: str = ""
    sources: list[str] = Field(default_factory=list)
    timeout_seconds: int = Field(default=DEFAULT_BETTERSTACK_TIMEOUT_S, gt=0)
    max_rows: int = Field(default=DEFAULT_BETTERSTACK_MAX_ROWS, gt=0, le=_MAX_ALLOWED_ROWS)
    integration_id: str = ""

    @field_validator("query_endpoint", mode="before")
    @classmethod
    def _normalize_query_endpoint(cls, value: Any) -> str:
        return str(value or "").strip().rstrip("/")

    @field_validator("username", mode="before")
    @classmethod
    def _normalize_username(cls, value: Any) -> str:
        return str(value or "").strip()

    @field_validator("password", mode="before")
    @classmethod
    def _normalize_password(cls, value: Any) -> str:
        # ``StrictConfigModel`` already strips strings as a wildcard validator;
        # this step only coerces ``None`` / non-string inputs into ``""``.
        return str(value or "")

    @field_validator("sources", mode="before")
    @classmethod
    def _normalize_sources(cls, value: Any) -> list[str]:
        if value in (None, ""):
            return []
        if isinstance(value, str):
            parts = [part.strip() for part in value.split(",")]
            return [p for p in parts if p]
        if isinstance(value, list):
            return [str(v).strip() for v in value if str(v).strip()]
        return []

    @property
    def is_configured(self) -> bool:
        return bool(self.query_endpoint and self.username)


def build_betterstack_config(raw: dict[str, Any] | None) -> BetterStackConfig:
    """Build a normalized Better Stack config object from env/store data."""
    return BetterStackConfig.model_validate(raw or {})


def betterstack_config_from_env() -> BetterStackConfig | None:
    """Load a Better Stack config from ``BETTERSTACK_*`` env vars.

    Returns ``None`` when either the endpoint or username is missing. The
    ``BETTERSTACK_SOURCES`` env var is an optional comma-separated hint list
    surfaced to the planner via :func:`betterstack_extract_params`; it is
    not required for availability.
    """
    endpoint = os.getenv("BETTERSTACK_QUERY_ENDPOINT", "").strip()
    username = os.getenv("BETTERSTACK_USERNAME", "").strip()
    if not endpoint or not username:
        return None
    return build_betterstack_config(
        {
            "query_endpoint": endpoint,
            "username": username,
            "password": os.getenv("BETTERSTACK_PASSWORD", ""),
            "sources": os.getenv("BETTERSTACK_SOURCES", ""),
        }
    )


def betterstack_is_available(sources: dict[str, dict]) -> bool:
    """Check if Better Stack credentials are present AND a source is derivable.

    The investigation executor invokes tools purely via
    ``action.run(**extract_params(...))``; there is no other path to inject a
    source identifier at call time. A betterstack integration with credentials
    but no way to derive ``source`` would always run with ``source=""`` and
    deterministically fail. So availability requires, beyond credentials,
    either (a) a configured ``sources`` hint list or (b) an alert-derived
    ``source_hint`` surfaced in the resolved integration config.
    """
    bs = sources.get("betterstack", {})
    if not (bs.get("query_endpoint") and bs.get("username")):
        return False
    has_sources = bool(bs.get("sources"))
    has_hint = bool(str(bs.get("source_hint") or "").strip())
    return has_sources or has_hint


def betterstack_extract_params(sources: dict[str, dict]) -> dict[str, Any]:
    """Extract Better Stack credentials and optional source hints for tool calls.

    Returns both the full ``sources`` hint list and a scalar ``source`` (derived
    from the resolved integration config when present). The executor invokes
    tools purely via ``action.run(**action.extract_params(...))``; we therefore
    have to surface the alert-derived target as a concrete kwarg here.
    """
    bs = sources.get("betterstack", {})
    source_hint = str(bs.get("source_hint", "") or "").strip()
    return {
        "query_endpoint": bs.get("query_endpoint", ""),
        "username": bs.get("username", ""),
        "password": bs.get("password", ""),
        "sources": list(bs.get("sources", []) or []),
        "source": source_hint,
    }


@dataclass(frozen=True)
class BetterStackValidationResult:
    """Outcome of validating a Better Stack integration against the SQL endpoint."""

    ok: bool
    detail: str


def _sql_client(config: BetterStackConfig) -> httpx.Client:
    """Build an authenticated ``httpx.Client`` scoped to the Better Stack SQL API."""
    return httpx.Client(
        auth=(config.username, config.password),
        timeout=float(config.timeout_seconds),
    )


def _post_sql(
    client: httpx.Client,
    endpoint: str,
    query: str,
) -> tuple[httpx.Response | None, str | None]:
    """POST a SQL statement to the Better Stack query endpoint.

    Returns ``(response, None)`` on a transport-level success (any HTTP status),
    or ``(None, error_message)`` on a transport-level failure (DNS, TLS,
    timeout, etc.). Callers are responsible for interpreting non-2xx status
    codes since the error phrasing depends on the operation (probe vs query).
    """
    try:
        response = client.post(
            endpoint,
            params=_REQUIRED_QUERY_PARAMS,
            content=query.encode("utf-8"),
            headers={"Content-Type": _REQUIRED_CONTENT_TYPE},
        )
    except httpx.RequestError as err:
        return None, f"Better Stack request failed: {err}"
    return response, None


def validate_betterstack_config(
    config: BetterStackConfig,
) -> BetterStackValidationResult:
    """Validate Better Stack reachability with a cheap ``SELECT 1`` probe."""
    if not config.is_configured:
        return BetterStackValidationResult(
            ok=False,
            detail="Better Stack query_endpoint and username are required.",
        )

    try:
        with _sql_client(config) as client:
            response, err = _post_sql(client, config.query_endpoint, _VALIDATION_PROBE_SQL)
    except Exception as err:
        report_validation_failure(
            err,
            logger=logger,
            integration="betterstack",
            method="validate_betterstack_config",
        )
        return BetterStackValidationResult(
            ok=False, detail=f"Better Stack connection failed: {err}"
        )

    if err is not None:
        return BetterStackValidationResult(ok=False, detail=err)
    assert response is not None

    status = response.status_code
    if status == 200:
        body = response.text.strip()
        if not body:
            return BetterStackValidationResult(
                ok=False,
                detail="Better Stack SQL endpoint returned an empty body for the probe.",
            )
        return BetterStackValidationResult(
            ok=True,
            detail=f"Connected to Better Stack SQL API at {config.query_endpoint}.",
        )
    if status == 401:
        return BetterStackValidationResult(
            ok=False,
            detail="Better Stack authentication failed (check BETTERSTACK_USERNAME / BETTERSTACK_PASSWORD).",
        )
    if status == 404:
        return BetterStackValidationResult(
            ok=False,
            detail=(
                "Better Stack endpoint not found — verify BETTERSTACK_QUERY_ENDPOINT "
                "matches your region (e.g. https://eu-nbg-2-connect.betterstackdata.com)."
            ),
        )
    return BetterStackValidationResult(
        ok=False,
        detail=f"Better Stack API returned HTTP {status}: {response.text[:200]}",
    )


# ---------------------------------------------------------------------------
# Log-query functions
# ---------------------------------------------------------------------------


def _error_evidence(error: str, *, source: str = "") -> dict[str, Any]:
    """Standard error-shape dict returned by query functions on failure."""
    return {
        "source": "betterstack",
        "available": False,
        "error": error,
        "betterstack_source": source,
        "rows": [],
        "row_count": 0,
    }


def _validate_source_name(source: str) -> str | None:
    """Return the source identifier if it is a safe bare identifier, else ``None``."""
    cleaned = (source or "").strip()
    if not cleaned or not _SOURCE_NAME_RE.fullmatch(cleaned):
        return None
    return cleaned


def _validate_iso_timestamp(value: str | None) -> str | None:
    """Pass-through the ISO-8601 timestamp when parseable; else ``None``.

    Used to reject injected SQL fragments before inlining into the WHERE clause.
    """
    if not value:
        return None
    try:
        datetime.fromisoformat(value.replace("Z", "+00:00"))
    except (TypeError, ValueError):
        return None
    return value


def _time_predicate(since: str | None, until: str | None) -> str:
    """Build the shared ``dt >= … AND dt <= …`` WHERE fragment (or empty)."""
    parts: list[str] = []
    if since:
        parts.append(f"dt >= parseDateTime64BestEffort('{since}', 3, 'UTC')")
    if until:
        parts.append(f"dt <= parseDateTime64BestEffort('{until}', 3, 'UTC')")
    return " AND ".join(parts)


def _build_logs_query(
    source: str,
    since: str | None,
    until: str | None,
    limit: int,
) -> str:
    """Build a UNION query over ``remote(<source>_logs)`` and ``s3Cluster(primary, <source>_s3)``.

    Matches the sample query Better Stack's dashboard generates on connection
    creation: recent logs come from the ``remote(...)`` table function; historical
    (archived) log rows live in S3 and are filtered by ``_row_type = 1`` to
    exclude spans and metrics that share the same ``_s3`` collection.
    """
    time_pred = _time_predicate(since, until)
    recent_where = f"\n  WHERE {time_pred}" if time_pred else ""
    hist_conds = ["_row_type = 1"]
    if time_pred:
        hist_conds.append(time_pred)
    hist_where = f"\n  WHERE {' AND '.join(hist_conds)}"
    return (
        "SELECT dt, raw FROM (\n"
        f"  SELECT dt, raw FROM remote({source}_logs){recent_where}\n"
        "  UNION ALL\n"
        f"  SELECT dt, raw FROM s3Cluster(primary, {source}_s3){hist_where}\n"
        ")\n"
        "ORDER BY dt DESC\n"
        f"LIMIT {limit}\n"
        "FORMAT JSONEachRow"
    )


def _parse_json_each_row(body: str) -> list[dict[str, Any]]:
    """Parse a ClickHouse ``FORMAT JSONEachRow`` body into a list of dicts."""
    rows: list[dict[str, Any]] = []
    for line in body.splitlines():
        stripped = line.strip()
        if not stripped:
            continue
        try:
            parsed = json.loads(stripped)
        except json.JSONDecodeError:
            continue
        if isinstance(parsed, dict):
            rows.append(parsed)
    return rows


def query_logs(
    config: BetterStackConfig,
    source: str,
    since: str | None = None,
    until: str | None = None,
    limit: int | None = None,
) -> dict[str, Any]:
    """Fetch log rows for a Better Stack source (recent + historical).

    Parameters:
        config: authenticated ``BetterStackConfig``.
        source: Better Stack source identifier (e.g. ``t123456_myapp_logs``);
            must match ``[A-Za-z0-9_]+``. Do NOT include the ``_logs`` or
            ``_s3`` suffix — this function appends them to build the
            ``remote(...)`` and ``s3Cluster(primary, ...)`` table functions.
        since: optional ISO-8601 lower-bound timestamp.
        until: optional ISO-8601 upper-bound timestamp.
        limit: optional row cap; clamped to ``config.max_rows``.
    """
    if not config.is_configured:
        return _error_evidence("Not configured.", source=source)

    safe_source = _validate_source_name(source)
    if safe_source is None:
        return _error_evidence(
            f"Invalid Better Stack source identifier: {source!r}. "
            "Expected a ClickHouse bare identifier (e.g. t123456_myapp).",
            source=source,
        )

    since_sql = _validate_iso_timestamp(since)
    if since and since_sql is None:
        return _error_evidence(
            f"Invalid 'since' timestamp: {since!r}. Expected ISO-8601.",
            source=safe_source,
        )
    until_sql = _validate_iso_timestamp(until)
    if until and until_sql is None:
        return _error_evidence(
            f"Invalid 'until' timestamp: {until!r}. Expected ISO-8601.",
            source=safe_source,
        )

    effective_limit = min(max(1, int(limit or config.max_rows)), config.max_rows)
    sql = _build_logs_query(safe_source, since_sql, until_sql, effective_limit)

    try:
        with _sql_client(config) as client:
            response, err = _post_sql(client, config.query_endpoint, sql)
    except Exception as err:
        report_validation_failure(
            err,
            logger=logger,
            integration="betterstack",
            method="query_logs",
        )
        return _error_evidence(f"Better Stack connection failed: {err}", source=safe_source)

    if err is not None or response is None:
        return _error_evidence(
            err or "Better Stack request returned no response.",
            source=safe_source,
        )

    if response.status_code == 401:
        return _error_evidence(
            "Better Stack authentication failed (check credentials).",
            source=safe_source,
        )
    if response.status_code != 200:
        return _error_evidence(
            f"Better Stack query returned HTTP {response.status_code}: {response.text[:200]}",
            source=safe_source,
        )

    rows = _parse_json_each_row(response.text)
    return {
        "source": "betterstack",
        "available": True,
        "betterstack_source": safe_source,
        "rows": rows,
        "row_count": len(rows),
        "limit": effective_limit,
    }


__all__ = [
    "DEFAULT_BETTERSTACK_MAX_ROWS",
    "DEFAULT_BETTERSTACK_TIMEOUT_S",
    "BetterStackConfig",
    "BetterStackValidationResult",
    "betterstack_config_from_env",
    "betterstack_extract_params",
    "betterstack_is_available",
    "build_betterstack_config",
    "query_logs",
    "validate_betterstack_config",
]


def classify(
    credentials: dict[str, Any], record_id: str
) -> tuple[BetterStackConfig | None, str | None]:
    try:
        cfg = build_betterstack_config(
            {
                "query_endpoint": credentials.get("query_endpoint", ""),
                "username": credentials.get("username", ""),
                "password": credentials.get("password", ""),
                "sources": credentials.get("sources", []),
                "integration_id": record_id,
            }
        )
    except Exception as exc:
        report_classify_failure(exc, logger=logger, integration="betterstack", record_id=record_id)
        return None, None
    if cfg.query_endpoint and cfg.username:
        return cfg, "betterstack"
    return None, None