Fix scripting of users on SQL DB (#2014)

* Fix scripting of users on SQL DB

* Minor typos

* Add null checks

src/Microsoft.SqlTools.ServiceLayer/Management/Common/ManagementActionBase.cs

@@ -31,6 +31,17 @@ namespace Microsoft.SqlTools.ServiceLayer.Management
         /// </summary>
         private ExecutionMode m_executionMode = ExecutionMode.Success;
 
+        /// <summary>
+        /// indicates that this is a database level operation, which on SQL DB will require an alternate 
+        /// execution manager to capture sql for scripting
+        /// </summary>
+        private bool isDatabaseOperation = false;
+
+        /// <summary>
+        /// Parent database object to use for database operations
+        /// </summary>
+        private Database parentDb = null;
+
         /// <summary>
         /// data container with initialization-related information
         /// </summary>
@@ -255,6 +266,26 @@ namespace Microsoft.SqlTools.ServiceLayer.Management
             }
         }
 
+        protected bool IsDatabaseOperation
+        {
+            get
+            {
+                return this.isDatabaseOperation;
+            }
+            set
+            {
+                this.isDatabaseOperation = value;
+            }
+        }
+
+        protected Database ParentDb
+        {
+            get
+            {
+                return this.parentDb;
+            }
+        }
+
         /// <summary>
         /// SMO Server connection that MUST be used for all enumerator calls
         /// We'll get this object out of CDataContainer, that must be initialized
@@ -385,6 +416,13 @@ namespace Microsoft.SqlTools.ServiceLayer.Management
                 sc = sqlDialogSubject.ExecutionManager.ConnectionContext.CapturedSql.Text;
             }
 
+            // for SQL DB database-level operations the script could be on the database object's execution manager
+            if (sc.Count == 0 && this.isDatabaseOperation && this.parentDb != null
+                && this.DataContainer.Server.ServerType == DatabaseEngineType.SqlAzureDatabase)
+            {
+                sc = this.parentDb.ExecutionManager.ConnectionContext.CapturedSql.Text;
+            }                
+
             StringBuilder script = new StringBuilder(4096);
             if (sc != null)
             {
@@ -452,6 +490,15 @@ namespace Microsoft.SqlTools.ServiceLayer.Management
                     sqlDialogSubject.ExecutionManager.ConnectionContext.SqlExecutionModes = newMode;
                 }
 
+                if (this.isDatabaseOperation)
+                {
+                    this.parentDb = this.DataContainer.Server.GetSmoObject(this.DataContainer.ParentUrn) as Database;
+                    if (this.parentDb != null && this.DataContainer.Server.ServerType == DatabaseEngineType.SqlAzureDatabase)
+                    {
+                        this.parentDb.ExecutionManager.ConnectionContext.SqlExecutionModes = newMode;
+                    }
+                }
+
                 executionResult = DoPreProcessExecutionAndRunViews(executionInfo.RunType);
 
                 if (isScripting)
@@ -465,10 +512,18 @@ namespace Microsoft.SqlTools.ServiceLayer.Management
             finally
             {
                 GetServerConnectionForScript().SqlExecutionModes = executionModeOriginal;
+                if (this.parentDb != null && this.isDatabaseOperation && this.DataContainer.Server.ServerType == DatabaseEngineType.SqlAzureDatabase)
+                {
+                    this.parentDb.ExecutionManager.ConnectionContext.SqlExecutionModes = executionModeOriginal;
+                }
 
                 if (isScripting)
                 {
                     GetServerConnectionForScript().CapturedSql.Clear();
+                    if (this.parentDb != null && this.isDatabaseOperation && this.DataContainer.Server.ServerType == DatabaseEngineType.SqlAzureDatabase)
+                    {
+                        this.parentDb.ExecutionManager.ConnectionContext.CapturedSql.Clear();
+                    }
                 }
 
                 if (sqlDialogSubject != null)

src/Microsoft.SqlTools.ServiceLayer/Security/UserActions.cs

@@ -441,6 +441,7 @@ namespace Microsoft.SqlTools.ServiceLayer.Security
             UserPrototypeData? originalData)
         {
             this.DataContainer = dataContainer;
+            this.IsDatabaseOperation = true;
             this.configAction = configAction;
 
             ExhaustiveUserTypes currentUserType;
@@ -475,7 +476,7 @@ namespace Microsoft.SqlTools.ServiceLayer.Security
         {
             if (this.configAction != ConfigAction.Drop)
             {
-                this.userPrototype.ApplyChanges();
+                this.userPrototype.ApplyChanges(this.ParentDb);
             }
         }
 

src/Microsoft.SqlTools.ServiceLayer/Security/UserData.cs

@@ -550,9 +550,9 @@ namespace Microsoft.SqlTools.ServiceLayer.Security
             }
         }
 
-        public User ApplyChanges()
+        public User ApplyChanges(Database parentDb)
         {
-            User user = this.GetUser();
+            User user = this.GetUser(parentDb);
 
             if (this.ChangesExist())
             {
@@ -564,10 +564,10 @@ namespace Microsoft.SqlTools.ServiceLayer.Security
                 //it will again generate the script corresponding to that.
                 user.Refresh();
 
-                this.ApplySchemaOwnershipChanges(user);
+                this.ApplySchemaOwnershipChanges(parentDb, user);
                 this.IsSchemaOwnershipChangesApplied = true;
 
-                this.ApplyRoleMembershipChanges(user);
+                this.ApplyRoleMembershipChanges(parentDb, user);
                 this.IsRoleMembershipChangesApplied = true;
             }
 
@@ -586,7 +586,7 @@ namespace Microsoft.SqlTools.ServiceLayer.Security
             }
         }
 
-        private void ApplySchemaOwnershipChanges(User user)
+        private void ApplySchemaOwnershipChanges(Database parentDb, User user)
         {
             IEnumerator<KeyValuePair<string, bool>>? enumerator = this.currentState.isSchemaOwned?.GetEnumerator();
             if (enumerator != null)
@@ -604,7 +604,7 @@ namespace Microsoft.SqlTools.ServiceLayer.Security
                     {
                         System.Diagnostics.Debug.Assert(!this.Exists || userIsOwner, "shouldn't have to unset ownership for new users");
 
-                        Schema schema = this.parent.Schemas[schemaName];
+                        Schema schema = parentDb.Schemas[schemaName];
                         schema.Owner = userIsOwner ? user.Name : nullString;
                         schema.Alter();
                     }
@@ -612,7 +612,7 @@ namespace Microsoft.SqlTools.ServiceLayer.Security
             }
         }
 
-        private void ApplyRoleMembershipChanges(User user)
+        private void ApplyRoleMembershipChanges(Database parentDb, User user)
         {
             IEnumerator<KeyValuePair<string, bool>>? enumerator = this.currentState.isMember?.GetEnumerator();
             if (enumerator != null)
@@ -628,7 +628,7 @@ namespace Microsoft.SqlTools.ServiceLayer.Security
                     {
                         System.Diagnostics.Debug.Assert(this.Exists || userIsMember, "shouldn't have to unset membership for new users");
 
-                        DatabaseRole role = this.parent.Roles[roleName];
+                        DatabaseRole role = parentDb.Roles[roleName];
 
                         if (userIsMember)
                         {
@@ -663,14 +663,14 @@ namespace Microsoft.SqlTools.ServiceLayer.Security
             }
         }
 
-        public User GetUser()
+        public User GetUser(Database parentDb)
         {
             User result;
 
             // if we think we exist, get the SMO user object
             if (this.Exists)
             {
-                result = this.parent.Users[this.originalState.name];
+                result = parentDb.Users[this.originalState.name];
                 result?.Refresh();
 
                 System.Diagnostics.Debug.Assert(0 == string.Compare(this.originalState.name, this.currentState.name, StringComparison.Ordinal), "name of existing user has changed");
@@ -681,7 +681,7 @@ namespace Microsoft.SqlTools.ServiceLayer.Security
             }
             else
             {
-                result = new User(this.parent, this.Name);
+                result = new User(parentDb, this.Name);
             }
 
             return result;
@@ -824,8 +824,10 @@ namespace Microsoft.SqlTools.ServiceLayer.Security
         {
             get
             {
+                Database? parentDb = this.context.Server.GetSmoObject(this.context.ParentUrn) as Database;
+                User user = this.GetUser(parentDb);
+
                 // Default Schema was not supported before Denali for windows group.
-                User user = this.GetUser();
                 if (this.Exists && user.LoginType == Microsoft.SqlServer.Management.Smo.LoginType.WindowsGroup)
                 {
                     return SqlMgmtUtils.IsSql11OrLater(this.context.Server.ConnectionContext.ServerVersion);