Handling the HierarchyId for edit data scenario (#709)

* Handling the HierarchyId for edit data scenario

* comments

src/Microsoft.SqlTools.ServiceLayer/EditData/EditColumnMetadata.cs

@@ -36,6 +36,22 @@ namespace Microsoft.SqlTools.ServiceLayer.EditData
         /// </summary>
         public string EscapedName { get; set; }
 
+        /// <summary>
+        /// The column's expression for select statement
+        /// </summary>
+        public string ExpressionForSelectStatement
+        {
+            get
+            {
+                return IsHierarchyId ? string.Format("{0}.ToString() AS {0}", EscapedName) : EscapedName;
+            }
+        }
+
+        /// <summary>
+        /// Whether or not the column's data type is HierarchyId
+        /// </summary>
+        public bool IsHierarchyId { get; set; }
+
         /// <summary>
         /// Whether or not the column is computed
         /// </summary>

src/Microsoft.SqlTools.ServiceLayer/EditData/EditDataService.cs

@@ -8,6 +8,7 @@ using System.Collections.Concurrent;
 using System.Data.Common;
 using System.Threading.Tasks;
 using Microsoft.SqlTools.Hosting.Protocol;
+using Microsoft.SqlTools.Hosting.Protocol.Contracts;
 using Microsoft.SqlTools.ServiceLayer.Connection;
 using Microsoft.SqlTools.ServiceLayer.EditData.Contracts;
 using Microsoft.SqlTools.ServiceLayer.Hosting;
@@ -149,11 +150,12 @@ namespace Microsoft.SqlTools.ServiceLayer.EditData
         internal async Task HandleInitializeRequest(EditInitializeParams initParams,
             RequestContext<EditInitializeResult> requestContext)
         {
-            Func<Exception, Task> executionFailureHandler = (e) => SendSessionReadyEvent(requestContext, initParams.OwnerUri, false, e.Message);
+            InitializeEditRequestContext context = new InitializeEditRequestContext(requestContext);
-            Func<Task> executionSuccessHandler = () => SendSessionReadyEvent(requestContext, initParams.OwnerUri, true, null);
+            Func<Exception, Task> executionFailureHandler = (e) => SendSessionReadyEvent(context, initParams.OwnerUri, false, e.Message);
+            Func<Task> executionSuccessHandler = () => SendSessionReadyEvent(context, initParams.OwnerUri, true, null);
 
             EditSession.Connector connector = () => connectionService.GetOrOpenConnection(initParams.OwnerUri, ConnectionType.Edit, alwaysPersistSecurity: true);
-            EditSession.QueryRunner queryRunner = q => SessionInitializeQueryRunner(initParams.OwnerUri, requestContext, q);
+            EditSession.QueryRunner queryRunner = q => SessionInitializeQueryRunner(initParams.OwnerUri, context, q);
 
             try
             {
@@ -169,6 +171,8 @@ namespace Microsoft.SqlTools.ServiceLayer.EditData
                     throw new InvalidOperationException(SR.EditDataSessionAlreadyExists);
                 }
 
+                context.ResultSetHandler = (ResultSetEventParams resultSetEventParams) => { session.UpdateColumnInformationWithMetadata(resultSetEventParams.ResultSetSummary.ColumnInfo); };
+
                 // Initialize the session
                 session.Initialize(initParams, connector, queryRunner, executionSuccessHandler, executionFailureHandler);
 
@@ -341,4 +345,29 @@ namespace Microsoft.SqlTools.ServiceLayer.EditData
         #endregion
 
     }
+
+    /// <summary>
+    /// Context for InitializeEditRequest, to provide a way to update the result set before sending it to UI.
+    /// </summary>
+    internal class InitializeEditRequestContext : IEventSender
+    {
+        private RequestContext<EditInitializeResult> _context;
+
+        public Action<ResultSetEventParams> ResultSetHandler { get; set; }
+
+        public InitializeEditRequestContext(RequestContext<EditInitializeResult> context)
+        {
+            this._context = context;
+        }
+
+        public Task SendEvent<TParams>(EventType<TParams> eventType, TParams eventParams)
+        {
+            if (eventParams is ResultSetEventParams && this.ResultSetHandler != null)
+            {
+                this.ResultSetHandler(eventParams as ResultSetEventParams);
+            }
+            return _context.SendEvent(eventType, eventParams);
+        }
+
+    }
 }

src/Microsoft.SqlTools.ServiceLayer/EditData/EditSession.cs

@@ -168,7 +168,8 @@ namespace Microsoft.SqlTools.ServiceLayer.EditData
                     throw new InvalidOperationException(SR.EditDataAliasesNotSupported);
                 }
 
-                if (col.IsExpression.HasTrue())
+                // We have changed HierarchyId column to an expression so that it can be displayed properly
+                if (!col.IsHierarchyId && col.IsExpression.HasTrue())
                 {
                     throw new InvalidOperationException(SR.EditDataExpressionsNotSupported);
                 }
@@ -480,6 +481,7 @@ namespace Microsoft.SqlTools.ServiceLayer.EditData
 
                 // Step 3) Setup the internal state
                 associatedResultSet = ValidateQueryForSession(state.Query);
+                UpdateColumnInformationWithMetadata(associatedResultSet.Columns);
                 CheckResultsForInvalidColumns(associatedResultSet, initParams.ObjectName);
 
                 NextRowId = associatedResultSet.RowCount;
@@ -554,7 +556,7 @@ namespace Microsoft.SqlTools.ServiceLayer.EditData
             }
 
             // Using the columns we know, add them to the query
-            var columns = metadata.Columns.Select(col => col.EscapedName);
+            var columns = metadata.Columns.Select(col => col.ExpressionForSelectStatement);
             var columnClause = string.Join(", ", columns);
             queryBuilder.Append(columnClause);
 
@@ -590,6 +592,20 @@ namespace Microsoft.SqlTools.ServiceLayer.EditData
             EditCache.TryRemove(rowId, out removedRow);
         }
 
+        internal void UpdateColumnInformationWithMetadata(DbColumnWrapper[] columns)
+        {
+            if (columns == null || this.objectMetadata == null)
+            {
+                return;
+            }
+
+            foreach (DbColumnWrapper col in columns)
+            {
+                var columnMetadata = objectMetadata.Columns.FirstOrDefault(cm => { return cm.EscapedName == ToSqlScript.FormatIdentifier(col.ColumnName); });
+                col.IsHierarchyId = columnMetadata != null && columnMetadata.IsHierarchyId;
+            }
+        }
+
         #endregion
 
         /// <summary>

src/Microsoft.SqlTools.ServiceLayer/EditData/SmoEditMetadataFactory.cs

@@ -118,6 +118,7 @@ namespace Microsoft.SqlTools.ServiceLayer.EditData
                     DefaultValue = defaultValue,
                     EscapedName = ToSqlScript.FormatIdentifier(smoColumn.Name),
                     Ordinal = i,
+                    IsHierarchyId = smoColumn.DataType.SqlDataType == SqlDataType.HierarchyId,
                 };
                 editColumns.Add(column);
             }

src/Microsoft.SqlTools.ServiceLayer/EditData/UpdateManagement/RowCreate.cs

@@ -103,7 +103,14 @@ namespace Microsoft.SqlTools.ServiceLayer.EditData.UpdateManagement
                 CellUpdate cell = newCells[i];
 
                 // Add the output columns regardless of whether the column is read only
+                if (metadata.IsHierarchyId)
+                {
+                    outClauseColumnNames.Add($"inserted.{metadata.EscapedName}.ToString() {metadata.EscapedName}");
+                }
+                else
+                {
                     outClauseColumnNames.Add($"inserted.{metadata.EscapedName}");
+                }
                 declareColumns.Add($"{metadata.EscapedName} {ToSqlScript.FormatColumnType(column, useSemanticEquivalent: true)}");
                 selectColumns.Add(metadata.EscapedName);
 
@@ -118,7 +125,16 @@ namespace Microsoft.SqlTools.ServiceLayer.EditData.UpdateManagement
 
                 // Add the input values as parameters
                 string paramName = $"@Value{RowId}_{i}";
+
+                if (metadata.IsHierarchyId)
+                {
+                    inValues.Add($"CONVERT(hierarchyid,{paramName})");
+                }
+                else
+                {
                     inValues.Add(paramName);
+                }
+
                 inParameters.Add(new SqlParameter(paramName, column.SqlDbType) { Value = cell.Value });
             }
 

src/Microsoft.SqlTools.ServiceLayer/EditData/UpdateManagement/RowUpdate.cs

@@ -100,7 +100,14 @@ namespace Microsoft.SqlTools.ServiceLayer.EditData.UpdateManagement
 
                 // Add the output columns regardless of whether the column is read only
                 declareColumns.Add($"{metadata.EscapedName} {ToSqlScript.FormatColumnType(metadata.DbColumn, useSemanticEquivalent: true)}");
+                if (metadata.IsHierarchyId)
+                {
+                    outClauseColumns.Add($"inserted.{metadata.EscapedName}.ToString() {metadata.EscapedName}");
+                }
+                else
+                {
                     outClauseColumns.Add($"inserted.{metadata.EscapedName}");
+                }
                 selectColumns.Add(metadata.EscapedName);
 
                 // If we have a new value for the column, proccess it now
@@ -108,7 +115,14 @@ namespace Microsoft.SqlTools.ServiceLayer.EditData.UpdateManagement
                 if (cellUpdates.TryGetValue(i, out cellUpdate))
                 {
                     string paramName = $"@Value{RowId}_{i}";
+                    if (metadata.IsHierarchyId)
+                    {
+                        setComponents.Add($"{metadata.EscapedName} = CONVERT(hierarchyid,{paramName})");
+                    }
+                    else
+                    {
                         setComponents.Add($"{metadata.EscapedName} = {paramName}");
+                    }
                     inParameters.Add(new SqlParameter(paramName, AssociatedResultSet.Columns[i].SqlDbType) { Value = cellUpdate.Value });
                 }
             }

src/Microsoft.SqlTools.ServiceLayer/QueryExecution/Contracts/DbColumnWrapper.cs

@@ -199,8 +199,8 @@ namespace Microsoft.SqlTools.ServiceLayer.QueryExecution.Contracts
                 object assemblyQualifiedName = column.UdtAssemblyQualifiedName;
                 const string hierarchyId = "MICROSOFT.SQLSERVER.TYPES.SQLHIERARCHYID";
 
-                if (assemblyQualifiedName != null &&
+                if (assemblyQualifiedName != null
-                    string.Equals(assemblyQualifiedName.ToString(), hierarchyId, StringComparison.OrdinalIgnoreCase))
+                && assemblyQualifiedName.ToString().StartsWith(hierarchyId, StringComparison.OrdinalIgnoreCase))
                 {
                     DataType = typeof(SqlBinary);
                 }
@@ -259,6 +259,11 @@ namespace Microsoft.SqlTools.ServiceLayer.QueryExecution.Contracts
         /// </summary>
         public SqlDbType SqlDbType { get; private set; }
 
+        /// <summary>
+        /// Whther this is a HierarchyId column
+        /// </summary>
+        public bool IsHierarchyId { get; set; }
+
         /// <summary>
         /// Whether or not the column is an XML Reader type.
         /// </summary>
@@ -286,10 +291,11 @@ namespace Microsoft.SqlTools.ServiceLayer.QueryExecution.Contracts
         /// </summary>
         /// <remarks>
         /// Logic taken from SSDT determination of updatable columns
+        /// Special treatment for HierarchyId since we are using an Expression for HierarchyId column and expression column is readonly.
         /// </remarks>
-        public bool IsUpdatable => !IsAutoIncrement.HasTrue() && 
+        public bool IsUpdatable => (!IsAutoIncrement.HasTrue() &&
                                    !IsReadOnly.HasTrue() &&
-                                   !IsSqlXmlType;
+                                   !IsSqlXmlType) || IsHierarchyId;
 
         #endregion