db_view/service/internal/modules/data_query/handler/handler.go

// Package handler 提供数据查询模块的HTTP请求处理器
// 负责处理数据查询、SQL执行等相关API请求，协调前端请求与后端服务层的交互
package handler

import (
	"encoding/json"
	"fmt"
	"net/http"

	"dbview/service/internal/common/databases/meta"
	"dbview/service/internal/common/response"
	api "dbview/service/internal/modules/data_query/api"
	svc "dbview/service/internal/modules/data_query/service"

	"github.com/gin-gonic/gin"
	"github.com/gorilla/websocket"
)

// Handler 数据查询处理器结构体
// 封装了数据查询服务层的调用逻辑，处理HTTP请求和响应
type Handler struct {
	svc *svc.DataService // 数据查询服务实例
}

// NewHandler 创建新的数据查询处理器
// 参数：
//   - s: 数据查询服务实例
//
// 返回：
//   - *Handler: 初始化后的处理器实例
func NewHandler(s *svc.DataService) *Handler {
	return &Handler{svc: s}
}

// QueryDataHandler 处理数据查询请求 POST /data/query
// 支持同步和异步两种查询模式，根据请求参数决定执行方式
// 请求体：api.QueryDataRequest 包含连接信息、查询路径、查询条件等
// 响应：统一的数据操作响应格式，包含查询结果或异步任务ID
func (h *Handler) QueryDataHandler(c *gin.Context) {
	// 解析请求参数，验证必填字段
	var req api.QueryDataRequest
	if err := c.ShouldBindJSON(&req); err != nil {
		// 参数绑定失败，返回客户端错误
		c.JSON(http.StatusBadRequest, response.Response{Code: 1, Msg: "请求参数错误", Error: err.Error()})
		return
	}

	// 调用服务层执行数据查询操作
	// 支持根据路径定位数据库对象，并应用查询条件进行数据筛选
	asyncMode := req.Async.Mode
	isAsync := asyncMode != api.AsyncModeSync
	// 将前端对大对象字段的控制传递到服务层
	includeLarge := req.IncludeLarge
	result, err := h.svc.QueryData(c.Request.Context(), req.ConnID, req.ConnInfo, req.Path, req.Query, isAsync, includeLarge)
	if err != nil {
		// 服务层调用失败，返回服务器内部错误
		c.JSON(http.StatusInternalServerError, response.Response{Code: 2, Msg: "查询数据失败", Error: err.Error()})
		return
	}

	// 根据查询模式处理不同的返回结果
	// 异步模式返回任务ID，同步模式返回查询结果
	if taskId, ok := result.(string); ok {
		// 异步查询：任务已提交到后台队列，返回任务ID供后续查询状态
		asyncResp := api.AsyncQueryResponse{
			Success: true,
			TaskID:  taskId,
			Message: "查询已提交到任务队列",
		}
		// 如果客户端以 WebSocket 模式请求异步，则返回可用的 ws 链接
		if req.Async.Mode == api.AsyncModeWS {
			scheme := "ws"
			if c.Request.TLS != nil {
				scheme = "wss"
			}
			asyncResp.WsURL = fmt.Sprintf("%s://%s/database/data/tasks/ws", scheme, c.Request.Host)
		}
		c.JSON(http.StatusOK, response.Response{Code: 0, Msg: "ok", Data: asyncResp})
		return
	}

	// 同步查询：直接返回查询结果数据
	if queryResult, ok := result.(meta.QueryResult); ok {
		// 构建统一的数据操作响应，包含查询结果、执行时间和成功状态
		resp := api.DataOperationResponse{
			Success: true,
			Data:    queryResult,
			Message: "查询完成",
		}
		c.JSON(http.StatusOK, response.Response{Code: 0, Msg: "ok", Data: resp})
		return
	}

	// 理论上不会到达这里，但为了代码健壮性保留错误处理
	c.JSON(http.StatusInternalServerError, response.Response{Code: 2, Msg: "未知的返回类型", Error: "service returned unexpected type"})
}

// GetTaskResultHandler 处理异步查询任务结果获取请求 GET /data/task/result
// 用于查询异步数据查询任务的执行状态和结果
// 当任务完成时返回查询结果数据，当任务进行中时只返回状态信息
// 请求体：api.GetTaskResultRequest 包含任务ID
// 响应：包含任务状态和可能的查询结果数据
func (h *Handler) GetTaskResultHandler(c *gin.Context) {
	// 解析请求参数，获取要查询的任务ID
	var req api.GetTaskResultRequest
	if err := c.ShouldBindJSON(&req); err != nil {
		// 参数验证失败，返回客户端错误
		c.JSON(http.StatusBadRequest, response.Response{Code: 1, Msg: "请求参数错误", Error: err.Error()})
		return
	}

	// 调用服务层查询指定任务的执行结果和状态
	// 返回查询结果数据、任务状态和可能的错误信息
	result, status, err := h.svc.GetTaskResult(c.Request.Context(), req.TaskID)
	if err != nil {
		// 获取任务结果失败，返回服务器内部错误
		c.JSON(http.StatusInternalServerError, response.Response{Code: 2, Msg: "获取任务结果失败", Error: err.Error()})
		return
	}

	// 构建基础响应结构，包含任务状态信息
	resp := api.GetTaskResultResponse{
		Success: true,
		Status:  status,
		Message: "获取任务结果成功",
	}

	// 如果任务已完成，填充查询结果数据
	// 只有在任务完成状态下才有实际的查询结果
	if status == "completed" {
		resp.Data = api.DataOperationResponse{
			Success: true,
			Data:    result,
			Message: "查询完成",
		}
	}

	// 返回任务状态和可能的查询结果
	c.JSON(http.StatusOK, response.Response{Code: 0, Msg: "ok", Data: resp})
}

// TaskWSHandler 处理 WebSocket 连接，用于订阅任务的实时推送
// 客户端需要先建立 websocket 连接，然后发送 JSON 消息订阅/退订任务：
// {"action":"subscribe","taskId":"..."} / {"action":"unsubscribe","taskId":"..."}
func (h *Handler) TaskWSHandler(c *gin.Context) {
	upgrader := websocket.Upgrader{
		CheckOrigin: func(r *http.Request) bool { return true },
	}
	conn, err := upgrader.Upgrade(c.Writer, c.Request, nil)
	if err != nil {
		return
	}

	// Register connection with manager to start pumps and receive inbound messages
	inbound, err := h.svc.RegisterWSConn(conn)
	if err != nil {
		_ = conn.Close()
		return
	}

	subscribed := make(map[string]struct{})
	// 监听 inbound 通道，处理客户端订阅/退订消息
	for msgB := range inbound {
		var msg struct {
			Action string `json:"action"`
			TaskID string `json:"taskId"`
		}
		if err := json.Unmarshal(msgB, &msg); err != nil {
			// 忽略解析错误
			continue
		}
		switch msg.Action {
		case "subscribe":
			if msg.TaskID == "" {
				continue
			}
			if err := h.svc.SubscribeTaskWS(conn, msg.TaskID); err == nil {
				subscribed[msg.TaskID] = struct{}{}
				_ = h.svc.SendWSMessage(conn, map[string]interface{}{"taskId": msg.TaskID, "action": "subscribed"})
			}
		case "unsubscribe":
			if msg.TaskID == "" {
				continue
			}
			_ = h.svc.UnsubscribeTaskWS(conn, msg.TaskID)
			delete(subscribed, msg.TaskID)
			_ = h.svc.SendWSMessage(conn, map[string]interface{}{"taskId": msg.TaskID, "action": "unsubscribed"})
		}
	}
	// inbound closed => connection closed; cleanup subscriptions
	for tid := range subscribed {
		_ = h.svc.UnsubscribeTaskWS(conn, tid)
	}
}

// ExecuteSQLHandler 处理SQL执行请求 POST /data/execute
// 支持执行各种类型的SQL语句（SELECT、INSERT、UPDATE、DELETE、CREATE、DROP等）
// 支持同步和异步两种执行模式，根据SQL复杂度和数据量自动选择或由客户端指定
// 请求体：api.ExecuteSQLRequest 包含连接信息、SQL语句、参数等
// 响应：统一的执行结果，包含执行状态、影响行数、可能的查询结果等信息
func (h *Handler) ExecuteSQLHandler(c *gin.Context) {
	// 解析请求参数，验证SQL语句和连接信息
	var req api.ExecuteSQLRequest
	if err := c.ShouldBindJSON(&req); err != nil {
		// 参数绑定失败，返回客户端错误
		c.JSON(http.StatusBadRequest, response.Response{Code: 1, Msg: "请求参数错误", Error: err.Error()})
		return
	}

	// 调用服务层执行SQL语句
	// 支持参数化查询防止SQL注入，根据路径定位数据库对象
	isAsyncSQL := req.Async.Mode != api.AsyncModeSync
	result, err := h.svc.ExecuteSQL(c.Request.Context(), req.ConnID, req.ConnInfo, req.Path, req.SQL, req.Params, isAsyncSQL, req.IncludeLarge)
	if err != nil {
		// SQL执行失败，返回服务器内部错误
		c.JSON(http.StatusInternalServerError, response.Response{Code: 2, Msg: "执行SQL失败", Error: err.Error()})
		return
	}

	// 根据执行模式处理不同的返回结果
	// 异步模式返回任务ID，同步模式返回执行结果
	if taskId, ok := result.(string); ok {
		// 异步执行：SQL已提交到后台队列，返回任务ID供后续查询执行状态
		asyncResp := api.AsyncExecuteResponse{
			Success: true,
			TaskID:  taskId,
			Message: "SQL已提交到任务队列",
		}
		if req.Async.Mode == api.AsyncModeWS {
			scheme := "ws"
			if c.Request.TLS != nil {
				scheme = "wss"
			}
			asyncResp.WsURL = fmt.Sprintf("%s://%s/database/data/tasks/ws", scheme, c.Request.Host)
		}
		c.JSON(http.StatusOK, response.Response{Code: 0, Msg: "ok", Data: asyncResp})
		return
	}

	// 同步执行：直接返回SQL执行结果
	if executeResult, ok := result.(meta.ExecuteResult); ok {
		// 将内部的ExecuteResult转换为统一的DataOperationResponse格式
		// 便于前端统一处理不同类型的数据库操作结果
		resp := api.DataOperationResponse{
			Success:       executeResult.Success,
			SQLType:       executeResult.SQLType,       // SQL语句类型（SELECT/INSERT/UPDATE等）
			AffectedRows:  executeResult.AffectedRows,  // 受影响的行数
			LastInsertID:  executeResult.LastInsertID,  // 自增主键值（INSERT语句）
			ErrorMessage:  executeResult.ErrorMessage,  // 执行错误信息
			ExecutionTime: executeResult.ExecutionTime, // 执行耗时（毫秒）
			RowsReturned:  executeResult.RowsReturned,  // 返回的行数（SELECT语句）
			Message:       "SQL执行完成",
		}
		// 只有SELECT语句才包含查询结果数据，其他语句（如INSERT/UPDATE）没有Data字段
		if executeResult.Data != nil {
			resp.Data = *executeResult.Data
		}
		c.JSON(http.StatusOK, response.Response{Code: 0, Msg: "ok", Data: resp})
		return
	}

	// 理论上不会到达这里，但为了代码健壮性保留错误处理
	c.JSON(http.StatusInternalServerError, response.Response{Code: 2, Msg: "未知的返回类型", Error: "service returned unexpected type"})
}

// GetExecuteTaskResultHandler 处理异步SQL执行任务结果获取请求 GET /data/execute/task/result
// 用于查询异步SQL执行任务的执行状态和结果
// 当任务完成时返回SQL执行结果（包括影响行数、可能的查询数据等）
// 当任务失败时返回错误信息，当任务进行中时只返回状态信息
// 请求体：api.GetExecuteTaskResultRequest 包含任务ID
// 响应：包含任务状态和可能的SQL执行结果数据
func (h *Handler) GetExecuteTaskResultHandler(c *gin.Context) {
	// 解析请求参数，获取要查询的任务ID
	var req api.GetExecuteTaskResultRequest
	if err := c.ShouldBindJSON(&req); err != nil {
		// 参数验证失败，返回客户端错误
		c.JSON(http.StatusBadRequest, response.Response{Code: 1, Msg: "请求参数错误", Error: err.Error()})
		return
	}

	// 调用服务层查询指定SQL执行任务的结果和状态
	// 返回执行结果数据、任务状态和可能的错误信息
	result, status, err := h.svc.GetExecuteTaskResult(c.Request.Context(), req.TaskID)
	if err != nil {
		// 获取任务结果失败，返回服务器内部错误
		c.JSON(http.StatusInternalServerError, response.Response{Code: 2, Msg: "获取任务结果失败", Error: err.Error()})
		return
	}

	// 构建基础响应结构，包含任务状态信息
	resp := api.GetExecuteTaskResultResponse{
		Success: true,
		Status:  status,
		Message: "获取任务结果成功",
	}

	// 如果任务已完成，填充SQL执行结果数据
	// 包含执行成功状态、影响行数、执行时间等详细信息
	if status == "completed" {
		// 将内部的ExecuteResult转换为统一的DataOperationResponse格式
		// 便于前端统一处理不同类型的数据库操作结果
		resp.Data = api.DataOperationResponse{
			Success:       result.Success,       // SQL执行是否成功
			SQLType:       result.SQLType,       // SQL语句类型（SELECT/INSERT/UPDATE等）
			AffectedRows:  result.AffectedRows,  // 受影响的行数
			LastInsertID:  result.LastInsertID,  // 自增主键值（INSERT语句）
			ErrorMessage:  result.ErrorMessage,  // 执行错误信息
			ExecutionTime: result.ExecutionTime, // 执行耗时（毫秒）
			RowsReturned:  result.RowsReturned,  // 返回的行数（SELECT语句）
		}
		// 只有SELECT语句才包含查询结果数据，其他语句（如INSERT/UPDATE）没有Data字段
		if result.Data != nil {
			resp.Data.Data = *result.Data
		}
	}

	// 返回任务状态和可能的SQL执行结果
	c.JSON(http.StatusOK, response.Response{Code: 0, Msg: "ok", Data: resp})
}