token_thief/db/migrate.go

package db

import (
	"context"
	"fmt"
	"strings"

	"github.com/ClickHouse/clickhouse-go/v2"
)

const schemaSQL = `
CREATE TABLE IF NOT EXISTS proxy_logs (
    request_id String,
    method String,
    path String,
    query String,
    client_ip String,
    request_headers String,
    request_body String,
    request_truncated Bool DEFAULT false,
    status_code Int32,
    response_headers String,
    response_body String,
    response_truncated Bool DEFAULT false,
    is_stream Bool DEFAULT false,
    latency_ms Int64,
    started_at DateTime64(3),
    finished_at DateTime64(3),
    error String
)
ENGINE = MergeTree
PARTITION BY toYYYYMM(started_at)
ORDER BY (started_at, request_id)
`

func migrate(ctx context.Context, conn clickhouse.Conn) error {
	if err := conn.Exec(ctx, schemaSQL); err != nil {
		return fmt.Errorf("create proxy_logs: %w", err)
	}

	columns, err := queryProxyLogColumns(ctx, conn)
	if err != nil {
		return err
	}
	table, err := queryProxyLogsTable(ctx, conn)
	if err != nil {
		return err
	}
	if err := validateProxyLogsTable(table); err != nil {
		return fmt.Errorf("incompatible proxy_logs schema: %w", err)
	}
	statements, err := missingProxyLogColumns(columns)
	if err != nil {
		return fmt.Errorf("plan proxy_logs migration: %w", err)
	}
	for _, statement := range statements {
		if err := conn.Exec(ctx, statement); err != nil {
			return fmt.Errorf("alter proxy_logs: %w", err)
		}
	}
	if len(statements) > 0 {
		columns, err = queryProxyLogColumns(ctx, conn)
		if err != nil {
			return err
		}
	}

	table, err = queryProxyLogsTable(ctx, conn)
	if err != nil {
		return err
	}
	if err := validateProxyLogsSchema(columns, table); err != nil {
		return fmt.Errorf("incompatible proxy_logs schema: %w", err)
	}
	return nil
}

func queryProxyLogsTable(ctx context.Context, conn clickhouse.Conn) (schemaTable, error) {
	var table schemaTable
	err := conn.QueryRow(ctx, `
SELECT engine, partition_key, sorting_key
FROM system.tables
WHERE database = currentDatabase() AND name = 'proxy_logs'`).Scan(
		&table.engine, &table.partitionKey, &table.sortingKey,
	)
	if err != nil {
		return schemaTable{}, fmt.Errorf("query proxy_logs table: %w", err)
	}
	return table, nil
}

func queryProxyLogColumns(ctx context.Context, conn clickhouse.Conn) ([]schemaColumn, error) {
	rows, err := conn.Query(ctx, `
SELECT name, type
FROM system.columns
WHERE database = currentDatabase() AND table = 'proxy_logs'
ORDER BY position`)
	if err != nil {
		return nil, fmt.Errorf("query proxy_logs columns: %w", err)
	}
	var columns []schemaColumn
	for rows.Next() {
		var column schemaColumn
		if err := rows.Scan(&column.name, &column.typ); err != nil {
			rows.Close()
			return nil, fmt.Errorf("scan proxy_logs columns: %w", err)
		}
		columns = append(columns, column)
	}
	if err := rows.Err(); err != nil {
		rows.Close()
		return nil, fmt.Errorf("read proxy_logs columns: %w", err)
	}
	rows.Close()
	return columns, nil
}

type schemaColumn struct {
	name string
	typ  string
}

type schemaTable struct {
	engine       string
	partitionKey string
	sortingKey   string
}

var proxyLogsColumns = []schemaColumn{
	{"request_id", "String"},
	{"method", "String"},
	{"path", "String"},
	{"query", "String"},
	{"client_ip", "String"},
	{"request_headers", "String"},
	{"request_body", "String"},
	{"request_truncated", "Bool"},
	{"status_code", "Int32"},
	{"response_headers", "String"},
	{"response_body", "String"},
	{"response_truncated", "Bool"},
	{"is_stream", "Bool"},
	{"latency_ms", "Int64"},
	{"started_at", "DateTime64(3)"},
	{"finished_at", "DateTime64(3)"},
	{"error", "String"},
}

var proxyLogColumnDDL = map[string]string{
	"request_id":         "request_id String",
	"method":             "method String",
	"path":               "path String",
	"query":              "query String",
	"client_ip":          "client_ip String",
	"request_headers":    "request_headers String",
	"request_body":       "request_body String",
	"request_truncated":  "request_truncated Bool DEFAULT false",
	"status_code":        "status_code Int32",
	"response_headers":   "response_headers String",
	"response_body":      "response_body String",
	"response_truncated": "response_truncated Bool DEFAULT false",
	"is_stream":          "is_stream Bool DEFAULT false",
	"latency_ms":         "latency_ms Int64",
	"started_at":         "started_at DateTime64(3)",
	"finished_at":        "finished_at DateTime64(3)",
	"error":              "error String",
}

func missingProxyLogColumns(columns []schemaColumn) ([]string, error) {
	existing := make(map[string]string, len(columns))
	for _, column := range columns {
		existing[column.name] = column.typ
	}
	var statements []string
	for _, required := range proxyLogsColumns {
		if typ, ok := existing[required.name]; ok {
			if typ != required.typ {
				return nil, fmt.Errorf("column %s has type %s, want %s", required.name, typ, required.typ)
			}
			continue
		}
		statements = append(statements, "ALTER TABLE proxy_logs ADD COLUMN IF NOT EXISTS "+proxyLogColumnDDL[required.name])
	}
	return statements, nil
}

func validateProxyLogsSchema(columns []schemaColumn, table schemaTable) error {
	existing := make(map[string]string, len(columns))
	for _, column := range columns {
		existing[column.name] = column.typ
	}
	for _, want := range proxyLogsColumns {
		if typ, ok := existing[want.name]; !ok || typ != want.typ {
			return fmt.Errorf("column %s is %s, want %s", want.name, typ, want.typ)
		}
	}
	return validateProxyLogsTable(table)
}

func validateProxyLogsTable(table schemaTable) error {
	if table.engine != "MergeTree" {
		return fmt.Errorf("engine is %q, want MergeTree", table.engine)
	}
	if compactExpression(table.partitionKey) != "toYYYYMM(started_at)" {
		return fmt.Errorf("partition key is %q, want toYYYYMM(started_at)", table.partitionKey)
	}
	if compactExpression(table.sortingKey) != "started_at,request_id" {
		return fmt.Errorf("sorting key is %q, want started_at, request_id", table.sortingKey)
	}
	return nil
}

func compactExpression(value string) string {
	return strings.Join(strings.Fields(value), "")
}