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>

src/Microsoft.SqlTools.ServiceLayer/Localization/sr.strings

@@ -236,6 +236,16 @@ EditDataNullNotAllowed = NULL is not allowed for this column
 
 EditDataValueTooLarge(string value, string columnType) = Value {0} is too large to fit in column of type {1}
 
+EditDataMultiTableNotSupported = EditData queries targeting multiple tables are not supported
+
+EditDataAliasesNotSupported = EditData queries with aliased columns are not supported
+
+EditDataExpressionsNotSupported = EditData queries with aggregate or expression columns are not supported
+
+EditDataDuplicateColumnsNotSupported = EditData queries with duplicate columns are not supported
+
+EditDataIncorrectTable(string tableName) = EditData queries must query the originally targeted table '{0}'
+
 ############################################################################
 # DacFx Resources
 

src/Microsoft.SqlTools.ServiceLayer/QueryExecution/Batch.cs

@@ -16,6 +16,7 @@ using Microsoft.SqlTools.ServiceLayer.QueryExecution.Contracts;
 using Microsoft.SqlTools.ServiceLayer.QueryExecution.DataStorage;
 using Microsoft.SqlTools.Utility;
 using System.Globalization;
+using System.Collections.ObjectModel;
 
 namespace Microsoft.SqlTools.ServiceLayer.QueryExecution
 {
@@ -339,6 +340,21 @@ namespace Microsoft.SqlTools.ServiceLayer.QueryExecution
                 dbCommand.CommandTimeout = 0;
                 executionStartTime = DateTime.Now;
 
+                // Fetch schema info separately, since CommandBehavior.KeyInfo will include primary
+                // key columns in the result set, even if they weren't part of the select statement.
+                // Extra key columns get added to the end, so just correlate via Column Ordinal.
+                List<DbColumn[]> columnSchemas = new List<DbColumn[]>();
+                using (DbDataReader reader = await dbCommand.ExecuteReaderAsync(CommandBehavior.KeyInfo | CommandBehavior.SchemaOnly, cancellationToken))
+                {
+                    if (reader != null && reader.CanGetColumnSchema())
+                    {
+                        do
+                        {
+                            columnSchemas.Add(reader.GetColumnSchema().ToArray());
+                        } while (await reader.NextResultAsync(cancellationToken));
+                    }
+                }
+
                 // Execute the command to get back a reader
                 using (DbDataReader reader = await dbCommand.ExecuteReaderAsync(cancellationToken))
                 {
@@ -375,6 +391,42 @@ namespace Microsoft.SqlTools.ServiceLayer.QueryExecution
                         await SendMessage(SR.QueryServiceCompletedSuccessfully, false);
                     }
                 }
+
+                if (columnSchemas != null)
+                {
+                    ExtendResultMetadata(columnSchemas, resultSets);
+                }
+            }
+        }
+
+        private void ExtendResultMetadata(List<DbColumn[]> columnSchemas, List<ResultSet> results)
+        {
+            if (columnSchemas.Count != results.Count) return;
+
+            for(int i = 0; i < results.Count; i++)
+            {
+                ResultSet result = results[i];
+                DbColumn[] columnSchema = columnSchemas[i];
+                if(result.Columns.Length > columnSchema.Length)
+                {
+                    throw new InvalidOperationException("Did not receive enough metadata columns.");
+                }
+
+                for(int j = 0; j < result.Columns.Length; j++)
+                {
+                    DbColumnWrapper resultCol = result.Columns[j];
+                    DbColumn schemaCol = columnSchema[j];
+
+                    if(!string.Equals(resultCol.DataTypeName, schemaCol.DataTypeName)
+                        || (!string.Equals(resultCol.ColumnName, schemaCol.ColumnName)
+                            && !string.IsNullOrEmpty(schemaCol.ColumnName)
+                            && !string.Equals(resultCol, SR.QueryServiceColumnNull)))
+                    {
+                        throw new InvalidOperationException("Inconsistent column metadata.");
+                    }
+
+                    result.Columns[j] = new DbColumnWrapper(schemaCol);
+                }
             }
         }
 

src/Microsoft.SqlTools.ServiceLayer/Utility/SqlScriptFormatters/FromSqlScript.cs

@@ -18,7 +18,8 @@ namespace Microsoft.SqlTools.ServiceLayer.Utility.SqlScriptFormatters
     {
         // Regex: optionally starts with N, captures string wrapped in single quotes
         private static readonly Regex StringRegex = new Regex("^N?'(.*)'$", RegexOptions.Compiled);
-        
+        private static readonly Regex BracketRegex = new Regex(@"^\[(.*)\]$", RegexOptions.Compiled);
+
         /// <summary>
         /// Decodes a multipart identifier as used in a SQL script into an array of the multiple
         /// parts of the identifier. Implemented as a state machine that iterates over the
@@ -128,6 +129,13 @@ namespace Microsoft.SqlTools.ServiceLayer.Utility.SqlScriptFormatters
             return literal;
         }
         
+        /// <summary>
+        /// Tests whether an identifier is escaped with brackets e.g. [Northwind].[dbo].[Orders]
+        /// </summary>
+        /// <param name="identifer">Identifier to check.</param>
+        /// <returns>Boolean indicating if identifier is escaped with brackets.</returns>
+        public static bool IsIdentifierBracketed(string identifer) => BracketRegex.IsMatch(identifer);
+
         #region Private Helpers
         
         private static bool HasNextCharacter(string haystack, char needle, int position)
@@ -143,7 +151,7 @@ namespace Microsoft.SqlTools.ServiceLayer.Utility.SqlScriptFormatters
             // Replace 2x of the escape character with 1x of the escape character
             return value.Replace(new string(escapeCharacter, 2), escapeCharacter.ToString());
         }
-        
+
         #endregion
     }
 }