Maintain Default Data Source Name from Original Config after Hot Reloading

src/Config/FileSystemRuntimeConfigLoader.cs

@@ -94,17 +94,32 @@ public string GetConfigFileName()
     /// <param name="config">The loaded <c>RuntimeConfig</c>, or null if none was loaded.</param>
     /// <param name="replaceEnvVar">Whether to replace environment variable with its
     /// value or not while deserializing.</param>
+    /// <param name="dataSourceName">If provided and not empty, this is the data source name that will be used in the loaded config.</param>
     /// <returns>True if the config was loaded, otherwise false.</returns>
     public bool TryLoadConfig(
         string path,
         [NotNullWhen(true)] out RuntimeConfig? config,
-        bool replaceEnvVar = false)
+        bool replaceEnvVar = false,
+        string dataSourceName = "")
     {
         if (_fileSystem.File.Exists(path))
         {
             Console.WriteLine($"Loading config file from {path}.");
             string json = _fileSystem.File.ReadAllText(path);
-            return TryParseConfig(json, out config, connectionString: _connectionString, replaceEnvVar: replaceEnvVar);
+
+            if (TryParseConfig(
+                json: json,
+                config: out config,
+                connectionString: _connectionString,
+                replaceEnvVar: replaceEnvVar))
+            {
+                if (!string.IsNullOrEmpty(dataSourceName))
+                {
+                    config.UpdateDefaultDataSourceName(dataSourceName);
+                }
+
+                return true;
+            }
         }
         else
         {
@@ -123,10 +138,11 @@ public bool TryLoadConfig(
     /// <param name="config">The loaded <c>RuntimeConfig</c>, or null if none was loaded.</param>
     /// <param name="replaceEnvVar">Whether to replace environment variable with its
     /// value or not while deserializing.</param>
+    /// <param name="dataSourceName">The data source name to be used in the loaded config.</param>
     /// <returns>True if the config was loaded, otherwise false.</returns>
-    public override bool TryLoadKnownConfig([NotNullWhen(true)] out RuntimeConfig? config, bool replaceEnvVar = false)
+    public override bool TryLoadKnownConfig([NotNullWhen(true)] out RuntimeConfig? config, bool replaceEnvVar = false, string dataSourceName = "")
     {
-        return TryLoadConfig(ConfigFilePath, out config, replaceEnvVar);
+        return TryLoadConfig(ConfigFilePath, out config, replaceEnvVar, dataSourceName);
     }
 
     /// <summary>

src/Config/ObjectModel/RuntimeConfig.cs

@@ -130,7 +130,7 @@ Runtime.GraphQL is null ||
         }
     }
 
-    private string _defaultDataSourceName;
+    public string _defaultDataSourceName;
 
     private Dictionary<string, DataSource> _dataSourceNameToDataSource;
 
@@ -273,6 +273,30 @@ public void UpdateDataSourceNameToDataSource(string dataSourceName, DataSource d
         _dataSourceNameToDataSource[dataSourceName] = dataSource;
     }
 
+    /// <summary>
+    /// In a Hot Reload scenario we should maintain the same default data source
+    /// name before the hot reload as after the hot reload. This is because we hold
+    /// references to the Data Source itself which depend on this data source name
+    /// for lookups. To correctly retrieve this information after a hot reload
+    /// we need the data source name to stay the same after hot reloading. This method takes
+    /// a default data source name, such as the one from before hot reload, and
+    /// replaces the current dictionary entries of this RuntimeConfig that were
+    /// built using a new, unique guid during the construction of this RuntimeConfig
+    /// with entries using the provided default data source name. We then update the DefaultDataSourceName.
+    /// </summary>
+    /// <param name="initialDefaultDataSourceName">The name used to update the dictionaries.</param>
+    public void UpdateDefaultDataSourceName(string initialDefaultDataSourceName)
+    {
+        _dataSourceNameToDataSource.Remove(_defaultDataSourceName);
+        _dataSourceNameToDataSource.Add(initialDefaultDataSourceName, this.DataSource);
+        foreach (KeyValuePair<string, Entity> entity in Entities)
+        {
+            _entityNameToDataSourceName[entity.Key] = initialDefaultDataSourceName;
+        }
+
+        _defaultDataSourceName = initialDefaultDataSourceName;
+    }
+
     /// <summary>
     /// Gets datasourceName from EntityNameToDatasourceName dictionary.
     /// </summary>

src/Config/RuntimeConfigLoader.cs

@@ -33,8 +33,9 @@ public RuntimeConfigLoader(string? connectionString = null)
     /// <param name="config">The loaded <c>RuntimeConfig</c>, or null if none was loaded.</param>
     /// <param name="replaceEnvVar">Whether to replace environment variable with its
     /// value or not while deserializing.</param>
+    /// <param name="dataSourceName">The data source name to be used in the loaded config.</param>
     /// <returns>True if the config was loaded, otherwise false.</returns>
-    public abstract bool TryLoadKnownConfig([NotNullWhen(true)] out RuntimeConfig? config, bool replaceEnvVar = false);
+    public abstract bool TryLoadKnownConfig([NotNullWhen(true)] out RuntimeConfig? config, bool replaceEnvVar = false, string dataSourceName = "");
 
     /// <summary>
     /// Returns the link to the published draft schema.

src/Core/Configurations/ConfigFileWatcher.cs

@@ -38,6 +38,7 @@ private void OnConfigFileChange(object sender, FileSystemEventArgs e)
 
             if (!_configProvider.IsLateConfigured && _configProvider.GetConfig().IsDevelopmentMode())
             {
+                // pass along the original default data source name for consistency within runtime config's data structures
                 _configProvider.HotReloadConfig();
             }
         }

src/Core/Configurations/RuntimeConfigProvider.cs

@@ -151,7 +151,8 @@ public bool TryGetLoadedConfig([NotNullWhen(true)] out RuntimeConfig? runtimeCon
     /// </summary>
     public void HotReloadConfig()
     {
-        ConfigLoader.TryLoadKnownConfig(out _runtimeConfig, replaceEnvVar: true);
+        // _runtimeconfig can not be null in a hot reload scenario
+        ConfigLoader.TryLoadKnownConfig(out _runtimeConfig, replaceEnvVar: true, _runtimeConfig!.GetDefaultDataSourceName());
     }
 
     /// <summary>

src/Service.Tests/Unittests/ConfigFileWatcherUnitTests.cs

@@ -13,31 +13,17 @@ namespace Azure.DataApiBuilder.Service.Tests.Unittests
     [TestClass]
     public class ConfigFileWatcherUnitTests
     {
-        /// <summary>
-        /// Use the file system (not mocked) to create a hot reload
-        /// scenario of the REST runtime options and verify that we
-        /// correctly hot reload those options.
-        /// </summary>
-        [TestMethod]
-        public void HotReloadConfigRestRuntimeOptions()
-        {
-            string initialRestPath = "/api";
-            string updatedRestPath = "/rest";
-            string initialGQLPath = "/api";
-            string updatedGQLPath = "/gql";
-
-            bool initialRestEnabled = true;
-            bool updatedRestEnabled = false;
-            bool initialGQLEnabled = true;
-            bool updatedGQLEnabled = false;
 
-            bool initialGQLIntrospection = true;
-            bool updatedGQLIntrospection = false;
-
-            HostMode initialMode = HostMode.Development;
-            HostMode updatedMode = HostMode.Production;
-
-            string initialConfig = @"
+        private static string GenerateRuntimeSectionStringFromParams(
+            string restPath,
+            string gqlPath,
+            bool restEnabled,
+            bool gqlEnabled,
+            bool gqlIntrospection,
+            HostMode mode
+            )
+        {
+            string runtimeString = @"
 {
   ""$schema"": ""https://github.com/Azure/data-api-builder/releases/download/vmajor.minor.patch/dab.draft.schema.json"",
   ""data-source"": {
@@ -49,48 +35,13 @@ public void HotReloadConfigRestRuntimeOptions()
   },
   ""runtime"": {
     ""rest"": {
-      ""enabled"": " + initialRestEnabled.ToString().ToLower() + @",
-      ""path"": """ + initialRestPath + @"""
-    },
-    ""graphql"": {
-      ""enabled"": " + initialGQLEnabled.ToString().ToLower() + @",
-      ""path"": """ + initialGQLPath + @""",
-      ""allow-introspection"": " + initialGQLIntrospection.ToString().ToLower() + @"
-    },
-    ""host"": {
-      ""cors"": {
-        ""origins"": [
-          ""http://localhost:5000""
-        ],
-        ""allow-credentials"": false
-      },
-      ""authentication"": {
-        ""provider"": ""StaticWebApps""
-      },
-      ""mode"": """ + initialMode + @"""
-    }
-  },
-  ""entities"": {}
-}";
-            string updatedConfig = @"
-{
-  ""$schema"": ""https://github.com/Azure/data-api-builder/releases/download/vmajor.minor.patch/dab.draft.schema.json"",
-  ""data-source"": {
-    ""database-type"": ""mssql"",
-    ""connection-string"": ""Server=test;Database=test;User ID=test"",
-    ""options"": {
-      ""set-session-context"": true
-    }
-  },
-  ""runtime"": {
-    ""rest"": {
-      ""enabled"": " + updatedRestEnabled.ToString().ToLower() + @",
-      ""path"": """ + updatedRestPath + @"""
+      ""enabled"": " + restEnabled.ToString().ToLower() + @",
+      ""path"": """ + restPath + @"""
     },
     ""graphql"": {
-      ""enabled"": " + updatedGQLEnabled.ToString().ToLower() + @",
-      ""path"": """ + updatedGQLPath + @""",
-      ""allow-introspection"": " + updatedGQLIntrospection.ToString().ToLower() + @"
+      ""enabled"": " + gqlEnabled.ToString().ToLower() + @",
+      ""path"": """ + gqlPath + @""",
+      ""allow-introspection"": " + gqlIntrospection.ToString().ToLower() + @"
     },
     ""host"": {
       ""cors"": {
@@ -102,11 +53,57 @@ public void HotReloadConfigRestRuntimeOptions()
       ""authentication"": {
         ""provider"": ""StaticWebApps""
       },
-      ""mode"": """ + updatedMode + @"""
+      ""mode"": """ + mode + @"""
     }
   },
   ""entities"": {}
 }";
+            return runtimeString;
+        }
+
+        /// <summary>
+        /// Use the file system (not mocked) to create a hot reload
+        /// scenario of the REST runtime options and verify that we
+        /// correctly hot reload those options.
+        /// </summary>
+        [TestMethod]
+        public void HotReloadConfigRestRuntimeOptions()
+        {
+            // Arrange
+            // 1. Setup the strings that are turned into our initital and runtime config to reload.
+            // 2. Generate initial runtimeconfig, start file watching, and assert we have valid initial values.
+            string initialRestPath = "/api";
+            string updatedRestPath = "/rest";
+            string initialGQLPath = "/api";
+            string updatedGQLPath = "/gql";
+
+            bool initialRestEnabled = true;
+            bool updatedRestEnabled = false;
+            bool initialGQLEnabled = true;
+            bool updatedGQLEnabled = false;
+
+            bool initialGQLIntrospection = true;
+            bool updatedGQLIntrospection = false;
+
+            HostMode initialMode = HostMode.Development;
+            HostMode updatedMode = HostMode.Production;
+
+            string initialConfig = GenerateRuntimeSectionStringFromParams(
+                initialRestPath,
+                initialGQLPath,
+                initialRestEnabled,
+                initialGQLEnabled,
+                initialGQLIntrospection,
+                initialMode);
+
+            string updatedConfig = GenerateRuntimeSectionStringFromParams(
+                updatedRestPath,
+                updatedGQLPath,
+                updatedRestEnabled,
+                updatedGQLEnabled,
+                updatedGQLIntrospection,
+                updatedMode);
+
             string configName = "hotreload-config.json";
             if (File.Exists(configName))
             {
@@ -118,8 +115,11 @@ public void HotReloadConfigRestRuntimeOptions()
             fileSystem.File.WriteAllText(configName, initialConfig);
             FileSystemRuntimeConfigLoader configLoader = new(fileSystem, configName, string.Empty);
             RuntimeConfigProvider configProvider = new(configLoader);
+
             // Must GetConfig() to start file watching
             RuntimeConfig runtimeConfig = configProvider.GetConfig();
+            string initialDefaultDataSourceName = runtimeConfig.GetDefaultDataSourceName();
+
             // assert we have a valid config
             Assert.IsNotNull(runtimeConfig);
             Assert.AreEqual(initialRestEnabled, runtimeConfig.Runtime.Rest.Enabled);
@@ -128,18 +128,29 @@ public void HotReloadConfigRestRuntimeOptions()
             Assert.AreEqual(initialGQLPath, runtimeConfig.Runtime.GraphQL.Path);
             Assert.AreEqual(initialGQLIntrospection, runtimeConfig.Runtime.GraphQL.AllowIntrospection);
             Assert.AreEqual(initialMode, runtimeConfig.Runtime.Host.Mode);
+
             // Simulate change to the config file
             fileSystem.File.WriteAllText(configName, updatedConfig);
+
             // Give ConfigFileWatcher enough time to hot reload the change
             System.Threading.Thread.Sleep(1000);
-            // Hot Reloaded config
+
+            // Act
+            // 1. Hot reload the runtime config
             runtimeConfig = configProvider.GetConfig();
+            string updatedDefaultDataSourceName = runtimeConfig.GetDefaultDataSourceName();
+
+            // Assert
+            // 1. Assert we have the correcr values after a hot reload.
             Assert.AreEqual(updatedRestEnabled, runtimeConfig.Runtime.Rest.Enabled);
             Assert.AreEqual(updatedRestPath, runtimeConfig.Runtime.Rest.Path);
             Assert.AreEqual(updatedGQLEnabled, runtimeConfig.Runtime.GraphQL.Enabled);
             Assert.AreEqual(updatedGQLPath, runtimeConfig.Runtime.GraphQL.Path);
             Assert.AreEqual(updatedGQLIntrospection, runtimeConfig.Runtime.GraphQL.AllowIntrospection);
             Assert.AreEqual(updatedMode, runtimeConfig.Runtime.Host.Mode);
+
+            // DefaultDataSourceName should not change after a hot reload.
+            Assert.AreEqual(initialDefaultDataSourceName, updatedDefaultDataSourceName);
             if (File.Exists(configName))
             {
                 File.Delete(configName);