go-otel/internal/storage/clickhouse.go

package storage

import (
	"context"
	"encoding/json"
	"fmt"
	"os"
	"strconv"
	"strings"
	"time"

	"github.com/ClickHouse/clickhouse-go/v2"
	"github.com/ClickHouse/clickhouse-go/v2/lib/driver"

	"git.tornberg.me/go-otel/internal/websocket"
)

// ClickHouseStore implements the Store interface backed by ClickHouse, so
// history search stays fast and unbounded (the FileStore rescans the whole
// log file on every query) and aggregate time-series - error trends, release
// comparisons - can be computed server-side instead of in the browser.
type ClickHouseStore struct {
	conn driver.Conn
	db   string
}

// ClickHouseConfig controls how the store connects.
type ClickHouseConfig struct {
	Addr            string        // host:port, native protocol (default port 9000)
	Database        string        // defaults to "otel"
	Username        string        // defaults to "default"
	Password        string
	MaxOpenConns    int           // defaults to 50
	MaxIdleConns    int           // defaults to 25
	DialTimeout     time.Duration // defaults to 30s (controls connection acquire timeout in driver pool)
	ConnMaxLifetime time.Duration // defaults to 10m
}

func parseEnvInt(keys []string, def int) int {
	for _, k := range keys {
		if v := os.Getenv(k); v != "" {
			if n, err := strconv.Atoi(v); err == nil && n > 0 {
				return n
			}
		}
	}
	return def
}

func parseEnvDuration(keys []string, def time.Duration) time.Duration {
	for _, k := range keys {
		if v := os.Getenv(k); v != "" {
			if d, err := time.ParseDuration(v); err == nil && d > 0 {
				return d
			}
		}
	}
	return def
}

// NewClickHouseStore connects, creates the schema if needed, and returns a
// ready store. Non-table errors during schema init are logged, not fatal:
// raw storage works even if the rollup view can't be created.
func NewClickHouseStore(ctx context.Context, cfg ClickHouseConfig) (*ClickHouseStore, error) {
	if cfg.Database == "" {
		cfg.Database = "otel"
	}
	if cfg.Username == "" {
		cfg.Username = "default"
	}

	maxOpen := cfg.MaxOpenConns
	if maxOpen <= 0 {
		maxOpen = parseEnvInt([]string{"OTEL_CLICKHOUSE_MAX_OPEN_CONNS", "CLICKHOUSE_MAX_OPEN_CONNS"}, 50)
	}
	maxIdle := cfg.MaxIdleConns
	if maxIdle <= 0 {
		maxIdle = parseEnvInt([]string{"OTEL_CLICKHOUSE_MAX_IDLE_CONNS", "CLICKHOUSE_MAX_IDLE_CONNS"}, 25)
	}
	dialTimeout := cfg.DialTimeout
	if dialTimeout <= 0 {
		dialTimeout = parseEnvDuration([]string{"OTEL_CLICKHOUSE_DIAL_TIMEOUT", "CLICKHOUSE_DIAL_TIMEOUT"}, 30*time.Second)
	}
	connMaxLifetime := cfg.ConnMaxLifetime
	if connMaxLifetime <= 0 {
		connMaxLifetime = parseEnvDuration([]string{"OTEL_CLICKHOUSE_CONN_MAX_LIFETIME", "CLICKHOUSE_CONN_MAX_LIFETIME"}, 10*time.Minute)
	}

	// 1. Connect to default database with short timeout to ensure the target database exists.
	bootConn, err := clickhouse.Open(&clickhouse.Options{
		Addr: []string{cfg.Addr},
		Auth: clickhouse.Auth{
			Database: "default",
			Username: cfg.Username,
			Password: cfg.Password,
		},
		DialTimeout: 5 * time.Second,
	})
	if err != nil {
		return nil, fmt.Errorf("clickhouse dial %s: %w", cfg.Addr, err)
	}
	if err := bootConn.Ping(ctx); err != nil {
		_ = bootConn.Close()
		return nil, fmt.Errorf("clickhouse ping %s: %w", cfg.Addr, err)
	}

	if err := bootConn.Exec(ctx,
		fmt.Sprintf("CREATE DATABASE IF NOT EXISTS %s", cfg.Database)); err != nil {
		_ = bootConn.Close()
		return nil, fmt.Errorf("clickhouse create database %s: %w", cfg.Database, err)
	}
	// Close bootstrap connection so it does not leak idle connections.
	_ = bootConn.Close()

	// 2. Open production connection pool targeting the actual database with full pool settings.
	conn, err := clickhouse.Open(&clickhouse.Options{
		Addr: []string{cfg.Addr},
		Auth: clickhouse.Auth{
			Database: cfg.Database,
			Username: cfg.Username,
			Password: cfg.Password,
		},
		MaxOpenConns:     maxOpen,
		MaxIdleConns:     maxIdle,
		DialTimeout:      dialTimeout,
		ConnMaxLifetime:  connMaxLifetime,
		ConnOpenStrategy: clickhouse.ConnOpenRoundRobin,
		ClientInfo: clickhouse.ClientInfo{
			Products: []struct {
				Name    string
				Version string
			}{
				{Name: "go-otel", Version: "1.0"},
			},
		},
	})
	if err != nil {
		return nil, fmt.Errorf("clickhouse open %s/%s: %w", cfg.Addr, cfg.Database, err)
	}
	if err := conn.Ping(ctx); err != nil {
		_ = conn.Close()
		return nil, fmt.Errorf("clickhouse ping %s/%s: %w", cfg.Addr, cfg.Database, err)
	}

	s := &ClickHouseStore{conn: conn, db: cfg.Database}
	if err := s.initSchema(ctx); err != nil {
		_ = conn.Close()
		return nil, err
	}
	return s, nil
}

// initSchema creates the raw entries table and the materialized hourly
// rollup used by the stats endpoints. Idempotent.
func (s *ClickHouseStore) initSchema(ctx context.Context) error {
	if err := EnsureNormalizedSchema(ctx, s.conn, s.db); err != nil {
		return fmt.Errorf("normalized schema: %w", err)
	}
	if err := EnsureSourceIPSchema(ctx, s.conn, s.db); err != nil {
		return fmt.Errorf("source ip schema: %w", err)
	}
	if err := EnsureCSPSchema(ctx, s.conn, s.db); err != nil {
		return fmt.Errorf("csp report schema: %w", err)
	}
	if err := s.conn.Exec(ctx, fmt.Sprintf(`CREATE TABLE IF NOT EXISTS %s.entries (
		timestamp DateTime64(9, 'UTC') NOT NULL,
		type      LowCardinality(String) NOT NULL,
		service   LowCardinality(String) NOT NULL DEFAULT 'unknown',
		is_error  UInt8 NOT NULL DEFAULT 0,
		release   LowCardinality(String) NOT NULL DEFAULT '',
		trace_id  String NOT NULL DEFAULT '',
		data      String NOT NULL
	) ENGINE = MergeTree
	PARTITION BY toYYYYMM(timestamp)
	ORDER BY (type, timestamp)
	TTL toDateTime(timestamp) + INTERVAL 90 DAY`, s.db)); err != nil {
		return fmt.Errorf("clickhouse create entries: %w", err)
	}

	// Hourly per-service/release rollup maintained on insert. is_error is a
	// dimension (0/1), so errors = sum(count) WHERE is_error=1 and
	// total = sum(count) unfiltered.
	if err := s.conn.Exec(ctx, fmt.Sprintf(`CREATE TABLE IF NOT EXISTS %s.hourly (
		hour     DateTime('UTC') NOT NULL,
		type     LowCardinality(String) NOT NULL,
		service  LowCardinality(String) NOT NULL,
		release  LowCardinality(String) NOT NULL DEFAULT '',
		is_error UInt8 NOT NULL DEFAULT 0,
		count    UInt64
	) ENGINE = SummingMergeTree
	ORDER BY (hour, type, service, release, is_error)`, s.db)); err != nil {
		return fmt.Errorf("clickhouse create hourly: %w", err)
	}

	mv := fmt.Sprintf(`CREATE MATERIALIZED VIEW IF NOT EXISTS %s.hourly_mv TO %s.hourly AS
		SELECT
			toStartOfHour(timestamp) AS hour,
			type,
			service,
			release,
			is_error,
			count() AS count
		FROM %s.entries
		GROUP BY hour, type, service, release, is_error`, s.db, s.db, s.db)
	if err := s.conn.Exec(ctx, mv); err != nil {
		// Raw storage works without the rollup; stats endpoints will just be
		// empty until this succeeds.
		fmt.Printf("clickhouse: failed to create rollup view (stats endpoints stay empty): %v\n", err)
	}
	return s.initTraceSchema(ctx)
}

// chPayload extracts rollup columns from the marshaled OTLP JSON. Field
// structs are deliberately minimal so unmarshal cost stays low.
type chPayload struct {
	ResourceSpans []chResourceSpan `json:"resourceSpans"`
	ResourceLogs  []chResourceLog  `json:"resourceLogs"`
}

type chResourceSpan struct {
	Resource   chResource `json:"resource"`
	ScopeSpans []struct {
		Spans []chSpan `json:"spans"`
	} `json:"scopeSpans"`
}

type chResourceLog struct {
	Resource  chResource `json:"resource"`
	ScopeLogs []struct {
		LogRecords []chLogRecord `json:"logRecords"`
	} `json:"scopeLogs"`
}

type chResource struct {
	Attributes []chAttr `json:"attributes"`
}

type chAttr struct {
	Key   string `json:"key"`
	Value struct {
		StringValue string `json:"stringValue"`
	} `json:"value"`
}

type chSpan struct {
	TraceID string `json:"traceId"`
	Status  struct {
		Code int `json:"code"`
	} `json:"status"`
	Attributes []chAttr `json:"attributes"`
}

type chLogRecord struct {
	SeverityNumber int  `json:"severityNumber"`
	SeverityText   string `json:"severityText"`
	TraceID        string `json:"traceId"`
	Attributes     []chAttr `json:"attributes"`
}

// extractHead pulls (service, is_error, release, trace_id) out of a raw OTLP
// JSON payload for the rollup columns. Best-effort: on parse failure the
// entry is still stored with defaults.
func extractHead(data []byte) (service string, isErr bool, release, traceID string) {
	service, release = "unknown", ""
	var p chPayload
	if err := json.Unmarshal(data, &p); err != nil {
		return service, false, release, traceID
	}

	attrString := func(attrs []chAttr, key string) string {
		for _, a := range attrs {
			if a.Key == key {
				return a.Value.StringValue
			}
		}
		return ""
	}

	// Traces: error if any span has status code 2 (StatusCodeError).
	for _, rs := range p.ResourceSpans {
		if s := attrString(rs.Resource.Attributes, "service.name"); s != "" && service == "unknown" {
			service = s
		}
		if r := attrString(rs.Resource.Attributes, "service.release"); r != "" && release == "" {
			release = r
		}
		for _, ss := range rs.ScopeSpans {
			for _, span := range ss.Spans {
				if span.TraceID != "" && traceID == "" {
					traceID = span.TraceID
				}
				if span.Status.Code == 2 {
					isErr = true
				}
				if r := attrString(span.Attributes, "service.release"); r != "" && release == "" {
					release = r
				}
			}
		}
	}

	// Logs: error if severity number >= 17 (Error) or text is error-ish.
	for _, rl := range p.ResourceLogs {
		if s := attrString(rl.Resource.Attributes, "service.name"); s != "" && service == "unknown" {
			service = s
		}
		if r := attrString(rl.Resource.Attributes, "service.release"); r != "" && release == "" {
			release = r
		}
		for _, sl := range rl.ScopeLogs {
			for _, rec := range sl.LogRecords {
				if rec.TraceID != "" && traceID == "" {
					traceID = rec.TraceID
				}
				if r := attrString(rec.Attributes, "service.release"); r != "" && release == "" {
					release = r
				}
				if rec.SeverityNumber >= 17 {
					isErr = true
				}
				switch strings.ToUpper(strings.TrimSpace(rec.SeverityText)) {
				case "ERROR", "FATAL", "CRITICAL":
					isErr = true
				}
			}
		}
	}
	return service, isErr, release, traceID
}

// Append stores one history entry with its extracted rollup columns.
// AsyncInsert lets ClickHouse batch the writes server-side. The severity
// column is intentionally omitted: one payload can hold many records with
// differing severities, so error-ness is a per-entry flag instead.
func (s *ClickHouseStore) Append(ctx context.Context, entry websocket.HistoryEntry) error {
	service, isErr, release, traceID := extractHead(entry.Data)

	// A batch routinely holds several traces, so keep every trace id (the
	// scalar trace_id column only records the first) for indexed lookups.
	var facts []spanFact
	var traceIDs []string
	switch entry.Type {
	case "traces":
		facts = extractSpanFacts(entry.Data)
		seen := make(map[string]struct{}, len(facts))
		for _, f := range facts {
			if _, ok := seen[f.TraceID]; !ok {
				seen[f.TraceID] = struct{}{}
				traceIDs = append(traceIDs, f.TraceID)
			}
		}
	case "logs":
		traceIDs = allTraceIDs(entry.Data)
	}
	if traceIDs == nil {
		traceIDs = []string{}
	}

	ts := time.Unix(0, entry.Timestamp).UTC()
	err := s.conn.AsyncInsert(ctx, fmt.Sprintf(
		`INSERT INTO %s.entries
		 (timestamp, type, service, is_error, release, trace_id, trace_ids, data)
		 VALUES (?, ?, ?, ?, ?, ?, ?, ?)`, s.db),
		false,
		ts, entry.Type, service, boolToUint8(isErr), release, traceID, traceIDs, string(entry.Data),
	)
	if err != nil {
		return fmt.Errorf("clickhouse append: %w", err)
	}
	if err := s.appendSpanFacts(ctx, ts, facts); err != nil {
		return fmt.Errorf("clickhouse span facts: %w", err)
	}
	s.appendNormalized(ctx, entry)
	return nil
}

func boolToUint8(b bool) uint8 {
	if b {
		return 1
	}
	return 0
}

// GetHistory returns the most recent limit entries, newest first.
func (s *ClickHouseStore) GetHistory(ctx context.Context, limit int) ([]websocket.HistoryEntry, error) {
	return s.Search(ctx, SearchOptions{}, limit)
}

// Search implements the same semantics as FileStore.Search: case-insensitive
// substring query over the raw payload, type/time filters, newest first,
// capped at limit.
func (s *ClickHouseStore) Search(ctx context.Context, opts SearchOptions, limit int) ([]websocket.HistoryEntry, error) {
	if limit <= 0 {
		limit = 50
	}

	where := []string{"1=1"}
	args := []interface{}{}
	if opts.Query != "" {
		where = append(where, "positionCaseInsensitive(data, ?) > 0")
		args = append(args, opts.Query)
	}
	if opts.Type != "" {
		where = append(where, "type = ?")
		args = append(args, opts.Type)
	}
	if opts.Since != 0 {
		where = append(where, "timestamp >= ?")
		args = append(args, time.Unix(0, opts.Since).UTC())
	}
	if opts.Until != 0 {
		where = append(where, "timestamp <= ?")
		args = append(args, time.Unix(0, opts.Until).UTC())
	}

	query := fmt.Sprintf(
		`SELECT timestamp, type, data
		 FROM %s.entries
		 WHERE %s
		 ORDER BY timestamp DESC
		 LIMIT ?`, s.db, strings.Join(where, " AND "))
	args = append(args, limit)

	rows, err := s.conn.Query(ctx, query, args...)
	if err != nil {
		return nil, fmt.Errorf("clickhouse search: %w", err)
	}
	defer rows.Close()

	var results []websocket.HistoryEntry
	for rows.Next() {
		var (
			ts  time.Time
			typ string
			raw string
		)
		if err := rows.Scan(&ts, &typ, &raw); err != nil {
			return nil, err
		}
		results = append(results, websocket.HistoryEntry{
			Type:      typ,
			Timestamp: ts.UnixNano(),
			Data:      json.RawMessage(raw),
		})
	}
	return results, rows.Err()
}

// StatsPoint is one bucket of a stats time-series.
type StatsPoint struct {
	Bucket string `json:"bucket"`
	Total  uint64 `json:"total"`
	Errors uint64 `json:"errors"`
}

// ErrorStats returns hourly (or daily) total/error entry counts over the
// lookback window, for the trends graphs. Empty buckets are absent; the UI
// fills gaps for display.
func (s *ClickHouseStore) ErrorStats(ctx context.Context, service string, hours int, daily bool) ([]StatsPoint, error) {
	bucket := "toStartOfHour(hour)"
	if daily {
		bucket = "toStartOfDay(hour)"
	}
	where := []string{"hour >= ?"}
	args := []interface{}{time.Now().UTC().Add(-time.Duration(hours) * time.Hour)}
	if service != "" {
		where = append(where, "service = ?")
		args = append(args, service)
	}

	query := fmt.Sprintf(
		`SELECT
			toString(%s) AS bucket,
			sum(count) AS total,
			sumIf(count, is_error = 1) AS errors
		 FROM %s.hourly
		 WHERE %s
		 GROUP BY bucket
		 ORDER BY bucket`, bucket, s.db, strings.Join(where, " AND "))

	rows, err := s.conn.Query(ctx, query, args...)
	if err != nil {
		return nil, fmt.Errorf("clickhouse error stats: %w", err)
	}
	defer rows.Close()

	var out []StatsPoint
	for rows.Next() {
		var p StatsPoint
		if err := rows.Scan(&p.Bucket, &p.Total, &p.Errors); err != nil {
			return nil, err
		}
		out = append(out, p)
	}
	return out, rows.Err()
}

// ReleaseSeries returns one error-count series per release:
// {release -> [{bucket, total, errors}]}, so the UI can overlay deploys.
func (s *ClickHouseStore) ReleaseSeries(ctx context.Context, service string, hours int) (map[string][]StatsPoint, error) {
	where := []string{"hour >= ?", "release != ''"}
	args := []interface{}{time.Now().UTC().Add(-time.Duration(hours) * time.Hour)}
	if service != "" {
		where = append(where, "service = ?")
		args = append(args, service)
	}

	query := fmt.Sprintf(
		`SELECT
			release,
			toString(toStartOfHour(hour)) AS bucket,
			sum(count) AS total,
			sumIf(count, is_error = 1) AS errors
		 FROM %s.hourly
		 WHERE %s
		 GROUP BY release, bucket
		 ORDER BY bucket`, s.db, strings.Join(where, " AND "))

	rows, err := s.conn.Query(ctx, query, args...)
	if err != nil {
		return nil, fmt.Errorf("clickhouse release series: %w", err)
	}
	defer rows.Close()

	out := map[string][]StatsPoint{}
	for rows.Next() {
		var (
			release string
			p       StatsPoint
		)
		if err := rows.Scan(&release, &p.Bucket, &p.Total, &p.Errors); err != nil {
			return nil, err
		}
		out[release] = append(out[release], p)
	}
	return out, rows.Err()
}

// Services returns the distinct services seen in the rollup, for filter UIs.
func (s *ClickHouseStore) Services(ctx context.Context) ([]string, error) {
	rows, err := s.conn.Query(ctx, fmt.Sprintf(
		`SELECT DISTINCT service FROM %s.hourly ORDER BY service`, s.db))
	if err != nil {
		return nil, err
	}
	defer rows.Close()

	var out []string
	for rows.Next() {
		var svc string
		if err := rows.Scan(&svc); err != nil {
			return nil, err
		}
		out = append(out, svc)
	}
	return out, rows.Err()
}

// Close closes the underlying connection.
// Conn and Database expose the connection for read-only analysis jobs such
// as the alert runner.
func (s *ClickHouseStore) Conn() driver.Conn { return s.conn }
func (s *ClickHouseStore) Database() string  { return s.db }

func (s *ClickHouseStore) Close() error {
	return s.conn.Close()
}