修复日志持久化可靠性与数据库安全默认值

db/migrate.go

@@ -38,43 +38,81 @@ func migrate(ctx context.Context, conn clickhouse.Conn) error {
 		return fmt.Errorf("create proxy_logs: %w", err)
 	}
 
-	rows, err := conn.Query(ctx, `
-SELECT name, type
-FROM system.columns
-WHERE database = currentDatabase() AND table = 'proxy_logs'
-ORDER BY position`)
+	columns, err := queryProxyLogColumns(ctx, conn)
 	if err != nil {
-		return fmt.Errorf("query proxy_logs columns: %w", err)
+		return err
 	}
-	var columns []schemaColumn
-	for rows.Next() {
-		var column schemaColumn
-		if err := rows.Scan(&column.name, &column.typ); err != nil {
-			rows.Close()
-			return fmt.Errorf("scan proxy_logs columns: %w", 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)
 		}
-		columns = append(columns, column)
 	}
-	if err := rows.Err(); err != nil {
-		rows.Close()
-		return fmt.Errorf("read proxy_logs columns: %w", err)
+	if len(statements) > 0 {
+		columns, err = queryProxyLogColumns(ctx, conn)
+		if err != nil {
+			return err
+		}
 	}
-	rows.Close()
 
+	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, `
+	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 fmt.Errorf("query proxy_logs table: %w", err)
+		return schemaTable{}, fmt.Errorf("query proxy_logs table: %w", err)
 	}
-	if err := validateProxyLogsSchema(columns, table); err != nil {
-		return fmt.Errorf("incompatible proxy_logs schema: %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)
 	}
-	return nil
+	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 {
@@ -108,15 +146,58 @@ var proxyLogsColumns = []schemaColumn{
 	{"error", "String"},
 }
 
-func validateProxyLogsSchema(columns []schemaColumn, table schemaTable) error {
-	if len(columns) != len(proxyLogsColumns) {
-		return fmt.Errorf("got %d columns, want %d", len(columns), len(proxyLogsColumns))
+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
 	}
-	for i, want := range proxyLogsColumns {
-		if columns[i] != want {
-			return fmt.Errorf("column %d is %s %s, want %s %s", i+1, columns[i].name, columns[i].typ, want.name, want.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)
 	}