Repository: dgraph-io/dgraph
File: dgraph/cmd/mcp/mcp_server.go (Go)

Code before:
			err := txn.Discard(ctx)
			if err != nil {
				glog.Errorf("failed to discard transaction: %v", err)
			}
		}()
		args := request.GetArguments()
		if args == nil {
			return mcp.NewToolResultError("Query must be present"), nil
		}
		queryArg, ok := args["query"]
		if !ok || queryArg == nil {
			return mcp.NewToolResultError("Query must be present"), nil
		}
		op, ok := queryArg.(string)
		if !ok {
			return mcp.NewToolResultError("Query must be a string"), nil
		}
		// variable handling
		// if variables arg is not present, use an empty map
		// if present, parse the JSON string parameter 'variables'
		// and pass the var map to QueryWithVars
		// return an error if the arg 'variables' is not a valid JSON
		vars := make(map[string]string)
		variablesArg, ok := args["variables"]
		if ok && variablesArg != nil {
			variables, _ := variablesArg.(string)
			// create a map of variables from JSON string
			if err := json.Unmarshal([]byte(variables), &vars); err != nil {
				return mcp.NewToolResultErrorFromErr("Error parsing variables", err), nil
			}
		}
		resp, err := txn.QueryWithVars(ctx, op, vars)
		if err != nil {
			return mcp.NewToolResultErrorFromErr("Error running query", err), nil
		}
		return mcp.NewToolResultText(string(resp.GetJson())), nil
	})

	s.AddTool(schemaTool, func(ctx context.Context, request mcp.CallToolRequest) (*mcp.CallToolResult, error) {
		conn, err := getConn(connectionString)
		if err != nil {
			return mcp.NewToolResultErrorFromErr("Error opening connection with Dgraph Alpha", err), nil
		}
		txn := conn.NewTxn()
		defer func() {
			err := txn.Discard(ctx)
			if err != nil {
				glog.Errorf("failed to discard transaction: %v", err)
			}
		}()
		resp, err := txn.Query(ctx, "schema {}")

Review comment (bug):
The type assertion ignores failures. If `variablesArg` isn't a string, this will silently produce an empty string. Use a two-value assertion (e.g., `variables, ok := variablesArg.(string)`) and return an error if `ok` is false.
```suggestion
			variables, ok := variablesArg.(string)
			if !ok {
				return mcp.NewToolResultError("Variables must be a string"), nil
			}
```

Code after:
		defer func() {
			err := txn.Discard(ctx)
			if err != nil {
				glog.Errorf("failed to discard transaction: %v", err)
			}
		}()
		args := request.GetArguments()
		if args == nil {
			return mcp.NewToolResultError("Query must be present"), nil
		}
		queryArg, ok := args["query"]
		if !ok || queryArg == nil {
			return mcp.NewToolResultError("Query must be present"), nil
		}
		op, ok := queryArg.(string)
		if !ok {
			return mcp.NewToolResultError("Query must be a string"), nil
		}
		vars := make(map[string]any)
		variablesArg, ok := args["variables"]
		if ok && variablesArg != nil {
			variables, ok := variablesArg.(string)
			if !ok {
				return mcp.NewToolResultError("Variables must be a JSON-formatted string"), nil
			}
			// create a map of variables from JSON string
			if err := json.Unmarshal([]byte(variables), &vars); err != nil {
				return mcp.NewToolResultErrorFromErr("Error parsing variables", err), nil
			}
		}
		// convert vars map[string]any to map[string]string as required by txn.QueryWithVars
		varsString := make(map[string]string)
		for k, v := range vars {
			switch val := v.(type) {
			case string:
				varsString[k] = val
			case float64:
				varsString[k] = strconv.FormatFloat(val, 'f', -1, 64)
			case bool:
				varsString[k] = strconv.FormatBool(val)
			case nil:
				varsString[k] = "null"
			default:
				return mcp.NewToolResultError(fmt.Sprintf("could not convert complex variable %q to string", k)), nil
			}
		}
		resp, err := txn.QueryWithVars(ctx, op, varsString)
		if err != nil {
			return mcp.NewToolResultErrorFromErr("Error running query", err), nil
		}
		return mcp.NewToolResultText(string(resp.GetJson())), nil