Add support for using generic SQL queries to filter EditData rows. (#605)

src/Microsoft.SqlTools.ServiceLayer/EditData/Contracts/EditInitializeRequest.cs

@@ -31,6 +31,11 @@ namespace Microsoft.SqlTools.ServiceLayer.EditData.Contracts
         /// The type of the object to use for generating an edit script
         /// </summary>
         public string ObjectType { get; set; }
+
+        /// <summary>
+        /// The query used to retrieve result set
+        /// </summary>
+        public string QueryString { get; set; }
     }
 
     /// <summary>

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

@@ -146,6 +146,62 @@ namespace Microsoft.SqlTools.ServiceLayer.EditData
             return query.Batches[0].ResultSets[0];
         }
 
+        /// <summary>
+        /// If the results contain any results that conflict with the table metadata, then
+        /// make all columns readonly so that the user cannot make an invalid update.
+        /// </summary>
+        public static void CheckResultsForInvalidColumns(ResultSet results, string tableName)
+        {
+            if(SchemaContainsMultipleItems(results.Columns, col => col.BaseCatalogName)
+                || SchemaContainsMultipleItems(results.Columns, col => col.BaseSchemaName)
+                || SchemaContainsMultipleItems(results.Columns, col => col.BaseTableName))
+            {
+                throw new InvalidOperationException(SR.EditDataMultiTableNotSupported);
+            }
+
+            // Check if any of the columns are invalid
+            HashSet<string> colNameTracker = new HashSet<string>();
+            foreach (DbColumnWrapper col in results.Columns)
+            {
+                if (col.IsAliased.HasTrue())
+                {
+                    throw new InvalidOperationException(SR.EditDataAliasesNotSupported);
+                }
+
+                if (col.IsExpression.HasTrue())
+                {
+                    throw new InvalidOperationException(SR.EditDataExpressionsNotSupported);
+                }
+
+                if (colNameTracker.Contains(col.ColumnName))
+                {
+                    throw new InvalidOperationException(SR.EditDataDuplicateColumnsNotSupported);
+                }
+                else
+                {
+                    colNameTracker.Add(col.ColumnName);
+                }
+            }
+
+            // Only one source table in the metadata, but check if results are from the original table.
+            if (results.Columns.Length > 0)
+            {
+                string resultTableName = results.Columns[0].BaseTableName;
+                if (!string.IsNullOrEmpty(resultTableName) && !string.Equals(resultTableName, tableName))
+                {
+                    throw new InvalidOperationException(SR.EditDataIncorrectTable(tableName));
+                }
+            }
+        }
+
+        private static bool SchemaContainsMultipleItems(DbColumn[] columns, Func<DbColumn, string> filter)
+        {
+            return columns
+                .Select(column => filter(column))
+                .Where(name => name != null)
+                .ToHashSet().Count > 1;
+        }
+
         /// <summary>
         /// Creates a new row update and adds it to the update cache
         /// </summary>
@@ -415,7 +471,7 @@ namespace Microsoft.SqlTools.ServiceLayer.EditData
                     initParams.ObjectType);
 
                 // Step 2) Get and execute a query for the rows in the object we're looking up
-                EditSessionQueryExecutionState state = await queryRunner(ConstructInitializeQuery(objectMetadata, initParams.Filters));
+                EditSessionQueryExecutionState state = await queryRunner(initParams.QueryString ?? ConstructInitializeQuery(objectMetadata, initParams.Filters));
                 if (state.Query == null)
                 {
                     string message = state.Message ?? SR.EditDataQueryFailed;
@@ -424,6 +480,8 @@ namespace Microsoft.SqlTools.ServiceLayer.EditData
 
                 // Step 3) Setup the internal state
                 associatedResultSet = ValidateQueryForSession(state.Query);
+                CheckResultsForInvalidColumns(associatedResultSet, initParams.ObjectName);
+
                 NextRowId = associatedResultSet.RowCount;
                 EditCache = new ConcurrentDictionary<long, RowEditBase>();
                 IsInitialized = true;

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

@@ -3,8 +3,11 @@
 // Licensed under the MIT license. See LICENSE file in the project root for full license information.
 //
 
+using System;
+using System.Collections.Generic;
 using System.Linq;
 using Microsoft.SqlTools.ServiceLayer.QueryExecution.Contracts;
+using Microsoft.SqlTools.ServiceLayer.Utility.SqlScriptFormatters;
 using Microsoft.SqlTools.Utility;
 
 namespace Microsoft.SqlTools.ServiceLayer.EditData
@@ -55,6 +58,40 @@ namespace Microsoft.SqlTools.ServiceLayer.EditData
 
         #endregion
 
+        /// <summary>
+        /// Filters out metadata that is not present in the result set, and matches metadata ordering to resultset.
+        /// </summary>
+        public static EditColumnMetadata[] FilterColumnMetadata(EditColumnMetadata[] metaColumns, DbColumnWrapper[] resultColumns)
+        {
+            if (metaColumns.Length == 0)
+            {
+                return metaColumns;
+            }
+
+            bool escapeColName = FromSqlScript.IsIdentifierBracketed(metaColumns[0].EscapedName);
+            Dictionary<string, int> columnNameOrdinalMap = new Dictionary<string, int>(capacity: resultColumns.Length);
+            for (int i = 0; i < resultColumns.Length; i++)
+            {
+                DbColumnWrapper column = resultColumns[i];
+                string columnName = column.ColumnName;
+                if (escapeColName && !FromSqlScript.IsIdentifierBracketed(columnName))
+                {
+                    columnName = ToSqlScript.FormatIdentifier(columnName);
+                }
+                columnNameOrdinalMap.Add(columnName, column.ColumnOrdinal ?? i);
+            }
+
+            HashSet<string> resultColumnNames = columnNameOrdinalMap.Keys.ToHashSet();
+            metaColumns = Array.FindAll(metaColumns, column => resultColumnNames.Contains(column.EscapedName));
+            foreach (EditColumnMetadata metaCol in metaColumns)
+            {
+                metaCol.Ordinal = columnNameOrdinalMap[metaCol.EscapedName];
+            }
+            Array.Sort(metaColumns, (x, y) => (Comparer<int>.Default).Compare(x.Ordinal, y.Ordinal));
+
+            return metaColumns;
+        }
+
         /// <summary>
         /// Extracts extended column properties from the database columns from SQL Client
         /// </summary>
@@ -63,6 +100,8 @@ namespace Microsoft.SqlTools.ServiceLayer.EditData
         {
             Validate.IsNotNull(nameof(dbColumnWrappers), dbColumnWrappers);
 
+            Columns = EditTableMetadata.FilterColumnMetadata(Columns, dbColumnWrappers);
+
             // Iterate over the column wrappers and improve the columns we have
             for (int i = 0; i < Columns.Length; i++)
             {

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

@@ -42,10 +42,10 @@ namespace Microsoft.SqlTools.ServiceLayer.EditData.UpdateManagement
         public RowCreate(long rowId, ResultSet associatedResultSet, EditTableMetadata associatedMetadata)
             : base(rowId, associatedResultSet, associatedMetadata)
         {
-            newCells = new CellUpdate[associatedResultSet.Columns.Length];
+            newCells = new CellUpdate[AssociatedResultSet.Columns.Length];
             
             // Process the default cell values. If the column is calculated, then the value is a placeholder
-            DefaultValues = associatedMetadata.Columns.Select((col, index) => col.IsCalculated.HasTrue()
+            DefaultValues = AssociatedObjectMetadata.Columns.Select((col, index) => col.IsCalculated.HasTrue()
                 ? SR.EditDataComputedColumnPlaceholder
                 : col.DefaultValue).ToArray();
         }

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

@@ -45,8 +45,10 @@ namespace Microsoft.SqlTools.ServiceLayer.EditData.UpdateManagement
             }
 
             RowId = rowId;
-            AssociatedResultSet = associatedResultSet;
             AssociatedObjectMetadata = associatedMetadata;
+            AssociatedResultSet = associatedResultSet;
+
+            AssociatedObjectMetadata.Columns = EditTableMetadata.FilterColumnMetadata(AssociatedObjectMetadata.Columns, AssociatedResultSet.Columns);
         }
 
         #region Properties

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

@@ -44,7 +44,7 @@ namespace Microsoft.SqlTools.ServiceLayer.EditData.UpdateManagement
             : base(rowId, associatedResultSet, associatedMetadata)
         {
             cellUpdates = new ConcurrentDictionary<int, CellUpdate>();
-            associatedRow = associatedResultSet.GetRow(rowId);
+            associatedRow = AssociatedResultSet.GetRow(rowId);
         }
 
         /// <summary>