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

package dbs

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

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

func init() {
	mcp.Register("get_mysql_columns", "根据表名称获取MySQL表的所有字段相关信息",
		map[string]interface{}{
			"type": "object",
			"properties": map[string]interface{}{
				"table_name": map[string]interface{}{
					"type":        "string",
					"description": "表名称",
				},
				"schema": map[string]interface{}{
					"type":        "string",
					"description": "数据库名称（默认为当前数据库）",
					"default":     "",
				},
				"include_comments": map[string]interface{}{
					"type":        "boolean",
					"description": "是否包含字段注释",
					"default":     true,
				},
				"database_key": map[string]interface{}{
					"type":        "string",
					"description": "数据库配置键名：warehouse（仓库数据库）或 business（业务数据库），可选，默认使用主数据库",
					"enum":        []string{"warehouse", "business"},
					"default":     "",
				},
			},
			"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
			}

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

			// 获取当前数据库名称
			currentDatabase := dbFactory.GetDatabaseName()
			schema := strings.TrimSpace(params.Schema)
			if schema == "" {
				schema = currentDatabase
			}

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

			// 构建查询SQL
			var query string
			if params.IncludeComments {
				query = `
					SELECT 
						ORDINAL_POSITION as column_id,
						COLUMN_NAME as column_name,
						COLUMN_TYPE as data_type,
						IS_NULLABLE as nullable,
						COLUMN_DEFAULT as column_default,
						COLUMN_COMMENT as column_comment,
						CHARACTER_MAXIMUM_LENGTH as character_maximum_length,
						NUMERIC_PRECISION as numeric_precision,
						NUMERIC_SCALE as numeric_scale,
						DATA_TYPE as simple_data_type,
						COLUMN_KEY as column_key,
						EXTRA as extra
					FROM information_schema.COLUMNS
					WHERE TABLE_SCHEMA = ?
					AND TABLE_NAME = ?
					ORDER BY ORDINAL_POSITION`
			} else {
				query = `
					SELECT 
						ORDINAL_POSITION as column_id,
						COLUMN_NAME as column_name,
						COLUMN_TYPE as data_type,
						IS_NULLABLE as nullable,
						COLUMN_DEFAULT as column_default,
						CHARACTER_MAXIMUM_LENGTH as character_maximum_length,
						NUMERIC_PRECISION as numeric_precision,
						NUMERIC_SCALE as numeric_scale,
						DATA_TYPE as simple_data_type,
						COLUMN_KEY as column_key,
						EXTRA as extra
					FROM information_schema.COLUMNS
					WHERE TABLE_SCHEMA = ?
					AND TABLE_NAME = ?
					ORDER BY 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 = ? AND TABLE_NAME = ?`
				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"
				}

				// 处理注释字段
				if params.IncludeComments {
					if comment, ok := results[i]["column_comment"]; ok && comment == nil {
						results[i]["column_comment"] = ""
					}
				}

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

				// 确定是否为键
				if columnKey, ok := results[i]["column_key"].(string); ok {
					results[i]["is_primary_key"] = columnKey == "PRI"
					results[i]["is_unique_key"] = columnKey == "UNI"
					delete(results[i], "column_key")
				}

				// 构建完整数据类型
				if dataType, ok := results[i]["data_type"].(string); ok {
					results[i]["full_data_type"] = dataType
				}
			}

			// 获取表注释
			tableCommentQuery := `SELECT TABLE_COMMENT as table_comment FROM information_schema.TABLES WHERE TABLE_SCHEMA = ? AND TABLE_NAME = ?`
			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),
			}, nil
		},
	)
}