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

Code before:
		}
		txn := conn.NewTxn()
		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
		}
		// 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)

Review comment (bug):
Unmarshaling JSON into `map[string]string` will fail if any values aren't strings. If you need to support numbers or booleans, consider using `map[string]interface{}` or a custom struct.
```suggestion
		vars := make(map[string]interface{})
```

Code after:
			return mcp.NewToolResultErrorFromErr("Error opening connection with Dgraph Alpha", err), nil
		}
		txn := conn.NewReadOnlyTxn()
		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 {