aiops/backend/controller/trace_datasource.go

package controller

import (
	"crypto/tls"
	"encoding/json"
	"io"
	"net/http"
	"strings"
	"time"

	"ezops/model"

	"github.com/gin-gonic/gin"
	"gorm.io/gorm"
)

// TraceDSHandler 链路追踪数据源管理 // Logan.Li
type TraceDSHandler struct {
	db *gorm.DB
}

func NewTraceDSHandler(db *gorm.DB) *TraceDSHandler {
	return &TraceDSHandler{db: db}
}

// List GET /api/v1/trace-datasources?type=jaeger
func (h *TraceDSHandler) List(c *gin.Context) {
	q := h.db.Order("is_default DESC, id ASC")
	if t := c.Query("type"); t != "" {
		q = q.Where("type = ?", t)
	}
	var sources []model.TraceDataSource
	if err := q.Find(&sources).Error; err != nil {
		c.JSON(500, model.Response{Code: 500, Message: err.Error()})
		return
	}
	if sources == nil {
		sources = []model.TraceDataSource{}
	}
	c.JSON(200, model.Response{Code: 200, Data: sources})
}

// Get GET /api/v1/trace-datasources/:id
func (h *TraceDSHandler) Get(c *gin.Context) {
	var ds model.TraceDataSource
	if err := h.db.First(&ds, c.Param("id")).Error; err != nil {
		c.JSON(404, model.Response{Code: 404, Message: "数据源不存在"})
		return
	}
	c.JSON(200, model.Response{Code: 200, Data: ds})
}

// Create POST /api/v1/trace-datasources
func (h *TraceDSHandler) Create(c *gin.Context) {
	var req model.CreateTraceDSReq
	if err := c.ShouldBindJSON(&req); err != nil {
		c.JSON(400, model.Response{Code: 400, Message: err.Error()})
		return
	}
	req.Type = strings.ToLower(req.Type)
	if req.Type != "jaeger" && req.Type != "tempo" {
		c.JSON(400, model.Response{Code: 400, Message: "type 必须为 jaeger 或 tempo"})
		return
	}
	if req.IsDefault {
		h.db.Model(&model.TraceDataSource{}).Where("type = ? AND is_default = ?", req.Type, true).Update("is_default", false)
	}
	ds := model.TraceDataSource{
		Name: req.Name, Type: req.Type, URL: strings.TrimRight(req.URL, "/"),
		BasicAuth: req.BasicAuth, BasicAuthUser: req.BasicAuthUser, BasicAuthPass: req.BasicAuthPass,
		SkipTLSVerify: req.SkipTLSVerify, IsDefault: req.IsDefault, Status: 1,
	}
	if err := h.db.Create(&ds).Error; err != nil {
		c.JSON(500, model.Response{Code: 500, Message: err.Error()})
		return
	}
	c.JSON(200, model.Response{Code: 200, Message: "创建成功", Data: ds})
}

// Update PUT /api/v1/trace-datasources/:id
func (h *TraceDSHandler) Update(c *gin.Context) {
	var ds model.TraceDataSource
	if err := h.db.First(&ds, c.Param("id")).Error; err != nil {
		c.JSON(404, model.Response{Code: 404, Message: "数据源不存在"})
		return
	}
	var req model.UpdateTraceDSReq
	if err := c.ShouldBindJSON(&req); err != nil {
		c.JSON(400, model.Response{Code: 400, Message: err.Error()})
		return
	}
	updates := map[string]interface{}{}
	if req.Name != nil {
		updates["name"] = *req.Name
	}
	if req.Type != nil {
		updates["type"] = strings.ToLower(*req.Type)
	}
	if req.URL != nil {
		updates["url"] = strings.TrimRight(*req.URL, "/")
	}
	if req.BasicAuth != nil {
		updates["basic_auth"] = *req.BasicAuth
	}
	if req.BasicAuthUser != nil {
		updates["basic_auth_user"] = *req.BasicAuthUser
	}
	if req.BasicAuthPass != nil && *req.BasicAuthPass != "" {
		updates["basic_auth_pass"] = *req.BasicAuthPass
	}
	if req.SkipTLSVerify != nil {
		updates["skip_tls_verify"] = *req.SkipTLSVerify
	}
	if req.IsDefault != nil {
		if *req.IsDefault {
			h.db.Model(&model.TraceDataSource{}).Where("type = ? AND is_default = ? AND id != ?", ds.Type, true, ds.ID).Update("is_default", false)
		}
		updates["is_default"] = *req.IsDefault
	}
	if req.Status != nil {
		updates["status"] = *req.Status
	}
	if len(updates) == 0 {
		c.JSON(200, model.Response{Code: 200, Data: ds})
		return
	}
	if err := h.db.Model(&ds).Updates(updates).Error; err != nil {
		c.JSON(500, model.Response{Code: 500, Message: err.Error()})
		return
	}
	h.db.First(&ds, ds.ID)
	c.JSON(200, model.Response{Code: 200, Message: "更新成功", Data: ds})
}

// Delete DELETE /api/v1/trace-datasources/:id
func (h *TraceDSHandler) Delete(c *gin.Context) {
	if err := h.db.Delete(&model.TraceDataSource{}, c.Param("id")).Error; err != nil {
		c.JSON(500, model.Response{Code: 500, Message: err.Error()})
		return
	}
	c.JSON(200, model.Response{Code: 200, Message: "删除成功"})
}

// SetDefault POST /api/v1/trace-datasources/:id/default
func (h *TraceDSHandler) SetDefault(c *gin.Context) {
	var ds model.TraceDataSource
	if err := h.db.First(&ds, c.Param("id")).Error; err != nil {
		c.JSON(404, model.Response{Code: 404, Message: "数据源不存在"})
		return
	}
	h.db.Model(&model.TraceDataSource{}).Where("type = ? AND is_default = ?", ds.Type, true).Update("is_default", false)
	h.db.Model(&ds).Update("is_default", true)
	c.JSON(200, model.Response{Code: 200, Message: "设置成功"})
}

// TestConnection POST /api/v1/trace-datasources/test
func (h *TraceDSHandler) TestConnection(c *gin.Context) {
	var req model.TestTraceDSReq
	if err := c.ShouldBindJSON(&req); err != nil {
		// 允许传 id 来测试已有数据源
		id := c.Query("id")
		if id != "" {
			var ds model.TraceDataSource
			if e := h.db.First(&ds, id).Error; e == nil {
				req.URL = ds.URL
				req.Type = ds.Type
				req.BasicAuth = ds.BasicAuth
				req.BasicAuthUser = ds.BasicAuthUser
				req.BasicAuthPass = ds.BasicAuthPass
				req.SkipTLSVerify = ds.SkipTLSVerify
			}
		}
		if req.URL == "" {
			c.JSON(400, model.Response{Code: 400, Message: "url 不能为空"})
			return
		}
	}
	url := strings.TrimRight(req.URL, "/")
	client := &http.Client{Timeout: 10 * time.Second}
	if req.SkipTLSVerify {
		client.Transport = &http.Transport{TLSClientConfig: &tls.Config{InsecureSkipVerify: true}}
	}
	// Jaeger/Tempo 都有 /api/services 端点
	target := url + "/api/services"
	httpReq, _ := http.NewRequest("GET", target, nil)
	if req.BasicAuth && req.BasicAuthUser != "" {
		httpReq.SetBasicAuth(req.BasicAuthUser, req.BasicAuthPass)
	}
	resp, err := client.Do(httpReq)
	if err != nil {
		c.JSON(200, model.Response{Code: 200, Data: map[string]interface{}{"ok": false, "error": err.Error()}})
		return
	}
	defer resp.Body.Close()
	body, _ := io.ReadAll(resp.Body)
	if resp.StatusCode != 200 {
		c.JSON(200, model.Response{Code: 200, Data: map[string]interface{}{"ok": false, "error": "HTTP " + resp.Status}})
		return
	}
	// 尝试解析 JSON 验证有效性
	var tmp interface{}
	if json.Unmarshal(body, &tmp) != nil {
		c.JSON(200, model.Response{Code: 200, Data: map[string]interface{}{"ok": false, "error": "响应非 JSON"}})
		return
	}
	c.JSON(200, model.Response{Code: 200, Data: map[string]interface{}{"ok": true, "message": "连接成功"}})
}

// getDefaultTraceDS 获取默认链路追踪数据源
func (h *TraceDSHandler) getDefaultTraceDS() (*model.TraceDataSource, error) {
	var ds model.TraceDataSource
	err := h.db.Where("is_default = ? AND status = 1", true).First(&ds).Error
	if err != nil {
		// fallback: 取第一个启用的
		err = h.db.Where("status = 1").First(&ds).Error
	}
	if err != nil {
		return nil, err
	}
	return &ds, nil
}