Maintain Default Data Source Name from Original Config after Hot Reloading

src/Cli.Tests/ModuleInitializer.cs

@@ -51,6 +51,8 @@ public static void Init()
         VerifierSettings.IgnoreMember<RuntimeConfig>(config => config.Schema);
         // Ignore the message as that's not serialized in our config file anyway.
         VerifierSettings.IgnoreMember<DataSource>(dataSource => dataSource.DatabaseTypeNotSupportedMessage);
+        // Ignore DefaultDataSourceName as that's not serialized in our config file.
+        VerifierSettings.IgnoreMember<RuntimeConfig>(config => config.DefaultDataSourceName);
         // Customise the path where we store snapshots, so they are easier to locate in a PR review.
         VerifyBase.DerivePathInfo(
             (sourceFile, projectDirectory, type, method) => new(

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,8 @@ Runtime.GraphQL is null ||
         }
     }
 
-    private string _defaultDataSourceName;
+    [JsonIgnore]
+    public string DefaultDataSourceName { get; private set; }
 
     private Dictionary<string, DataSource> _dataSourceNameToDataSource;
 
@@ -169,18 +170,18 @@ public RuntimeConfig(string? Schema, DataSource DataSource, RuntimeEntities Enti
         this.DataSource = DataSource;
         this.Runtime = Runtime;
         this.Entities = Entities;
-        _defaultDataSourceName = Guid.NewGuid().ToString();
+        this.DefaultDataSourceName = Guid.NewGuid().ToString();
 
         // we will set them up with default values
         _dataSourceNameToDataSource = new Dictionary<string, DataSource>
         {
-            { _defaultDataSourceName, this.DataSource }
+            { DefaultDataSourceName, this.DataSource }
         };
 
         _entityNameToDataSourceName = new Dictionary<string, string>();
         foreach (KeyValuePair<string, Entity> entity in Entities)
         {
-            _entityNameToDataSourceName.TryAdd(entity.Key, _defaultDataSourceName);
+            _entityNameToDataSourceName.TryAdd(entity.Key, DefaultDataSourceName);
         }
 
         // Process data source and entities information for each database in multiple database scenario.
@@ -241,7 +242,7 @@ public RuntimeConfig(string Schema, DataSource DataSource, RuntimeOptions Runtim
         this.DataSource = DataSource;
         this.Runtime = Runtime;
         this.Entities = Entities;
-        _defaultDataSourceName = DefaultDataSourceName;
+        this.DefaultDataSourceName = DefaultDataSourceName;
         _dataSourceNameToDataSource = DataSourceNameToDataSource;
         _entityNameToDataSourceName = EntityNameToDataSourceName;
         this.DataSourceFiles = DataSourceFiles;
@@ -273,6 +274,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>
@@ -303,17 +328,6 @@ public bool CheckDataSourceExists(string dataSourceName)
         return _dataSourceNameToDataSource.ContainsKey(dataSourceName);
     }
 
-    /// <summary>
-    /// Get the default datasource name.
-    /// </summary>
-    /// <returns>default datasourceName.</returns>
-#pragma warning disable CA1024 // Use properties where appropriate. Reason: Do not want datasource serialized and want to keep it private to restrict set;
-    public string GetDefaultDataSourceName()
-#pragma warning restore CA1024 // Use properties where appropriate
-    {
-        return _defaultDataSourceName;
-    }
-
     /// <summary>
     /// Serializes the RuntimeConfig object to JSON for writing to file.
     /// </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.
@@ -80,7 +81,7 @@ public static bool TryParseConfig(string json,
             // set dataSourceName to default if not provided
             if (string.IsNullOrEmpty(dataSourceName))
             {
-                dataSourceName = config.GetDefaultDataSourceName();
+                dataSourceName = config.DefaultDataSourceName;
             }
 
             if (!string.IsNullOrEmpty(connectionString))
@@ -113,7 +114,7 @@ public static bool TryParseConfig(string json,
                 ds = ds with { ConnectionString = updatedConnection };
                 config.UpdateDataSourceNameToDataSource(dataSourceName, ds);
 
-                if (string.Equals(dataSourceKey, config.GetDefaultDataSourceName(), StringComparison.OrdinalIgnoreCase))
+                if (string.Equals(dataSourceKey, config.DefaultDataSourceName, StringComparison.OrdinalIgnoreCase))
                 {
                     config = config with { DataSource = ds };
                 }

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!.DefaultDataSourceName);
     }
 
     /// <summary>
@@ -193,7 +194,7 @@ public async Task<bool> Initialize(
                 _runtimeConfig = HandleCosmosNoSqlConfiguration(schema, _runtimeConfig, _runtimeConfig.DataSource.ConnectionString);
             }
 
-            ManagedIdentityAccessToken[_runtimeConfig.GetDefaultDataSourceName()] = accessToken;
+            ManagedIdentityAccessToken[_runtimeConfig.DefaultDataSourceName] = accessToken;
         }
 
         bool configLoadSucceeded = await InvokeConfigLoadedHandlersAsync();
@@ -268,8 +269,8 @@ public async Task<bool> Initialize(
                 DatabaseType.CosmosDB_NoSQL => HandleCosmosNoSqlConfiguration(graphQLSchema, runtimeConfig, connectionString),
                 _ => runtimeConfig with { DataSource = runtimeConfig.DataSource with { ConnectionString = connectionString } }
             };
-            ManagedIdentityAccessToken[_runtimeConfig.GetDefaultDataSourceName()] = accessToken;
-            _runtimeConfig.UpdateDataSourceNameToDataSource(_runtimeConfig.GetDefaultDataSourceName(), _runtimeConfig.DataSource);
+            ManagedIdentityAccessToken[_runtimeConfig.DefaultDataSourceName] = accessToken;
+            _runtimeConfig.UpdateDataSourceNameToDataSource(_runtimeConfig.DefaultDataSourceName, _runtimeConfig.DataSource);
 
             return await InvokeConfigLoadedHandlersAsync();
         }
@@ -298,7 +299,7 @@ private static RuntimeConfig HandleCosmosNoSqlConfiguration(string? schema, Runt
     {
         if (string.IsNullOrEmpty(dataSourceName))
         {
-            dataSourceName = runtimeConfig.GetDefaultDataSourceName();
+            dataSourceName = runtimeConfig.DefaultDataSourceName;
         }
 
         DbConnectionStringBuilder dbConnectionStringBuilder = new()
@@ -341,7 +342,7 @@ private static RuntimeConfig HandleCosmosNoSqlConfiguration(string? schema, Runt
         // Update the connection string in the datasource with the one that was provided to the controller
         dataSource = dataSource with { Options = options, ConnectionString = connectionString };
 
-        if (dataSourceName == runtimeConfig.GetDefaultDataSourceName())
+        if (dataSourceName == runtimeConfig.DefaultDataSourceName)
         {
             // update default db.
             runtimeConfig = runtimeConfig with { DataSource = dataSource };

src/Core/Resolvers/MsSqlQueryExecutor.cs

@@ -107,7 +107,7 @@ public override async Task SetManagedIdentityAccessTokenIfAnyAsync(DbConnection
             // using default datasource name for first db - maintaining backward compatibility for single db scenario.
             if (string.IsNullOrEmpty(dataSourceName))
             {
-                dataSourceName = ConfigProvider.GetConfig().GetDefaultDataSourceName();
+                dataSourceName = ConfigProvider.GetConfig().DefaultDataSourceName;
             }
 
             _dataSourceAccessTokenUsage.TryGetValue(dataSourceName, out bool setAccessToken);
@@ -198,7 +198,7 @@ public override string GetSessionParamsQuery(HttpContext? httpContext, IDictiona
         {
             if (string.IsNullOrEmpty(dataSourceName))
             {
-                dataSourceName = ConfigProvider.GetConfig().GetDefaultDataSourceName();
+                dataSourceName = ConfigProvider.GetConfig().DefaultDataSourceName;
             }
 
             if (httpContext is null || !_dataSourceToSessionContextUsage[dataSourceName])

src/Core/Resolvers/MySqlQueryExecutor.cs

@@ -97,7 +97,7 @@ public override async Task SetManagedIdentityAccessTokenIfAnyAsync(DbConnection
             // using default datasource name for first db - maintaining backward compatibility for single db scenario.
             if (string.IsNullOrEmpty(dataSourceName))
             {
-                dataSourceName = ConfigProvider.GetConfig().GetDefaultDataSourceName();
+                dataSourceName = ConfigProvider.GetConfig().DefaultDataSourceName;
             }
 
             _dataSourceAccessTokenUsage.TryGetValue(dataSourceName, out bool setAccessToken);

src/Core/Resolvers/PostgreSqlExecutor.cs

@@ -93,7 +93,7 @@ public override async Task SetManagedIdentityAccessTokenIfAnyAsync(DbConnection
             // using default datasource name for first db - maintaining backward compatibility for single db scenario.
             if (string.IsNullOrEmpty(dataSourceName))
             {
-                dataSourceName = ConfigProvider.GetConfig().GetDefaultDataSourceName();
+                dataSourceName = ConfigProvider.GetConfig().DefaultDataSourceName;
             }
 
             _dataSourceAccessTokenUsage.TryGetValue(dataSourceName, out bool setAccessToken);

src/Core/Resolvers/QueryExecutor.cs

@@ -71,7 +71,7 @@ public QueryExecutor(DbExceptionParser dbExceptionParser,
         {
             if (string.IsNullOrEmpty(dataSourceName))
             {
-                dataSourceName = ConfigProvider.GetConfig().GetDefaultDataSourceName();
+                dataSourceName = ConfigProvider.GetConfig().DefaultDataSourceName;
             }
 
             if (!ConnectionStringBuilders.ContainsKey(dataSourceName))

src/Core/Resolvers/SqlMutationEngine.cs

@@ -373,7 +373,7 @@ await queryExecutor.ExecuteQueryAsync(
         public async Task<IActionResult?> ExecuteAsync(RestRequestContext context)
         {
             // for REST API scenarios, use the default datasource
-            string dataSourceName = _runtimeConfigProvider.GetConfig().GetDefaultDataSourceName();
+            string dataSourceName = _runtimeConfigProvider.GetConfig().DefaultDataSourceName;
 
             Dictionary<string, object?> parameters = PrepareParameters(context);
             ISqlMetadataProvider sqlMetadataProvider = _sqlMetadataProviderFactory.GetMetadataProvider(dataSourceName);
@@ -1194,7 +1194,7 @@ private static TransactionScope ConstructTransactionScopeWithSpecifiedIsolationL
         private string GetValidatedDataSourceName(string dataSourceName)
         {
             // For rest scenarios - no multiple db support. Hence to maintain backward compatibility, we will use the default db.
-            return string.IsNullOrEmpty(dataSourceName) ? _runtimeConfigProvider.GetConfig().GetDefaultDataSourceName() : dataSourceName;
+            return string.IsNullOrEmpty(dataSourceName) ? _runtimeConfigProvider.GetConfig().DefaultDataSourceName : dataSourceName;
         }
 
         /// <summary>

src/Core/Resolvers/SqlQueryEngine.cs

@@ -139,7 +139,7 @@ await ExecuteAsync(structure, dataSourceName),
         public async Task<JsonDocument?> ExecuteAsync(FindRequestContext context)
         {
             // for REST API scenarios, use the default datasource
-            string dataSourceName = _runtimeConfigProvider.GetConfig().GetDefaultDataSourceName();
+            string dataSourceName = _runtimeConfigProvider.GetConfig().DefaultDataSourceName;
 
             ISqlMetadataProvider sqlMetadataProvider = _sqlMetadataProviderFactory.GetMetadataProvider(dataSourceName);
             SqlQueryStructure structure = new(

src/Core/Services/OpenAPI/OpenApiDocumentor.cs

@@ -173,7 +173,7 @@ private OpenApiPaths BuildPaths()
         {
             OpenApiPaths pathsCollection = new();
 
-            string defaultDataSourceName = _runtimeConfig.GetDefaultDataSourceName();
+            string defaultDataSourceName = _runtimeConfig.DefaultDataSourceName;
             ISqlMetadataProvider metadataProvider = _metadataProviderFactory.GetMetadataProvider(defaultDataSourceName);
             foreach (KeyValuePair<string, DatabaseObject> entityDbMetadataMap in metadataProvider.EntityToDatabaseObject)
             {
@@ -952,7 +952,7 @@ private Dictionary<string, OpenApiSchema> CreateComponentSchemas()
             Dictionary<string, OpenApiSchema> schemas = new();
 
             // for rest scenario we need the default datasource name.
-            string defaultDataSourceName = _runtimeConfig.GetDefaultDataSourceName();
+            string defaultDataSourceName = _runtimeConfig.DefaultDataSourceName;
             ISqlMetadataProvider metadataProvider = _metadataProviderFactory.GetMetadataProvider(defaultDataSourceName);
 
             foreach (KeyValuePair<string, DatabaseObject> entityDbMetadataMap in metadataProvider.EntityToDatabaseObject)

src/Core/Services/RestService.cs

@@ -212,7 +212,7 @@ RequestValidator requestValidator
         private async Task<IActionResult> DispatchQuery(RestRequestContext context, DatabaseType databaseType)
         {
             IQueryEngine queryEngine = _queryEngineFactory.GetQueryEngine(databaseType);
-            string defaultDataSourceName = _runtimeConfigProvider.GetConfig().GetDefaultDataSourceName();
+            string defaultDataSourceName = _runtimeConfigProvider.GetConfig().DefaultDataSourceName;
 
             if (context is FindRequestContext findRequestContext)
             {
@@ -237,7 +237,7 @@ private async Task<IActionResult> DispatchQuery(RestRequestContext context, Data
         private Task<IActionResult?> DispatchMutation(RestRequestContext context, DatabaseType databaseType)
         {
             IMutationEngine mutationEngine = _mutationEngineFactory.GetMutationEngine(databaseType);
-            string defaultDataSourceName = _runtimeConfigProvider.GetConfig().GetDefaultDataSourceName();
+            string defaultDataSourceName = _runtimeConfigProvider.GetConfig().DefaultDataSourceName;
             return context switch
             {
                 StoredProcedureRequestContext => mutationEngine.ExecuteAsync((StoredProcedureRequestContext)context, defaultDataSourceName),
@@ -437,7 +437,7 @@ public bool TryGetRestRouteFromConfig([NotNullWhen(true)] out string? configured
         public (string, string) GetEntityNameAndPrimaryKeyRouteFromRoute(string routeAfterPathBase)
         {
 
-            string dataSourceName = _runtimeConfigProvider.GetConfig().GetDefaultDataSourceName();
+            string dataSourceName = _runtimeConfigProvider.GetConfig().DefaultDataSourceName;
             ISqlMetadataProvider sqlMetadataProvider = _sqlMetadataProviderFactory.GetMetadataProvider(dataSourceName);
 
             // Split routeAfterPath on the first occurrence of '/', if we get back 2 elements