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

package dbs

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

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

func init() {
	mcp.Register("get_sqlserver_columns", "根据表名称获取SQL Server表的所有字段相关信息",
		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":     "",
				},
				"include_comments": map[string]interface{}{
					"type":        "boolean",
					"description": "是否包含字段注释",
					"default":     true,
				},
				"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"`
				IncludeComments bool   `json:"include_comments"`
				DatabaseKey     string `json:"database_key"`
			}

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

			// 获取数据库工厂
			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("表名称不能为空")
			}

			// 构建查询SQL
			var query string
			if params.IncludeComments {
				query = `
					SELECT 
						c.ORDINAL_POSITION as column_id,
						c.COLUMN_NAME as column_name,
						c.DATA_TYPE as data_type,
						c.IS_NULLABLE as nullable,
						c.COLUMN_DEFAULT as column_default,
						c.CHARACTER_MAXIMUM_LENGTH,
						c.NUMERIC_PRECISION,
						c.NUMERIC_SCALE,
						CAST(
							CASE WHEN pk.COLUMN_NAME IS NOT NULL THEN 1 ELSE 0 END 
						AS BIT) as is_primary_key,
						'' as column_comment  -- 先占位，后面单独查询
					FROM INFORMATION_SCHEMA.COLUMNS c
					LEFT JOIN (
						SELECT 
							ku.COLUMN_NAME,
							tc.TABLE_SCHEMA,
							tc.TABLE_NAME
						FROM INFORMATION_SCHEMA.TABLE_CONSTRAINTS tc
						JOIN INFORMATION_SCHEMA.KEY_COLUMN_USAGE ku 
							ON tc.CONSTRAINT_NAME = ku.CONSTRAINT_NAME 
							AND tc.TABLE_SCHEMA = ku.TABLE_SCHEMA
						WHERE tc.CONSTRAINT_TYPE = 'PRIMARY KEY'
					) pk ON c.TABLE_SCHEMA = pk.TABLE_SCHEMA 
						AND c.TABLE_NAME = pk.TABLE_NAME 
						AND c.COLUMN_NAME = pk.COLUMN_NAME
					WHERE c.TABLE_SCHEMA = @p1
					AND c.TABLE_NAME = @p2
					ORDER BY c.ORDINAL_POSITION`
			} else {
				query = `
					SELECT 
						c.ORDINAL_POSITION as column_id,
						c.COLUMN_NAME as column_name,
						c.DATA_TYPE as data_type,
						c.IS_NULLABLE as nullable,
						c.COLUMN_DEFAULT as column_default,
						c.CHARACTER_MAXIMUM_LENGTH,
						c.NUMERIC_PRECISION,
						c.NUMERIC_SCALE,
						CAST(
							CASE WHEN pk.COLUMN_NAME IS NOT NULL THEN 1 ELSE 0 END 
						AS BIT) as is_primary_key
					FROM INFORMATION_SCHEMA.COLUMNS c
					LEFT JOIN (
						SELECT 
							ku.COLUMN_NAME,
							tc.TABLE_SCHEMA,
							tc.TABLE_NAME
						FROM INFORMATION_SCHEMA.TABLE_CONSTRAINTS tc
						JOIN INFORMATION_SCHEMA.KEY_COLUMN_USAGE ku 
							ON tc.CONSTRAINT_NAME = ku.CONSTRAINT_NAME 
							AND tc.TABLE_SCHEMA = ku.TABLE_SCHEMA
						WHERE tc.CONSTRAINT_TYPE = 'PRIMARY KEY'
					) pk ON c.TABLE_SCHEMA = pk.TABLE_SCHEMA 
						AND c.TABLE_NAME = pk.TABLE_NAME 
						AND c.COLUMN_NAME = pk.COLUMN_NAME
					WHERE c.TABLE_SCHEMA = @p1
					AND c.TABLE_NAME = @p2
					ORDER BY c.ORDINAL_POSITION`
			}

			// 执行查询
			results, err := dbFactory.QuerySliceMapWithParams(query, schema, tableName)
			if err != nil {
				return nil, fmt.Errorf("查询表字段信息失败: %v", err)
			}

			if len(results) == 0 {
				// 尝试查询表是否存在
				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 && len(tableCheckResults) > 0 {
					if count, ok := tableCheckResults[0]["table_count"].(int64); ok && count == 0 {
						return nil, fmt.Errorf("表 '%s' 不存在于模式 '%s' 中", tableName, schema)
					}
				}

				return nil, fmt.Errorf("表 '%s' 存在但没有字段信息或表为空", tableName)
			}

			// 处理字段信息
			for i := range results {
				// 处理nullable字段
				if nullable, ok := results[i]["nullable"].(string); ok {
					results[i]["is_nullable"] = nullable == "YES"
					delete(results[i], "nullable")
				}

				// 处理注释字段
				if params.IncludeComments {
					// 查询扩展属性获取字段注释
					columnName := results[i]["column_name"].(string)
					descQuery := `
						SELECT value as column_comment
						FROM fn_listextendedproperty ('MS_Description', 'SCHEMA', @p1, 'TABLE', @p2, 'COLUMN', @p3)`
					descResults, err := dbFactory.QuerySliceMapWithParams(descQuery, schema, tableName, columnName)
					if err == nil && len(descResults) > 0 {
						if desc, ok := descResults[0]["column_comment"].(string); ok && desc != "" {
							results[i]["column_comment"] = desc
						} else {
							results[i]["column_comment"] = ""
						}
					} else {
						results[i]["column_comment"] = ""
					}
				}

				// 处理默认值
				if defaultValue, ok := results[i]["column_default"]; ok && defaultValue == nil {
					results[i]["column_default"] = ""
				}

				// 处理主键字段
				if isPK, ok := results[i]["is_primary_key"].(bool); ok {
					results[i]["is_primary_key"] = isPK
				}

				// 构建完整数据类型
				if dataType, ok := results[i]["data_type"].(string); ok {
					fullType := dataType
					if length, ok := results[i]["character_maximum_length"].(int64); ok && length > 0 && length != 2147483647 {
						if length == -1 {
							fullType = fmt.Sprintf("%s(max)", dataType)
						} else {
							fullType = fmt.Sprintf("%s(%d)", dataType, length)
						}
					} else if precision, ok := results[i]["numeric_precision"].(int64); ok && precision > 0 {
						if scale, ok := results[i]["numeric_scale"].(int64); ok && scale > 0 {
							fullType = fmt.Sprintf("%s(%d,%d)", dataType, precision, scale)
						} else {
							fullType = fmt.Sprintf("%s(%d)", dataType, precision)
						}
					}
					results[i]["full_data_type"] = fullType
				}
			}

			// 获取表注释
			tableCommentQuery := `
				SELECT value as table_comment
				FROM fn_listextendedproperty ('MS_Description', 'SCHEMA', @p1, 'TABLE', @p2, NULL, NULL)`
			commentResults, err := dbFactory.QuerySliceMapWithParams(tableCommentQuery, schema, tableName)
			tableComment := ""
			if err == nil && len(commentResults) > 0 {
				if comment, ok := commentResults[0]["table_comment"].(string); ok {
					tableComment = comment
				}
			}

			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,
				"table_comment":    tableComment,
				"include_comments": params.IncludeComments,
				"columns":          results,
				"total_columns":    len(results),
				"timestamp":        time.Now().Format(time.RFC3339),
				"note":             "SQL Server字段注释通过扩展属性MS_Description设置",
			}, nil
		},
	)
}