CreateRecordTool.cs

// Copyright (c) Microsoft Corporation.
// Licensed under the MIT License.

using System.Text.Json;
using Azure.DataApiBuilder.Auth;
using Azure.DataApiBuilder.Config.DatabasePrimitives;
using Azure.DataApiBuilder.Config.ObjectModel;
using Azure.DataApiBuilder.Core.Configurations;
using Azure.DataApiBuilder.Core.Models;
using Azure.DataApiBuilder.Core.Resolvers;
using Azure.DataApiBuilder.Core.Resolvers.Factories;
using Azure.DataApiBuilder.Core.Services;
using Azure.DataApiBuilder.Mcp.Model;
using Azure.DataApiBuilder.Mcp.Utils;
using Microsoft.AspNetCore.Http;
using Microsoft.AspNetCore.Mvc;
using Microsoft.Extensions.DependencyInjection;
using Microsoft.Extensions.Logging;
using ModelContextProtocol.Protocol;
using static Azure.DataApiBuilder.Mcp.Model.McpEnums;

namespace Azure.DataApiBuilder.Mcp.BuiltInTools
{
    public class CreateRecordTool : IMcpTool
    {
        public ToolType ToolType { get; } = ToolType.BuiltIn;

        public bool IsEnabled(RuntimeConfig config) => config.McpDmlTools?.CreateRecord ?? true;

        public Tool GetToolMetadata()
        {
            return new Tool
            {
                Name = "create_record",
                Description = "STEP 1: describe_entities -> find entities with CREATE permission and their fields. STEP 2: call this tool with matching field names and values.",
                InputSchema = JsonSerializer.Deserialize<JsonElement>(
                    @"{
                        ""type"": ""object"",
                        ""properties"": {
                            ""entity"": {
                                ""type"": ""string"",
                                ""description"": ""Entity name with CREATE permission.""
                            },
                            ""data"": {
                                ""type"": ""object"",
                                ""description"": ""Required fields and values for the new record.""
                            }
                        },
                        ""required"": [""entity"", ""data""]
                    }"
                )
            };
        }

        public async Task<CallToolResult> ExecuteAsync(
            JsonDocument? arguments,
            IServiceProvider serviceProvider,
            CancellationToken cancellationToken = default)
        {
            ILogger<CreateRecordTool>? logger = serviceProvider.GetService<ILogger<CreateRecordTool>>();
            string toolName = GetToolMetadata().Name;
            if (arguments == null)
            {
                return McpResponseBuilder.BuildErrorResult(toolName, "InvalidArguments", "No arguments provided.", logger);
            }

            RuntimeConfigProvider runtimeConfigProvider = serviceProvider.GetRequiredService<RuntimeConfigProvider>();
            if (!runtimeConfigProvider.TryGetConfig(out RuntimeConfig? runtimeConfig))
            {
                return McpResponseBuilder.BuildErrorResult(toolName, "InvalidConfiguration", "Runtime configuration not available.", logger);
            }

            if (runtimeConfig.McpDmlTools?.CreateRecord != true)
            {
                return McpErrorHelpers.ToolDisabled(toolName, logger);
            }

            try
            {
                cancellationToken.ThrowIfCancellationRequested();
                JsonElement root = arguments.RootElement;

                if (!McpArgumentParser.TryParseEntityAndData(root, out string entityName, out JsonElement dataElement, out string parseError))
                {
                    return McpResponseBuilder.BuildErrorResult(toolName, "InvalidArguments", parseError, logger);
                }

                // Check entity-level DML tool configuration
                if (runtimeConfig.Entities?.TryGetValue(entityName, out Entity? entity) == true &&
                    entity.Mcp?.DmlToolEnabled == false)
                {
                    return McpErrorHelpers.ToolDisabled(toolName, logger, $"DML tools are disabled for entity '{entityName}'.");
                }

                if (!McpMetadataHelper.TryResolveMetadata(
                        entityName,
                        runtimeConfig,
                        serviceProvider,
                        out ISqlMetadataProvider sqlMetadataProvider,
                        out DatabaseObject dbObject,
                        out string dataSourceName,
                        out string metadataError))
                {
                    return McpResponseBuilder.BuildErrorResult(toolName, "EntityNotFound", metadataError, logger);
                }

                // Create an HTTP context for authorization
                IHttpContextAccessor httpContextAccessor = serviceProvider.GetRequiredService<IHttpContextAccessor>();
                HttpContext httpContext = httpContextAccessor.HttpContext ?? new DefaultHttpContext();
                IAuthorizationResolver authorizationResolver = serviceProvider.GetRequiredService<IAuthorizationResolver>();

                if (!McpAuthorizationHelper.ValidateRoleContext(httpContext, authorizationResolver, out string roleCtxError))
                {
                    return McpErrorHelpers.PermissionDenied(toolName, entityName, "create", roleCtxError, logger);
                }

                if (!McpAuthorizationHelper.TryResolveAuthorizedRole(
                        httpContext,
                        authorizationResolver,
                        entityName,
                        EntityActionOperation.Create,
                        out string? effectiveRole,
                        out string authError))
                {
                    return McpErrorHelpers.PermissionDenied(toolName, entityName, "create", authError, logger);
                }

                // Column-level authorization: ensure the caller's effective role is permitted to write
                // every column present in the request payload (fields.include/fields.exclude enforcement).
                IEnumerable<string> requestedColumns = dataElement.ValueKind == JsonValueKind.Object
                    ? dataElement.EnumerateObject().Select(property => property.Name)
                    : Enumerable.Empty<string>();

                try
                {
                    if (!McpAuthorizationHelper.AreColumnsAuthorizedForOperation(
                            authorizationResolver,
                            entityName,
                            effectiveRole!,
                            EntityActionOperation.Create,
                            requestedColumns,
                            out string columnAuthError))
                    {
                        return McpErrorHelpers.PermissionDenied(toolName, entityName, "create", columnAuthError, logger);
                    }
                }
                catch (Azure.DataApiBuilder.Service.Exceptions.DataApiBuilderException dabEx)
                {
                    return McpResponseBuilder.BuildErrorResult(toolName, "ValidationFailed", $"Request validation failed: {dabEx.Message}", logger);
                }

                JsonElement insertPayloadRoot = dataElement.Clone();

                // Validate it's a table or view - stored procedures use execute_entity
                if (dbObject.SourceType != EntitySourceType.Table && dbObject.SourceType != EntitySourceType.View)
                {
                    return McpResponseBuilder.BuildErrorResult(toolName, "InvalidEntity", $"Entity '{entityName}' is not a table or view. For stored procedures, use the execute_entity tool instead.", logger);
                }

                InsertRequestContext insertRequestContext = new(
                    entityName,
                    dbObject,
                    insertPayloadRoot,
                    EntityActionOperation.Insert);

                // Only validate tables. For views, skip validation and let the database handle any errors.
                if (dbObject.SourceType is EntitySourceType.Table)
                {
                    RequestValidator requestValidator = serviceProvider.GetRequiredService<RequestValidator>();
                    try
                    {
                        requestValidator.ValidateInsertRequestContext(insertRequestContext);
                    }
                    catch (Exception ex)
                    {
                        return McpResponseBuilder.BuildErrorResult(toolName, "ValidationFailed", $"Request validation failed: {ex.Message}", logger);
                    }
                }

                IMutationEngineFactory mutationEngineFactory = serviceProvider.GetRequiredService<IMutationEngineFactory>();
                DatabaseType databaseType = sqlMetadataProvider.GetDatabaseType();
                IMutationEngine mutationEngine = mutationEngineFactory.GetMutationEngine(databaseType);

                IActionResult? result = await mutationEngine.ExecuteAsync(insertRequestContext);

                if (result is CreatedResult createdResult)
                {
                    return McpResponseBuilder.BuildSuccessResult(
                        new Dictionary<string, object?>
                        {
                            ["entity"] = entityName,
                            ["result"] = createdResult.Value,
                            ["message"] = $"Successfully created record in entity '{entityName}'"
                        },
                        logger,
                        $"Successfully created record in entity '{entityName}'");
                }
                else if (result is ObjectResult objectResult)
                {
                    bool isError = objectResult.StatusCode.HasValue && objectResult.StatusCode.Value >= 400 && objectResult.StatusCode.Value != 403;
                    if (isError)
                    {
                        return McpResponseBuilder.BuildErrorResult(
                            toolName,
                            "CreateFailed",
                            $"Failed to create record in entity '{entityName}'. Error: {JsonSerializer.Serialize(objectResult.Value)}",
                            logger);
                    }
                    else
                    {
                        return McpResponseBuilder.BuildSuccessResult(
                            new Dictionary<string, object?>
                            {
                                ["entity"] = entityName,
                                ["result"] = objectResult.Value,
                                ["message"] = $"Successfully created record in entity '{entityName}'. Unable to perform read-back of inserted records."
                            },
                            logger,
                            $"Successfully created record in entity '{entityName}'. Unable to perform read-back of inserted records.");
                    }
                }
                else
                {
                    if (result is null)
                    {
                        return McpResponseBuilder.BuildErrorResult(
                            toolName,
                            "UnexpectedError",
                            $"Mutation engine returned null result for entity '{entityName}'",
                            logger);
                    }
                    else
                    {
                        return McpResponseBuilder.BuildSuccessResult(
                            new Dictionary<string, object?>
                            {
                                ["entity"] = entityName,
                                ["message"] = $"Create operation completed with unexpected result type: {result.GetType().Name}"
                            },
                            logger,
                            $"Create operation completed for entity '{entityName}' with unexpected result type: {result.GetType().Name}");
                    }
                }
            }
            catch (Exception ex)
            {
                return McpResponseBuilder.BuildErrorResult(toolName, "Error", $"Error: {ex.Message}", logger);
            }
        }
    }
}