HttpUtilities.cs

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

using System;
using System.Collections.Generic;
using System.Data.Common;
using System.Linq;
using System.Net.Http;
using System.Text;
using System.Threading.Tasks;
using Azure.DataApiBuilder.Config.DatabasePrimitives;
using Azure.DataApiBuilder.Config.ObjectModel;
using Azure.DataApiBuilder.Core.Authorization;
using Azure.DataApiBuilder.Core.Configurations;
using Azure.DataApiBuilder.Core.Services;
using Azure.DataApiBuilder.Core.Services.MetadataProviders;
using Azure.DataApiBuilder.Service.GraphQLBuilder;
using Microsoft.Extensions.Logging;

namespace Azure.DataApiBuilder.Service.HealthCheck
{
    /// <summary>
    /// HttpUtilities creates and executes the HTTP requests for HealthEndpoint.
    /// </summary>
    public class HttpUtilities
    {
        private readonly ILogger<HttpUtilities> _logger;
        private HttpClient _httpClient;
        private IMetadataProviderFactory _metadataProviderFactory;
        private RuntimeConfigProvider _runtimeConfigProvider;

        /// <summary>
        /// HttpUtility constructor.
        /// </summary>
        /// <param name="logger">Logger</param>
        /// <param name="metadataProviderFactory">MetadataProviderFactory</param>
        /// <param name="runtimeConfigProvider">RuntimeConfigProvider</param>
        /// <param name="httpClientFactory">HttpClientFactory</param>
        public HttpUtilities(
            ILogger<HttpUtilities> logger,
            IMetadataProviderFactory metadataProviderFactory,
            RuntimeConfigProvider runtimeConfigProvider,
            IHttpClientFactory httpClientFactory)
        {
            _logger = logger;
            _metadataProviderFactory = metadataProviderFactory;
            _runtimeConfigProvider = runtimeConfigProvider;
            _httpClient = httpClientFactory.CreateClient("ContextConfiguredHealthCheckClient");
        }

        // Executes the DB query by establishing a connection to the DB.
        public async Task<string?> ExecuteDbQueryAsync(string query, string connectionString, DbProviderFactory providerFactory, DatabaseType databaseType)
        {
            string? errorMessage = null;
            // Execute the query on DB and return the response time.
            DbConnection? connection = providerFactory.CreateConnection();
            if (connection == null)
            {
                errorMessage = "Failed to create database connection.";
                _logger.LogError(errorMessage);
                return errorMessage;
            }

            using (connection)
            {
                try
                {
                    connection.ConnectionString = Utilities.NormalizeConnectionString(connectionString, databaseType, _logger);
                    using (DbCommand command = connection.CreateCommand())
                    {
                        command.CommandText = query;
                        await connection.OpenAsync();
                        using (DbDataReader reader = await command.ExecuteReaderAsync())
                        {
                            _logger.LogTrace("The health check query for datasource executed successfully.");
                        }
                    }
                }
                catch (Exception ex)
                {
                    _logger.LogError($"An exception occurred while executing the health check query: {ex.Message}");
                    errorMessage = ex.Message;
                }
            }

            return errorMessage;
        }

        // Executes the REST query by sending a GET request to the API.
        public async Task<string?> ExecuteRestQueryAsync(string restUriSuffix, string entityName, int first, string incomingRoleHeader, string incomingRoleToken)
        {
            string? errorMessage = null;
            try
            {
                // Base URL of the API that handles SQL operations
                string apiRoute = $"{restUriSuffix}{Utilities.CreateHttpRestQuery(entityName, first)}";
                if (string.IsNullOrEmpty(apiRoute))
                {
                    _logger.LogError("The API route is not available, hence HealthEndpoint is not available.");
                    return errorMessage;
                }

                if (!Program.CheckSanityOfUrl($"{_httpClient.BaseAddress}{restUriSuffix}"))
                {
                    _logger.LogError("Blocked outbound request due to invalid or unsafe URI.");
                    return "Blocked outbound request due to invalid or unsafe URI.";
                }

                HttpRequestMessage message = new(method: HttpMethod.Get, requestUri: apiRoute);
                if (!string.IsNullOrEmpty(incomingRoleToken))
                {
                    message.Headers.Add(AuthenticationOptions.CLIENT_PRINCIPAL_HEADER, incomingRoleToken);
                }

                if (!string.IsNullOrEmpty(incomingRoleHeader))
                {
                    message.Headers.Add(AuthorizationResolver.CLIENT_ROLE_HEADER, incomingRoleHeader);
                }

                HttpResponseMessage response = await _httpClient.SendAsync(message);
                if (response.IsSuccessStatusCode)
                {
                    _logger.LogTrace($"The REST HealthEndpoint query executed successfully with code {response.IsSuccessStatusCode}.");
                }
                else
                {
                    errorMessage = $"The REST HealthEndpoint query failed with code: {response.StatusCode}.";
                }

                return errorMessage;
            }
            catch (Exception ex)
            {
                _logger.LogError($"An exception occurred while executing the health check REST query: {ex.Message}");
                return ex.Message;
            }
        }

        // Executes the MCP query by sending an initialize JSON-RPC POST request to the MCP endpoint.
        public async Task<string?> ExecuteMcpQueryAsync(string mcpUriSuffix, string incomingRoleHeader, string incomingRoleToken)
        {
            string? errorMessage = null;
            try
            {
                if (string.IsNullOrWhiteSpace(mcpUriSuffix))
                {
                    const string msg = "MCP path is not configured.";
                    _logger.LogError(msg);
                    return msg;
                }

                if (_httpClient.BaseAddress is null)
                {
                    const string msg = "Health check HTTP client BaseAddress is not configured.";
                    _logger.LogError(msg);
                    return msg;
                }

                // Ensure the configured MCP path cannot override the host (e.g. absolute URIs).
                if (!Uri.TryCreate(mcpUriSuffix, UriKind.Relative, out _))
                {
                    _logger.LogError("Blocked outbound request due to invalid or unsafe URI.");
                    return "Blocked outbound request due to invalid or unsafe URI.";
                }

                Uri requestUri = new(_httpClient.BaseAddress, mcpUriSuffix);

                if (!Program.CheckSanityOfUrl(requestUri.AbsoluteUri))
                {
                    _logger.LogError("Blocked outbound request due to invalid or unsafe URI.");
                    return "Blocked outbound request due to invalid or unsafe URI.";
                }

                string jsonPayload = Utilities.CreateHttpMcpQuery();
                HttpContent content = new StringContent(jsonPayload, Encoding.UTF8, Utilities.JSON_CONTENT_TYPE);

                HttpRequestMessage message = new(method: HttpMethod.Post, requestUri: requestUri)
                {
                    Content = content
                }

                // The MCP Streamable HTTP transport requires the client to accept both
                // JSON and SSE responses.
                message.Headers.Add("Accept", "application/json, text/event-stream");

                if (!string.IsNullOrEmpty(incomingRoleToken))
                {
                    message.Headers.Add(AuthenticationOptions.CLIENT_PRINCIPAL_HEADER, incomingRoleToken);
                }

                if (!string.IsNullOrEmpty(incomingRoleHeader))
                {
                    message.Headers.Add(AuthorizationResolver.CLIENT_ROLE_HEADER, incomingRoleHeader);
                }

                HttpResponseMessage response = await _httpClient.SendAsync(message);
                if (response.IsSuccessStatusCode)
                {
                    _logger.LogTrace($"The MCP HealthEndpoint query executed successfully with code {response.StatusCode}.");
                }
                else
                {
                    errorMessage = $"The MCP HealthEndpoint query failed with code: {response.StatusCode}.";
                }

                return errorMessage;
            }
            catch (Exception ex)
            {
                _logger.LogError($"An exception occurred while executing the health check MCP query: {ex.Message}");
                return ex.Message;
            }
        }

        // Executes the GraphQL query by sending a POST request to the API.
        // Internally calls the metadata provider to fetch the column names to create the graphql payload.
        public async Task<string?> ExecuteGraphQLQueryAsync(string graphqlUriSuffix, string entityName, Entity entity, string incomingRoleHeader, string incomingRoleToken)
        {
            string? errorMessage = null;

            if (!Program.CheckSanityOfUrl($"{_httpClient.BaseAddress}{graphqlUriSuffix}"))
            {
                _logger.LogError("Blocked outbound request due to invalid or unsafe URI.");
                return "Blocked outbound request due to invalid or unsafe URI.";
            }

            try
            {
                string dataSourceName = _runtimeConfigProvider.GetConfig().GetDataSourceNameFromEntityName(entityName);

                // Fetch Column Names from Metadata Provider
                ISqlMetadataProvider sqlMetadataProvider = _metadataProviderFactory.GetMetadataProvider(dataSourceName);
                DatabaseObject dbObject = sqlMetadataProvider.GetDatabaseObjectByKey(entityName);
                List<string> columnNames = dbObject.SourceDefinition.Columns.Keys.ToList();

                // In case of GraphQL API, use the plural value specified in [entity.graphql.type.plural].
                // Further, we need to camel case this plural value to match the GraphQL object name.
                string graphqlObjectName = GraphQLNaming.GenerateListQueryName(entityName, entity);

                // In case any primitive column names are present, execute the query
                if (columnNames.Any())
                {
                    string jsonPayload = Utilities.CreateHttpGraphQLQuery(graphqlObjectName, columnNames, entity.EntityFirst);
                    HttpContent content = new StringContent(jsonPayload, Encoding.UTF8, Utilities.JSON_CONTENT_TYPE);

                    HttpRequestMessage message = new(method: HttpMethod.Post, requestUri: graphqlUriSuffix)
                    {
                        Content = content
                    };

                    if (!string.IsNullOrEmpty(incomingRoleToken))
                    {
                        message.Headers.Add(AuthenticationOptions.CLIENT_PRINCIPAL_HEADER, incomingRoleToken);
                    }

                    if (!string.IsNullOrEmpty(incomingRoleHeader))
                    {
                        message.Headers.Add(AuthorizationResolver.CLIENT_ROLE_HEADER, incomingRoleHeader);
                    }

                    HttpResponseMessage response = await _httpClient.SendAsync(message);

                    if (response.IsSuccessStatusCode)
                    {
                        _logger.LogTrace("The GraphQL HealthEndpoint query executed successfully.");
                    }
                    else
                    {
                        errorMessage = $"The GraphQL HealthEndpoint query failed with code: {response.StatusCode}.";
                    }
                }

                return errorMessage;
            }
            catch (Exception ex)
            {
                _logger.LogError($"An exception occurred while executing the GraphQL health check query: {ex.Message}");
                return ex.Message;
            }
        }
    }
}