go-svc-mcp/internal/tools/dbs/get_sqlserver_data.go

package dbs

import (
	"encoding/json"
	"fmt"
	"strings"
	"time"

	"git.x2erp.com/qdy/go-svc-mcp/internal/mcp"
)

func init() {
	mcp.Register("get_sqlserver_data", "通过表名称获取SQL Server表数据：只能接收表名称和分页查询参数分页参数数据获取数据不能超过10条",
		map[string]interface{}{
			"type": "object",
			"properties": map[string]interface{}{
				"table_name": map[string]interface{}{
					"type":        "string",
					"description": "表名称",
				},
				"schema": map[string]interface{}{
					"type":        "string",
					"description": "模式名称（默认为dbo）",
					"default":     "",
				},
				"page": map[string]interface{}{
					"type":        "integer",
					"description": "页码（从1开始）",
					"default":     1,
					"minimum":     1,
				},
				"page_size": map[string]interface{}{
					"type":        "integer",
					"description": "每页记录数（最大10条）",
					"default":     10,
					"minimum":     1,
					"maximum":     10,
				},
				"order_by": map[string]interface{}{
					"type":        "string",
					"description": "排序字段（例如：column1 ASC, column2 DESC）",
					"default":     "",
				},
				"where_clause": map[string]interface{}{
					"type":        "string",
					"description": "WHERE条件（例如：status = 'ACTIVE'）",
					"default":     "",
				},
				"database_key": map[string]interface{}{
					"type":        "string",
					"description": "数据库配置键名（如：business），可选，默认使用主数据库",
					"enum":        []string{"warehouse", "business"},
					"default":     "warehouse",
				},
			},
			"required": []string{"table_name"},
		},
		func(input json.RawMessage, deps *mcp.ToolDependencies) (interface{}, error) {
			var params struct {
				TableName   string `json:"table_name"`
				Schema      string `json:"schema"`
				Page        int    `json:"page"`
				PageSize    int    `json:"page_size"`
				OrderBy     string `json:"order_by"`
				WhereClause string `json:"where_clause"`
				DatabaseKey string `json:"database_key"`
			}

			if len(input) > 0 {
				if err := json.Unmarshal(input, &params); err != nil {
					return nil, err
				}
			}

			// 设置默认值
			if params.Page == 0 {
				params.Page = 1
			}
			if params.PageSize == 0 {
				params.PageSize = 10
			}

			// 确保每页记录数不超过10条
			if params.PageSize > 10 {
				params.PageSize = 10
			}

			// 获取数据库工厂
			dbFactory, err := GetDBFactory(params.DatabaseKey, deps)
			if err != nil {
				return nil, err
			}

			// 获取数据库类型，确保是SQL Server
			dbType := dbFactory.GetDBType()
			if dbType != "sqlserver" {
				return nil, fmt.Errorf("当前数据库类型为 %s，此工具仅支持SQL Server数据库", dbType)
			}

			// 设置默认模式
			schema := strings.TrimSpace(params.Schema)
			if schema == "" {
				schema = "dbo"
			}

			tableName := strings.TrimSpace(params.TableName)
			if tableName == "" {
				return nil, fmt.Errorf("表名称不能为空")
			}

			// 验证表是否存在
			tableExistsQuery := `SELECT COUNT(*) as table_count FROM INFORMATION_SCHEMA.TABLES WHERE TABLE_SCHEMA = @p1 AND TABLE_NAME = @p2`
			tableCheckResults, err := dbFactory.QuerySliceMapWithParams(tableExistsQuery, schema, tableName)
			if err != nil {
				return nil, fmt.Errorf("检查表是否存在失败: %v", err)
			}

			tableExists := false
			if len(tableCheckResults) > 0 {
				if count, ok := tableCheckResults[0]["table_count"].(int64); ok && count > 0 {
					tableExists = true
				}
			}

			if !tableExists {
				return nil, fmt.Errorf("表 '%s' 不存在于模式 '%s' 中", tableName, schema)
			}

			// 构建带模式前缀的表名
			fullTableName := fmt.Sprintf("[%s].[%s]", schema, tableName)

			// 构建WHERE子句
			whereClause := ""
			if params.WhereClause != "" {
				whereClause = fmt.Sprintf("WHERE %s", params.WhereClause)
			}

			// 构建ORDER BY子句
			orderByClause := ""
			if params.OrderBy != "" {
				orderByClause = fmt.Sprintf("ORDER BY %s", params.OrderBy)
			}

			// 计算分页偏移量
			offset := (params.Page - 1) * params.PageSize

			// 先获取总记录数
			countQuery := fmt.Sprintf("SELECT COUNT(*) as total_count FROM %s %s", fullTableName, whereClause)
			countResults, err := dbFactory.QuerySliceMap(countQuery)
			if err != nil {
				return nil, fmt.Errorf("查询总记录数失败: %v", err)
			}

			totalCount := int64(0)
			if len(countResults) > 0 {
				if count, ok := countResults[0]["total_count"].(int64); ok {
					totalCount = count
				}
			}

			// 计算总页数
			totalPages := 0
			if totalCount > 0 {
				totalPages = int((totalCount + int64(params.PageSize) - 1) / int64(params.PageSize))
			}

			// 如果请求的页码大于总页数，返回空结果
			if params.Page > totalPages && totalPages > 0 {
				return map[string]interface{}{
					"tenant_id":       deps.ReqCtx.TenantID,
					"user_id":         deps.ReqCtx.UserID,
					"database_type":   dbType,
					"database_name":   dbFactory.GetDatabaseName(),
					"schema":          schema,
					"table_name":      tableName,
					"full_table_name": fullTableName,
					"page":            params.Page,
					"page_size":       params.PageSize,
					"total_count":     totalCount,
					"total_pages":     totalPages,
					"where_clause":    params.WhereClause,
					"order_by":        params.OrderBy,
					"data":            []map[string]interface{}{},
					"data_count":      0,
					"has_more":        false,
					"timestamp":       time.Now().Format(time.RFC3339),
				}, nil
			}

			// 构建数据查询（SQL Server分页查询 - 使用OFFSET FETCH，需要SQL Server 2012+）
			dataQuery := ""
			if orderByClause == "" {
				// 如果没有指定ORDER BY，需要指定一个默认排序
				orderByClause = "ORDER BY (SELECT NULL)"
			}

			dataQuery = fmt.Sprintf(`
				SELECT * FROM %s 
				%s 
				%s 
				OFFSET %d ROWS 
				FETCH NEXT %d ROWS ONLY`,
				fullTableName, whereClause, orderByClause, offset, params.PageSize)

			// 执行数据查询
			dataResults, err := dbFactory.QuerySliceMap(dataQuery)
			if err != nil {
				// 如果OFFSET FETCH失败，可能是旧版本SQL Server，使用ROW_NUMBER作为备选方案
				fallbackQuery := fmt.Sprintf(`
					SELECT * FROM (
						SELECT *, ROW_NUMBER() OVER (%s) as row_num 
						FROM %s 
						%s
					) as numbered 
					WHERE row_num > %d AND row_num <= %d`,
					orderByClause, fullTableName, whereClause, offset, offset+params.PageSize)

				dataResults, err = dbFactory.QuerySliceMap(fallbackQuery)
				if err != nil {
					return nil, fmt.Errorf("查询表数据失败: %v", err)
				}

				// 移除row_num字段
				for i := range dataResults {
					delete(dataResults[i], "row_num")
				}
			}

			// 检查是否有更多数据
			hasMore := (int64(offset+params.PageSize) < totalCount)

			return map[string]interface{}{
				"tenant_id":       deps.ReqCtx.TenantID,
				"user_id":         deps.ReqCtx.UserID,
				"database_type":   dbType,
				"database_name":   dbFactory.GetDatabaseName(),
				"schema":          schema,
				"table_name":      tableName,
				"full_table_name": fullTableName,
				"page":            params.Page,
				"page_size":       params.PageSize,
				"total_count":     totalCount,
				"total_pages":     totalPages,
				"where_clause":    params.WhereClause,
				"order_by":        params.OrderBy,
				"data":            dataResults,
				"data_count":      len(dataResults),
				"has_more":        hasMore,
				"timestamp":       time.Now().Format(time.RFC3339),
				"note":            "数据获取限制：最多返回10条记录",
			}, nil
		},
	)
}