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), "") }