From 37469dcbbd316cb0e152f83ce353ffd751fb8de9 Mon Sep 17 00:00:00 2001
From: Hai Cao <hacao@microsoft.com>
Date: Fri, 2 Jun 2023 10:59:32 -0700
Subject: [PATCH] Fix schema level search and refactor search code for user
 management (#2084)

---
 .../ObjectManagementService.cs                | 50 ++++++++++++-------
 .../ObjectManagement/SecurableUtils.cs        | 15 ++++++
 2 files changed, 48 insertions(+), 17 deletions(-)

diff --git a/src/Microsoft.SqlTools.ServiceLayer/ObjectManagement/ObjectManagementService.cs b/src/Microsoft.SqlTools.ServiceLayer/ObjectManagement/ObjectManagementService.cs
index 8014a719..71043b7e 100644
--- a/src/Microsoft.SqlTools.ServiceLayer/ObjectManagement/ObjectManagementService.cs
+++ b/src/Microsoft.SqlTools.ServiceLayer/ObjectManagement/ObjectManagementService.cs
@@ -12,6 +12,8 @@ using Microsoft.SqlTools.ServiceLayer.ObjectManagement.Contracts;
 using System.Collections.Generic;
 using System.Collections.Concurrent;
 using Microsoft.SqlTools.ServiceLayer.Management;
+using System.Linq;
+using Microsoft.SqlTools.ServiceLayer.ObjectManagement.PermissionsData;
 
 namespace Microsoft.SqlTools.ServiceLayer.ObjectManagement
 {
@@ -131,6 +133,7 @@ namespace Microsoft.SqlTools.ServiceLayer.ObjectManagement
             CDataContainer dataContainer = CDataContainer.CreateDataContainer(connInfo, databaseExists: true);
 
             List<SearchResultItem> res = new List<SearchResultItem>();
+            var schemaTypes = SecurableUtils.GetSchemaTypes(dataContainer.Server).Select(Securable.GetSearchableObjectType).ToHashSet();
 
             foreach (string type in requestParams.ObjectTypes)
             {
@@ -142,31 +145,44 @@ namespace Microsoft.SqlTools.ServiceLayer.ObjectManagement
                     continue;
                 }
 
-                SearchableObjectTypeDescription desc = SearchableObjectTypeDescription.GetDescription(searchableObjectType);
-
-                if (requestParams.SearchText != null)
+                // only schema-Scoped securableTypes support schema level search
+                if (!string.IsNullOrEmpty(requestParams.Schema) && !schemaTypes.Contains(searchableObjectType))
                 {
-                    if (desc.IsDatabaseObject)
+                    continue;
+                }
+
+                SearchableObjectTypeDescription desc = SearchableObjectTypeDescription.GetDescription(searchableObjectType);
+                
+                if (desc.IsDatabaseObject)
+                {
+                    if (!string.IsNullOrEmpty(requestParams.Schema))
                     {
-                        SearchableObject.Search(result, searchableObjectType, dataContainer.ConnectionInfo, context.Parameters.Database, requestParams.SearchText, false, true);
+                        SearchableObject.Search(result, searchableObjectType, dataContainer.ConnectionInfo, context.Parameters.Database, requestParams.SearchText ?? string.Empty, false, requestParams.Schema, true, true);
+                    }
+                    else
+                    {
+                        if (!string.IsNullOrEmpty(requestParams.SearchText))
+                        {
+                            SearchableObject.Search(result, searchableObjectType, dataContainer.ConnectionInfo, context.Parameters.Database, requestParams.SearchText, false, true);
+                        }
+                        else
+                        {
+                            SearchableObject.Search(result, searchableObjectType, dataContainer.ConnectionInfo, context.Parameters.Database, true);
+                        }
+                    }
+                }
+                else
+                {
+                    // server object
+                    if (string.IsNullOrEmpty(requestParams.SearchText))
+                    {
+                        SearchableObject.Search(result, searchableObjectType, dataContainer.ConnectionInfo, true);
                     }
                     else
                     {
                         SearchableObject.Search(result, searchableObjectType, dataContainer.ConnectionInfo, requestParams.SearchText, false, true);
                     }
                 }
-                else
-                {
-                    if (desc.IsDatabaseObject)
-                    {
-                        SearchableObject.Search(result, searchableObjectType, dataContainer.ConnectionInfo, context.Parameters.Database, true);
-                    }
-                    else
-                    {
-                        SearchableObject.Search(result, searchableObjectType, dataContainer.ConnectionInfo, true);
-                    }
-
-                }
 
                 foreach (SearchableObject obj in result)
                 {
diff --git a/src/Microsoft.SqlTools.ServiceLayer/ObjectManagement/SecurableUtils.cs b/src/Microsoft.SqlTools.ServiceLayer/ObjectManagement/SecurableUtils.cs
index efe451c6..ebbeb367 100644
--- a/src/Microsoft.SqlTools.ServiceLayer/ObjectManagement/SecurableUtils.cs
+++ b/src/Microsoft.SqlTools.ServiceLayer/ObjectManagement/SecurableUtils.cs
@@ -582,5 +582,20 @@ namespace Microsoft.SqlTools.ServiceLayer.ObjectManagement
                 }
             }
         }
+
+        /// <summary>
+        /// The list of Schema-Scoped SecurableTypes
+        /// </summary>
+        internal static IEnumerable<SecurableType> GetSchemaTypes(Server server)
+        {
+            return Enum.GetValues(typeof(SecurableType))
+                .Cast<SecurableType>()
+                .Where(
+                    t =>
+                        t.IsValidSchemaBoundSecurable(
+                        new ServerVersion(server.ServerVersion.Major, server.ServerVersion.Minor, server.ServerVersion.BuildNumber),
+                            server.DatabaseEngineEdition,
+                            server.DatabaseEngineType));
+        }
     }
 }
\ No newline at end of file