From 2d248ecb846b9e94a6382db7c6f27712e1a5728d Mon Sep 17 00:00:00 2001
From: Mark Carrington <MarkMpn@users.noreply.github.com>
Date: Wed, 30 Mar 2022 21:19:56 +0100
Subject: [PATCH] Use actual list of tables from TDS Endpoint to determine
 compatibility Always return a SqlNode for the execution plan even if the
 entire batch can be executed by TDS Endpoint

---
 .../ExecutionPlanBuilder.cs                   | 60 ++++++++++++-------
 .../TDSEndpointCompatibilityVisitor.cs        | 41 ++++++++++---
 2 files changed, 74 insertions(+), 27 deletions(-)

diff --git a/MarkMpn.Sql4Cds.Engine/ExecutionPlanBuilder.cs b/MarkMpn.Sql4Cds.Engine/ExecutionPlanBuilder.cs
index 19f1f600..54b898ee 100644
--- a/MarkMpn.Sql4Cds.Engine/ExecutionPlanBuilder.cs
+++ b/MarkMpn.Sql4Cds.Engine/ExecutionPlanBuilder.cs
@@ -87,13 +87,30 @@ public IRootExecutionPlanNode[] Build(string sql, IDictionary<string, DataTypeRe
 
             if (TDSEndpoint.CanUseTDSEndpoint(Options, DataSources[Options.PrimaryDataSource].Connection))
             {
-                var tdsEndpointCompatibilityVisitor = new TDSEndpointCompatibilityVisitor(DataSources[Options.PrimaryDataSource].Metadata);
-                fragment.Accept(tdsEndpointCompatibilityVisitor);
-
-                if (tdsEndpointCompatibilityVisitor.IsCompatible)
+                using (var con = TDSEndpoint.Connect(DataSources[Options.PrimaryDataSource].Connection))
                 {
-                    useTDSEndpointDirectly = true;
-                    return null;
+                    var tdsEndpointCompatibilityVisitor = new TDSEndpointCompatibilityVisitor(con, DataSources[Options.PrimaryDataSource].Metadata);
+                    fragment.Accept(tdsEndpointCompatibilityVisitor);
+
+                    if (tdsEndpointCompatibilityVisitor.IsCompatible)
+                    {
+                        useTDSEndpointDirectly = true;
+                        var sqlNode = new SqlNode
+                        {
+                            DataSource = Options.PrimaryDataSource,
+                            Sql = sql,
+                            Index = 0,
+                            Length = sql.Length
+                        };
+
+                        if (parameters != null)
+                        {
+                            foreach (var param in parameters.Keys)
+                                sqlNode.Parameters.Add(param);
+                        }
+
+                        return new IRootExecutionPlanNode[] { sqlNode };
+                    }
                 }
             }
 
@@ -1296,25 +1313,28 @@ private IRootExecutionPlanNodeInternal ConvertSelectStatement(SelectStatement se
         {
             if (TDSEndpoint.CanUseTDSEndpoint(Options, DataSources[Options.PrimaryDataSource].Connection))
             {
-                var tdsEndpointCompatibilityVisitor = new TDSEndpointCompatibilityVisitor(DataSources[Options.PrimaryDataSource].Metadata, false);
-                select.Accept(tdsEndpointCompatibilityVisitor);
-
-                if (tdsEndpointCompatibilityVisitor.IsCompatible)
+                using (var con = TDSEndpoint.Connect(DataSources[Options.PrimaryDataSource].Connection))
                 {
-                    select.ScriptTokenStream = null;
-                    var sql = new SqlNode
+                    var tdsEndpointCompatibilityVisitor = new TDSEndpointCompatibilityVisitor(con, DataSources[Options.PrimaryDataSource].Metadata, false);
+                    select.Accept(tdsEndpointCompatibilityVisitor);
+
+                    if (tdsEndpointCompatibilityVisitor.IsCompatible)
                     {
-                        DataSource = Options.PrimaryDataSource,
-                        Sql = select.ToSql()
-                    };
+                        select.ScriptTokenStream = null;
+                        var sql = new SqlNode
+                        {
+                            DataSource = Options.PrimaryDataSource,
+                            Sql = select.ToSql()
+                        };
 
-                    var variables = new VariableCollectingVisitor();
-                    select.Accept(variables);
+                        var variables = new VariableCollectingVisitor();
+                        select.Accept(variables);
 
-                    foreach (var variable in variables.Variables)
-                        sql.Parameters.Add(variable.Name);
+                        foreach (var variable in variables.Variables)
+                            sql.Parameters.Add(variable.Name);
 
-                    return sql;
+                        return sql;
+                    }
                 }
             }
 
diff --git a/MarkMpn.Sql4Cds.Engine/Visitors/TDSEndpointCompatibilityVisitor.cs b/MarkMpn.Sql4Cds.Engine/Visitors/TDSEndpointCompatibilityVisitor.cs
index c7e4dadb..05ddb2f2 100644
--- a/MarkMpn.Sql4Cds.Engine/Visitors/TDSEndpointCompatibilityVisitor.cs
+++ b/MarkMpn.Sql4Cds.Engine/Visitors/TDSEndpointCompatibilityVisitor.cs
@@ -1,5 +1,6 @@
 using System;
 using System.Collections.Generic;
+using System.Data;
 using System.Linq;
 using System.Text;
 using Microsoft.SqlServer.TransactSql.ScriptDom;
@@ -14,17 +15,24 @@ namespace MarkMpn.Sql4Cds.Engine.Visitors
     /// <seealso cref="https://docs.microsoft.com/en-us/powerapps/developer/data-platform/how-dataverse-sql-differs-from-transact-sql?tabs=not-supported"/>
     class TDSEndpointCompatibilityVisitor : TSqlFragmentVisitor
     {
+        private readonly IDbConnection _con;
         private readonly IAttributeMetadataCache _metadata;
         private readonly Dictionary<string, string> _tableNames;
+        private readonly HashSet<string> _supportedTables;
         private bool? _isEntireBatch;
         private TSqlFragment _root;
 
         /// <summary>
         /// Creates a new <see cref="TDSEndpointCompatibilityVisitor"/>
         /// </summary>
+        /// <param name="con">A connection to the TDS Endpoint</param>
         /// <param name="metadata">The metadata cache for the primary data source</param>
-        public TDSEndpointCompatibilityVisitor(IAttributeMetadataCache metadata, bool? isEntireBatch = null, Dictionary<string, string> outerTableNames = null)
+        /// <param name="isEntireBatch">Indicates if this query is the entire SQL batch, or a single statement within it</param>
+        /// <param name="outerTableNames">A mapping of table aliases to table names available from the outer query</param>
+        /// <param name="supportedTables">A pre-calculated list of supported tables</param>
+        public TDSEndpointCompatibilityVisitor(IDbConnection con, IAttributeMetadataCache metadata, bool? isEntireBatch = null, Dictionary<string, string> outerTableNames = null, HashSet<string> supportedTables = null)
         {
+            _con = con;
             _metadata = metadata;
             _tableNames = new Dictionary<string, string>(StringComparer.OrdinalIgnoreCase);
             _isEntireBatch = isEntireBatch;
@@ -35,6 +43,26 @@ public TDSEndpointCompatibilityVisitor(IAttributeMetadataCache metadata, bool? i
                     _tableNames[kvp.Key] = kvp.Value;
             }
 
+            if (supportedTables != null)
+            {
+                _supportedTables = supportedTables;
+            }
+            else
+            {
+                _supportedTables = new HashSet<string>(StringComparer.OrdinalIgnoreCase);
+
+                using (var cmd = con.CreateCommand())
+                {
+                    cmd.CommandText = "SELECT name FROM sys.tables";
+
+                    using (var reader = cmd.ExecuteReader())
+                    {
+                        while (reader.Read())
+                            _supportedTables.Add(reader.GetString(0));
+                    }
+                }
+            }
+
             IsCompatible = true;
         }
 
@@ -65,10 +93,9 @@ public override void Visit(NamedTableReference node)
                 return;
             }
 
-            var entity = TryGetEntity(node.SchemaObject.BaseIdentifier.Value);
-            if (entity?.DataProviderId != null)
+            if (!_supportedTables.Contains(node.SchemaObject.BaseIdentifier.Value))
             {
-                // No access to entities with custom data providers
+                // Table does not exist in TDS endpoint
                 IsCompatible = false;
                 return;
             }
@@ -184,7 +211,7 @@ public override void Visit(ScalarSubquery node)
             // Name resolution needs to be scoped to the query, so create a new sub-visitor
             if (IsCompatible && _root != node)
             {
-                var subVisitor = new TDSEndpointCompatibilityVisitor(_metadata, _isEntireBatch, _tableNames);
+                var subVisitor = new TDSEndpointCompatibilityVisitor(_con, _metadata, _isEntireBatch, _tableNames, _supportedTables);
                 node.Accept(subVisitor);
 
                 if (!subVisitor.IsCompatible)
@@ -201,7 +228,7 @@ public override void Visit(QueryDerivedTable node)
             // Name resolution needs to be scoped to the query, so create a new sub-visitor
             if (IsCompatible && _root != node)
             {
-                var subVisitor = new TDSEndpointCompatibilityVisitor(_metadata, _isEntireBatch);
+                var subVisitor = new TDSEndpointCompatibilityVisitor(_con, _metadata, _isEntireBatch, supportedTables: _supportedTables);
                 node.Accept(subVisitor);
 
                 if (!subVisitor.IsCompatible)
@@ -218,7 +245,7 @@ public override void Visit(SelectStatement node)
             // Name resolution needs to be scoped to the query, so create a new sub-visitor
             if (IsCompatible && _root != node)
             {
-                var subVisitor = new TDSEndpointCompatibilityVisitor(_metadata, _isEntireBatch);
+                var subVisitor = new TDSEndpointCompatibilityVisitor(_con, _metadata, _isEntireBatch, supportedTables: _supportedTables);
                 node.Accept(subVisitor);
 
                 if (!subVisitor.IsCompatible)