fix(surrealdb): improve surrealdb healthchecks (#795)

* fix(surrealdb): improve surrealdb healthchecks

* fix(surrealdb): escape password value

* Adding some logging statements in the health check for diagnostics

* Bit more diag messaging

* fix: new client for each healthcheck call

* fix: wait until ns created

---------

Co-authored-by: Aaron Powell <[email protected]>

Author: Odonno

src/CommunityToolkit.Aspire.Hosting.SurrealDb/SurrealDbBuilderExtensions.cs

@@ -72,16 +72,16 @@ public static IResourceBuilder<SurrealDbServerResource> AddSurrealServer(
         {
             args.Add("--strict");
         }
-        
+
         // The password must be at least 8 characters long and contain characters from three of the following four sets: Uppercase letters, Lowercase letters, Base 10 digits, and Symbols
         var passwordParameter = password?.Resource ?? ParameterResourceBuilderExtensions.CreateDefaultPasswordParameter(builder, $"{name}-password", minLower: 1, minUpper: 1, minNumeric: 1);
 
         string imageTag = builder.ExecutionContext.IsRunMode
             ? $"{SurrealDbContainerImageTags.Tag}-dev"
             : SurrealDbContainerImageTags.Tag;
-        
+
         var surrealServer = new SurrealDbServerResource(name, userName?.Resource, passwordParameter);
-        
+
         return builder.AddResource(surrealServer)
                       .WithEndpoint(port: port, targetPort: SurrealDbPort, name: SurrealDbServerResource.PrimaryEndpointName)
                       .WithImage(SurrealDbContainerImageTags.Image, imageTag)
@@ -99,42 +99,65 @@ public static IResourceBuilder<SurrealDbServerResource> AddSurrealServer(
                           {
                               return;
                           }
-                          
+
                           var connectionString = await surrealServer.GetConnectionStringAsync(ct).ConfigureAwait(false);
                           if (connectionString is null)
                           {
                               throw new DistributedApplicationException($"ResourceReadyEvent was published for the '{surrealServer.Name}' resource but the connection string was null.");
                           }
-                          
-                          var options = new SurrealDbOptionsBuilder().FromConnectionString(connectionString).Build();
-                          await using var surrealClient = new SurrealDbClient(options);
-                          
-                          foreach (var nsResourceName in surrealServer.Namespaces.Keys)
-                          {
-                              if (builder.Resources.FirstOrDefault(n =>
-                                      string.Equals(n.Name, nsResourceName, StringComparison.OrdinalIgnoreCase)) is
-                                  SurrealDbNamespaceResource surrealDbNamespace)
-                              {
-                                  await CreateNamespaceAsync(surrealClient, surrealDbNamespace, @event.Services, ct)
-                                      .ConfigureAwait(false);
-                              
-                                  await surrealClient.Use(surrealDbNamespace.NamespaceName, null!, ct).ConfigureAwait(false);
-                          
-                                  foreach (var dbResourceName in surrealDbNamespace.Databases.Keys)
-                                  {
-                                      if (builder.Resources.FirstOrDefault(n =>
-                                              string.Equals(n.Name, dbResourceName, StringComparison.OrdinalIgnoreCase)) is
-                                          SurrealDbDatabaseResource surrealDbDatabase)
-                                      {
-                                          await CreateDatabaseAsync(surrealClient, surrealDbDatabase, @event.Services, ct)
-                                              .ConfigureAwait(false);
-                                      }
-                                  }
-                              }
-                          }
+
+                          await EnsuresNsDbCreated(builder, connectionString, surrealServer, @event.Services, ct);
                       });
     }
 
+    private static async Task EnsuresNsDbCreated(
+        IDistributedApplicationBuilder builder,
+        string connectionString,
+        SurrealDbServerResource surrealServer,
+        IServiceProvider services,
+        CancellationToken ct
+    )
+    {
+        var options = new SurrealDbOptionsBuilder().FromConnectionString(connectionString).Build();
+        await using var surrealClient = new SurrealDbClient(options);
+
+        foreach (var nsResourceName in surrealServer.Namespaces.Keys)
+        {
+            if (builder.Resources.FirstOrDefault(n =>
+                    string.Equals(n.Name, nsResourceName, StringComparison.OrdinalIgnoreCase)) is
+                SurrealDbNamespaceResource surrealDbNamespace)
+            {
+                await CreateNamespaceAsync(surrealClient, surrealDbNamespace, services, ct)
+                    .ConfigureAwait(false);
+
+                // 💡 Wait until the Namespace is really created?!
+                while (!ct.IsCancellationRequested)
+                {
+                    try
+                    {
+                        await surrealClient.Use(surrealDbNamespace.NamespaceName, null!, ct).ConfigureAwait(false);
+                        break;
+                    }
+                    catch
+                    {
+                        await Task.Delay(200, ct).ConfigureAwait(false);
+                    }
+                }
+
+                foreach (var dbResourceName in surrealDbNamespace.Databases.Keys)
+                {
+                    if (builder.Resources.FirstOrDefault(n =>
+                            string.Equals(n.Name, dbResourceName, StringComparison.OrdinalIgnoreCase)) is
+                        SurrealDbDatabaseResource surrealDbDatabase)
+                    {
+                        await CreateDatabaseAsync(surrealClient, surrealDbDatabase, services, ct)
+                            .ConfigureAwait(false);
+                    }
+                }
+            }
+        }
+    }
+
     /// <summary>
     /// Adds a SurrealDB namespace to the application model. This is a child resource of a <see cref="SurrealDbServerResource"/>.
     /// </summary>
@@ -174,7 +197,7 @@ public static IResourceBuilder<SurrealDbNamespaceResource> AddNamespace(
         var surrealServerNamespace = new SurrealDbNamespaceResource(name, namespaceName, builder.Resource);
         return builder.ApplicationBuilder.AddResource(surrealServerNamespace);
     }
-    
+
     /// <summary>
     /// Defines the SQL script used to create the namespace.
     /// </summary>
@@ -232,38 +255,34 @@ public static IResourceBuilder<SurrealDbDatabaseResource> AddDatabase(
         builder.Resource.AddDatabase(name, databaseName);
         var surrealServerDatabase = new SurrealDbDatabaseResource(name, databaseName, builder.Resource);
 
-        SurrealDbClient? surrealDbClient = null;
-
-        builder.ApplicationBuilder.Eventing.Subscribe<ConnectionStringAvailableEvent>(surrealServerDatabase, async (@event, ct) =>
-        {
-            var connectionString = await surrealServerDatabase.ConnectionStringExpression.GetValueAsync(ct).ConfigureAwait(false);
-            if (connectionString is null)
-            {
-                throw new DistributedApplicationException($"ConnectionStringAvailableEvent was published for the '{surrealServerDatabase}' resource but the connection string was null.");
-            }
-
-            var options = new SurrealDbOptionsBuilder().FromConnectionString(connectionString).Build();
-            surrealDbClient = new SurrealDbClient(options);
-        });
+        SurrealDbOptions? surrealDbOptions = null;
 
         string namespaceName = builder.Resource.Name;
         string serverName = builder.Resource.Parent.Name;
 
         string healthCheckKey = $"{serverName}_{namespaceName}_{name}_check";
-        // TODO : Bug to be fixed
-        //builder.ApplicationBuilder.Services.AddHealthChecks().AddSurreal(_ => surrealDbClient!, healthCheckKey);
         builder.ApplicationBuilder.Services.AddHealthChecks().Add(new HealthCheckRegistration(
                 name: healthCheckKey,
-                _ => new SurrealDbHealthCheck(surrealDbClient!),
+                sp => new SurrealDbHealthCheck(surrealDbOptions!, sp.GetRequiredService<ILogger<SurrealDbHealthCheck>>()),
                 failureStatus: null,
                 tags: null
             )
         );
-        
+
         return builder.ApplicationBuilder.AddResource(surrealServerDatabase)
-            .WithHealthCheck(healthCheckKey);
+            .WithHealthCheck(healthCheckKey)
+            .OnConnectionStringAvailable(async (_, _, ct) =>
+            {
+                var connectionString = await surrealServerDatabase.ConnectionStringExpression.GetValueAsync(ct).ConfigureAwait(false);
+                if (connectionString is null)
+                {
+                    throw new DistributedApplicationException($"ConnectionStringAvailableEvent was published for the '{surrealServerDatabase}' resource but the connection string was null.");
+                }
+
+                surrealDbOptions = new SurrealDbOptionsBuilder().FromConnectionString(connectionString).Build();
+            });
     }
-    
+
     /// <summary>
     /// Defines the SQL script used to create the database.
     /// </summary>
@@ -347,7 +366,7 @@ public static IResourceBuilder<SurrealDbServerResource> WithDataBindMount(this I
 
         return builder.WithBindMount(source, "/data");
     }
-    
+
     /// <summary>
     /// Copies init files into a SurrealDB container resource.
     /// </summary>
@@ -367,10 +386,10 @@ public static IResourceBuilder<SurrealDbServerResource> WithInitFiles(this IReso
         {
             throw new DistributedApplicationException($"Unable to determine the file name for '{source}'.");
         }
-        
+
         string fileName = Path.GetFileName(importFullPath);
         string initFilePath = $"{initPath}/{fileName}";
-        
+
         return builder
             .WithContainerFiles(initPath, importFullPath)
             .WithEnvironment(context =>
@@ -413,16 +432,16 @@ public static IResourceBuilder<T> WithSurrealist<T>(
             .WithHttpEndpoint(targetPort: 8080, name: "http")
             .WithRelationship(builder.Resource, "Surrealist")
             .ExcludeFromManifest();
-        
+
         surrealistContainerBuilder.WithContainerFiles(
             destinationPath: "/usr/share/nginx/html",
             callback: async (_, cancellationToken) =>
             {
-                var surrealDbServerInstances = 
+                var surrealDbServerInstances =
                     builder.ApplicationBuilder.Resources.OfType<SurrealDbServerResource>().ToList();
-                var surrealDbNamespaceResources = 
+                var surrealDbNamespaceResources =
                     builder.ApplicationBuilder.Resources.OfType<SurrealDbNamespaceResource>().ToList();
-                var surrealDbDatabaseResources = 
+                var surrealDbDatabaseResources =
                     builder.ApplicationBuilder.Resources.OfType<SurrealDbDatabaseResource>().ToList();
 
                 return [
@@ -431,8 +450,8 @@ public static IResourceBuilder<T> WithSurrealist<T>(
                         Name = "instance.json",
                         Contents = await WriteSurrealistInstanceJson(
                             surrealDbServerInstances,
-                            surrealDbNamespaceResources, 
-                            surrealDbDatabaseResources, 
+                            surrealDbNamespaceResources,
+                            surrealDbDatabaseResources,
                             cancellationToken
                         ).ConfigureAwait(false),
                     },
@@ -443,7 +462,7 @@ public static IResourceBuilder<T> WithSurrealist<T>(
 
         return builder;
     }
-    
+
     private static async Task<string> WriteSurrealistInstanceJson(
         IList<SurrealDbServerResource> surrealDbServerInstances,
         IList<SurrealDbNamespaceResource> surrealDbNamespaceResources,
@@ -522,16 +541,16 @@ CancellationToken cancellationToken
         writer.WriteEndArray();
 
         writer.WriteEndObject();
-        
+
         await writer.FlushAsync(cancellationToken);
-        
+
         return Encoding.UTF8.GetString(stream.ToArray());
     }
-    
+
     private static async Task CreateNamespaceAsync(
-        SurrealDbClient surrealClient, 
-        SurrealDbNamespaceResource namespaceResource, 
-        IServiceProvider serviceProvider, 
+        SurrealDbClient surrealClient,
+        SurrealDbNamespaceResource namespaceResource,
+        IServiceProvider serviceProvider,
         CancellationToken cancellationToken
     )
     {
@@ -543,7 +562,7 @@ CancellationToken cancellationToken
         try
         {
             var response = await surrealClient.RawQuery(
-                scriptAnnotation?.Script ?? $"DEFINE NAMESPACE IF NOT EXISTS `{namespaceResource.NamespaceName}`;", 
+                scriptAnnotation?.Script ?? $"DEFINE NAMESPACE IF NOT EXISTS `{namespaceResource.NamespaceName}`;",
                 cancellationToken: cancellationToken
             ).ConfigureAwait(false);
 
@@ -556,11 +575,11 @@ CancellationToken cancellationToken
             logger.LogError(e, "Failed to create namespace '{NamespaceName}'", namespaceResource.NamespaceName);
         }
     }
-    
+
     private static async Task CreateDatabaseAsync(
-        SurrealDbClient surrealClient, 
-        SurrealDbDatabaseResource databaseResource, 
-        IServiceProvider serviceProvider, 
+        SurrealDbClient surrealClient,
+        SurrealDbDatabaseResource databaseResource,
+        IServiceProvider serviceProvider,
         CancellationToken cancellationToken
     )
     {
@@ -572,7 +591,7 @@ CancellationToken cancellationToken
         try
         {
             var response = await surrealClient.RawQuery(
-                scriptAnnotation?.Script ?? $"DEFINE DATABASE IF NOT EXISTS `{databaseResource.DatabaseName}`;", 
+                scriptAnnotation?.Script ?? $"DEFINE DATABASE IF NOT EXISTS `{databaseResource.DatabaseName}`;",
                 cancellationToken: cancellationToken
             ).ConfigureAwait(false);
 

src/CommunityToolkit.Aspire.Hosting.SurrealDb/SurrealDbServerResource.cs

@@ -54,7 +54,7 @@ UserNameParameter is not null ?
 
     private ReferenceExpression ConnectionString =>
         ReferenceExpression.Create(
-            $"Server={SchemeUri}://{PrimaryEndpoint.Property(EndpointProperty.Host)}:{PrimaryEndpoint.Property(EndpointProperty.Port)}/rpc;User={UserNameReference};Password={PasswordParameter}");
+            $"Server={SchemeUri}://{PrimaryEndpoint.Property(EndpointProperty.Host)}:{PrimaryEndpoint.Property(EndpointProperty.Port)}/rpc;User={UserNameReference};Password='{PasswordParameter}'");
 
     /// <summary>
     /// Gets the connection string expression for the SurrealDB instance.

src/CommunityToolkit.Aspire.SurrealDb/AspireSurrealDbExtensions.cs

@@ -6,6 +6,7 @@
 using Microsoft.Extensions.Configuration;
 using Microsoft.Extensions.DependencyInjection;
 using Microsoft.Extensions.Diagnostics.HealthChecks;
+using Microsoft.Extensions.Logging;
 using SurrealDb.Net;
 
 namespace Microsoft.Extensions.Hosting;
@@ -16,7 +17,7 @@ namespace Microsoft.Extensions.Hosting;
 public static class AspireSurrealDbExtensions
 {
     private const string DefaultConfigSectionName = "Aspire:Surreal:Client";
-    
+
     /// <summary>
     /// Registers <see cref="SurrealDbClient" /> in the services provided by the <paramref name="builder"/>.
     /// </summary>
@@ -97,9 +98,9 @@ private static void AddSurrealClient(
 
             builder.TryAddHealthCheck(new HealthCheckRegistration(
                 healthCheckName,
-                sp => new SurrealDbHealthCheck(serviceKey is null ?
-                    sp.GetRequiredService<SurrealDbClient>() :
-                    sp.GetRequiredKeyedService<SurrealDbClient>(serviceKey)),
+                sp => new SurrealDbHealthCheck(
+                    settings.Options,
+                    sp.GetRequiredService<ILogger<SurrealDbHealthCheck>>()),
                 failureStatus: null,
                 tags: null,
                 timeout: settings.HealthCheckTimeout > 0 ? TimeSpan.FromMilliseconds(settings.HealthCheckTimeout.Value) : null

src/CommunityToolkit.Aspire.SurrealDb/SurrealDbHealthCheck.cs

@@ -1,34 +1,47 @@
 // Licensed to the .NET Foundation under one or more agreements.
 // The .NET Foundation licenses this file to you under the MIT license.
 
+using Microsoft.Extensions.DependencyInjection;
 using SurrealDb.Net;
 using Microsoft.Extensions.Diagnostics.HealthChecks;
+using Microsoft.Extensions.Logging;
 
 namespace CommunityToolkit.Aspire.SurrealDb;
 
 internal sealed class SurrealDbHealthCheck : IHealthCheck
 {
-    private readonly ISurrealDbClient _surrealdbClient;
+    private readonly SurrealDbOptions _options;
+    private readonly ILogger<SurrealDbHealthCheck> _logger;
 
-    public SurrealDbHealthCheck(ISurrealDbClient surrealdbClient)
+    public SurrealDbHealthCheck(SurrealDbOptions options, ILogger<SurrealDbHealthCheck> logger)
     {
-        ArgumentNullException.ThrowIfNull(surrealdbClient, nameof(surrealdbClient));
-        _surrealdbClient = surrealdbClient;
+        _options = options;
+        _logger = logger;
     }
 
     /// <inheritdoc />
     public async Task<HealthCheckResult> CheckHealthAsync(HealthCheckContext context, CancellationToken cancellationToken = default)
     {
+        bool isHealthy = false;
+        
         try
         {
-            bool isHealthy = await _surrealdbClient.Health(cancellationToken).ConfigureAwait(false);
+            await using var surrealdbClient = new SurrealDbClient(_options);
+            
+            isHealthy = await surrealdbClient.Health(cancellationToken).ConfigureAwait(false);
+            var response = await surrealdbClient.RawQuery("RETURN 1", cancellationToken: cancellationToken).ConfigureAwait(false);
+            response.EnsureAllOks();
+
+            _logger.LogInformation("SurrealDB health check passed. Response: {Response}", response);
+            _logger.LogInformation("SurrealDB health check outcome: {Outcome}", isHealthy ? "Healthy" : "Unhealthy");
 
             return isHealthy
                 ? HealthCheckResult.Healthy()
                 : new HealthCheckResult(context.Registration.FailureStatus);
         }
         catch (Exception ex)
         {
+            _logger.LogError(ex, "SurrealDB health check raised an exception. Health check had previously reported: {Outcome}. CancellationToken status: {CancellationTokenStatus}", isHealthy ? "Healthy" : "Unhealthy", cancellationToken.IsCancellationRequested ? "Canceled" : "Active");
             return new HealthCheckResult(context.Registration.FailureStatus, exception: ex);
         }
     }