From 822213a655a66fcec4515e614d00c538e014dd94 Mon Sep 17 00:00:00 2001
From: Justin M <63619224+JustinMDotNet@users.noreply.github.com>
Date: Tue, 10 Aug 2021 11:00:45 -0700
Subject: [PATCH] AzureMonitor Fix Dashboard Tables (#1232)

* Changed results for AzureMonitor GetChildObjects for Dashboard

* Changed MonitorDataSource to use SortedDictionary instead of List when storing nodes.
---
 .../DataSource/KustoQueryUtils.cs              | 14 ++++++++++----
 .../DataSource/Monitor/MonitorDataSource.cs    | 18 ++++++++++--------
 2 files changed, 20 insertions(+), 12 deletions(-)

diff --git a/src/Microsoft.Kusto.ServiceLayer/DataSource/KustoQueryUtils.cs b/src/Microsoft.Kusto.ServiceLayer/DataSource/KustoQueryUtils.cs
index fcf1e796..a7d8da6e 100644
--- a/src/Microsoft.Kusto.ServiceLayer/DataSource/KustoQueryUtils.cs
+++ b/src/Microsoft.Kusto.ServiceLayer/DataSource/KustoQueryUtils.cs
@@ -91,16 +91,22 @@ namespace Microsoft.Kusto.ServiceLayer.DataSource
                 dictionary[key] = new Dictionary<string, T> {{metadata.Name, metadata}};
             }
         }
-        
-        public static void SafeAdd(this Dictionary<string, List<DataSourceObjectMetadata>> dictionary, string key, DataSourceObjectMetadata node)
+
+        public static void SafeAdd<T>(this Dictionary<string, SortedDictionary<string, DataSourceObjectMetadata>> dictionary, string key,
+            T node) where T : DataSourceObjectMetadata
         {
             if (dictionary.ContainsKey(key))
             {
-                dictionary[key].Add(node);
+                if (dictionary[key].ContainsKey(node.PrettyName))
+                {
+                    return;
+                }
+
+                dictionary[key].Add(node.PrettyName, node);
             }
             else
             {
-                dictionary[key] = new List<DataSourceObjectMetadata> {node};
+                dictionary[key] = new SortedDictionary<string, DataSourceObjectMetadata> {{node.PrettyName, node}};
             }
         }
         
diff --git a/src/Microsoft.Kusto.ServiceLayer/DataSource/Monitor/MonitorDataSource.cs b/src/Microsoft.Kusto.ServiceLayer/DataSource/Monitor/MonitorDataSource.cs
index a6b7bd50..b32fd428 100644
--- a/src/Microsoft.Kusto.ServiceLayer/DataSource/Monitor/MonitorDataSource.cs
+++ b/src/Microsoft.Kusto.ServiceLayer/DataSource/Monitor/MonitorDataSource.cs
@@ -22,7 +22,8 @@ namespace Microsoft.Kusto.ServiceLayer.DataSource.Monitor
         private readonly MonitorClient _monitorClient;
         private readonly IntellisenseClientBase _intellisenseClient;
         private WorkspaceResponse _metadata;
-        private Dictionary<string, List<DataSourceObjectMetadata>> _nodes;
+        private Dictionary<string, SortedDictionary<string, DataSourceObjectMetadata>> _nodes;
+        private const string DatabaseKeyPrefix = "OnlyTables";
         
         public override string ClusterName => _monitorClient.WorkspaceId;
         public override string DatabaseName { get; set; }
@@ -31,7 +32,7 @@ namespace Microsoft.Kusto.ServiceLayer.DataSource.Monitor
         {
             _monitorClient = monitorClient;
             _intellisenseClient = intellisenseClient;
-            _nodes = new Dictionary<string, List<DataSourceObjectMetadata>>();
+            _nodes = new Dictionary<string, SortedDictionary<string, DataSourceObjectMetadata>>(StringComparer.OrdinalIgnoreCase);
             _metadata = _monitorClient.LoadMetadata();
             DataSourceType = DataSourceType.LogAnalytics;
             SetupTableGroups(monitorClient.WorkspaceId);
@@ -52,11 +53,11 @@ namespace Microsoft.Kusto.ServiceLayer.DataSource.Monitor
 
                 _nodes.SafeAdd($"{workspace.Id}", tableGroupNodeInfo);
                 
-                SetupTables(tableGroupNodeInfo, workspaceTableGroup);
+                SetupTables(tableGroupNodeInfo, workspaceTableGroup, workspace.Id);
             }
         }
         
-        private void SetupTables(DataSourceObjectMetadata tableGroupNodeInfo, TableGroupsModel workspaceTableGroup)
+        private void SetupTables(DataSourceObjectMetadata tableGroupNodeInfo, TableGroupsModel workspaceTableGroup, string workspaceId)
         {
             var tableGroupTables = _metadata.Tables.Where(x => workspaceTableGroup.Tables.Contains(x.Id));
             
@@ -66,6 +67,7 @@ namespace Microsoft.Kusto.ServiceLayer.DataSource.Monitor
                     $"{tableGroupNodeInfo.Urn}.{metadataTable.Name}");
 
                 _nodes.SafeAdd(tableGroupNodeInfo.Urn, tableNodeInfo);
+                _nodes.SafeAdd($"{DatabaseKeyPrefix}.{workspaceId}", tableNodeInfo);
 
                 SetupColumns(metadataTable, tableNodeInfo);
             }
@@ -108,17 +110,17 @@ namespace Microsoft.Kusto.ServiceLayer.DataSource.Monitor
 
             if (parentMetadata.MetadataType == DataSourceMetadataType.Cluster && includeSizeDetails)
             {
-                var child = _nodes[parentMetadata.Urn].FirstOrDefault();
-                return child == null ? Enumerable.Empty<DataSourceObjectMetadata>() : _nodes[child.Urn];
+                string newKey = $"{DatabaseKeyPrefix}.{parentMetadata.Urn}";
+                return _nodes[newKey].Values;
             }
             
-            return _nodes[parentMetadata.Urn].OrderBy(x => x.PrettyName, StringComparer.OrdinalIgnoreCase);
+            return _nodes[parentMetadata.Urn].Values;
         }
 
         public override void Refresh(bool includeDatabase)
         {
             // reset the data source
-            _nodes = new Dictionary<string, List<DataSourceObjectMetadata>>();
+            _nodes = new Dictionary<string, SortedDictionary<string, DataSourceObjectMetadata>>(StringComparer.OrdinalIgnoreCase);
             _metadata = _monitorClient.LoadMetadata();
             SetupTableGroups(_monitorClient.WorkspaceId);
         }