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

package dbs

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

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

func init() {
	mcp.Register("get_postgresql_tables", "获取PostgreSQL数据库中的所有表和描述",
		map[string]interface{}{
			"type": "object",
			"properties": map[string]interface{}{
				"include_system_tables": map[string]interface{}{
					"type":        "boolean",
					"description": "是否包含系统表",
					"default":     false,
				},
				"schema": map[string]interface{}{
					"type":        "string",
					"description": "模式名称（默认为public）",
					"default":     "",
				},
				"database_key": map[string]interface{}{
					"type":        "string",
					"description": "数据库配置键名（如：business），可选，默认使用主数据库",
					"enum":        []string{"warehouse", "business"},
					"default":     "warehouse",
				},
			},
			"required": []string{},
		},
		func(input json.RawMessage, deps *mcp.ToolDependencies) (interface{}, error) {
			var params struct {
				IncludeSystemTables bool   `json:"include_system_tables"`
				Schema              string `json:"schema"`
				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
			}

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

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

			// 构建查询SQL
			var query string
			if params.IncludeSystemTables {
				query = `
					SELECT 
						table_schema as schema_name,
						table_name,
						obj_description(pgc.oid, 'pg_class') as table_description
					FROM information_schema.tables t
					JOIN pg_catalog.pg_class pgc ON pgc.relname = t.table_name
					JOIN pg_catalog.pg_namespace pgn ON pgn.oid = pgc.relnamespace AND pgn.nspname = t.table_schema
					WHERE table_type = 'BASE TABLE'
					AND table_schema = $1
					ORDER BY table_name`
			} else {
				query = `
					SELECT 
						table_schema as schema_name,
						table_name,
						obj_description(pgc.oid, 'pg_class') as table_description
					FROM information_schema.tables t
					JOIN pg_catalog.pg_class pgc ON pgc.relname = t.table_name
					JOIN pg_catalog.pg_namespace pgn ON pgn.oid = pgc.relnamespace AND pgn.nspname = t.table_schema
					WHERE table_type = 'BASE TABLE'
					AND table_schema = $1
					AND table_schema NOT IN ('pg_catalog', 'information_schema')
					ORDER BY table_name`
			}

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

			// 如果没有描述信息，设置为空字符串
			for i := range results {
				if results[i]["table_description"] == nil {
					results[i]["table_description"] = ""
				}
			}

			return map[string]interface{}{
				"tenant_id":             deps.ReqCtx.TenantID,
				"user_id":               deps.ReqCtx.UserID,
				"database_type":         dbType,
				"database_name":         dbFactory.GetDatabaseName(),
				"schema":                schema,
				"include_system_tables": params.IncludeSystemTables,
				"tables":                results,
				"total_tables":          len(results),
				"timestamp":             time.Now().Format(time.RFC3339),
			}, nil
		},
	)
}