From a82afed6d9e154bb09f27b50f7547be6dbcece9b Mon Sep 17 00:00:00 2001
From: rick-nguyen <97256426+rick-nguyen@users.noreply.github.com>
Date: Thu, 29 Aug 2024 11:10:49 -0700
Subject: [PATCH] Fix UDFs not properly handled in TryGetDataSource recursion
 (#2615)

`FilterUDF():DataSource = Filter(DataSource, Value > 10);
FilterUDF2():DataSource = Filter(FilterUDF(), Value > 11);`

`TryGetDsInfo(Callnode)` would not get a datasource for `FilterUDF()`,
leading to Binder not marking this node async. This led to an issue of
delegation not being preserved on `FilterUDF2()`.
---
 .../Functions/UserDefinedFunction.cs          | 13 ++--
 .../RecalcEngineTests.cs                      | 60 +++++++++++++------
 2 files changed, 50 insertions(+), 23 deletions(-)

diff --git a/src/libraries/Microsoft.PowerFx.Core/Functions/UserDefinedFunction.cs b/src/libraries/Microsoft.PowerFx.Core/Functions/UserDefinedFunction.cs
index 49b6a13f9e..812122724f 100644
--- a/src/libraries/Microsoft.PowerFx.Core/Functions/UserDefinedFunction.cs
+++ b/src/libraries/Microsoft.PowerFx.Core/Functions/UserDefinedFunction.cs
@@ -47,10 +47,10 @@ internal class UserDefinedFunction : TexlFunction, IExternalPageableSymbol, IExt
         public bool IsDelegatable => _binding.IsDelegatable(_binding.Top);
 
         public override bool IsServerDelegatable(CallNode callNode, TexlBinding binding)
-        {
-            Contracts.AssertValue(callNode);
-            Contracts.AssertValue(binding);
-            Contracts.Assert(binding.GetInfo(callNode).Function is UserDefinedFunction udf && udf.Binding != null);
+        {
+            Contracts.AssertValue(callNode);
+            Contracts.AssertValue(binding);
+            Contracts.Assert(binding.GetInfo(callNode).Function is UserDefinedFunction udf && udf.Binding != null);
 
             return base.IsServerDelegatable(callNode, binding) || IsDelegatable;
         }
@@ -70,6 +70,11 @@ public bool TryGetExternalDataSource(out IExternalDataSource dataSource)
             return ArgValidators.DelegatableDataSourceInfoValidator.TryGetValidValue(_binding.Top, _binding, out dataSource);
         }
 
+        public override bool TryGetDataSource(CallNode callNode, TexlBinding binding, out IExternalDataSource dsInfo)
+        {
+            return TryGetExternalDataSource(out dsInfo);
+        }
+
         /// <summary>
         /// Initializes a new instance of the <see cref="UserDefinedFunction"/> class.
         /// </summary>
diff --git a/src/tests/Microsoft.PowerFx.Interpreter.Tests.Shared/RecalcEngineTests.cs b/src/tests/Microsoft.PowerFx.Interpreter.Tests.Shared/RecalcEngineTests.cs
index 188c722042..4240122887 100644
--- a/src/tests/Microsoft.PowerFx.Interpreter.Tests.Shared/RecalcEngineTests.cs
+++ b/src/tests/Microsoft.PowerFx.Interpreter.Tests.Shared/RecalcEngineTests.cs
@@ -10,6 +10,8 @@
 using System.Threading.Tasks;
 using Microsoft.PowerFx.Core;
 using Microsoft.PowerFx.Core.Functions;
+using Microsoft.PowerFx.Core.Functions.Delegation;
+using Microsoft.PowerFx.Core.Functions.Delegation.DelegationMetadata;
 using Microsoft.PowerFx.Core.IR;
 using Microsoft.PowerFx.Core.Localization;
 using Microsoft.PowerFx.Core.Tests;
@@ -679,34 +681,54 @@ public void ImperativeUserDefinedFunctionTest(string udfExpression, string expre
         [Fact]
 
         public void DelegatableUDFTest()
-        {
-            var config = new PowerFxConfig();
-            config.EnableSetFunction();
-
+        {
+            var symbolTable = new DelegatableSymbolTable();
             var schema = DType.CreateTable(
                 new TypedName(DType.Guid, new DName("ID")),
                 new TypedName(DType.Number, new DName("Value")));
-            config.SymbolTable.AddEntity(new TestDataSource("MyDataSource", schema));
-            config.SymbolTable.AddType(new DName("MyDataSourceTableType"), FormulaType.Build(schema));
+            symbolTable.AddEntity(new TestDelegableDataSource(
+                "MyDataSource",
+                schema,
+                new TestDelegationMetadata(
+                        new DelegationCapability(DelegationCapability.Filter),
+                        schema,
+                        new FilterOpMetadata(
+                            schema,
+                            new Dictionary<DPath, DelegationCapability>(),
+                            new Dictionary<DPath, DelegationCapability>(),
+                            new DelegationCapability(DelegationCapability.GreaterThan),
+                            null)),
+                true));
+            symbolTable.AddType(new DName("MyDataSourceTableType"), FormulaType.Build(schema));
+            var config = new PowerFxConfig()
+            {
+                SymbolTable = symbolTable
+            };
+            config.EnableSetFunction();
 
             var recalcEngine = new RecalcEngine(config);
 
-            recalcEngine.AddUserDefinedFunction("A():MyDataSourceTableType = Filter(Sort(MyDataSource,Value), Value > 10);C():MyDataSourceTableType = A();", CultureInfo.InvariantCulture, symbolTable: recalcEngine.EngineSymbols, allowSideEffects: true);
-            var func = recalcEngine.Functions.WithName("A");
+            recalcEngine.AddUserDefinedFunction("A():MyDataSourceTableType = Filter(MyDataSource, Value > 10);C():MyDataSourceTableType = A(); B():MyDataSourceTableType = Filter(C(), Value > 11); D():MyDataSourceTableType = { Filter(B(), Value > 12); };", CultureInfo.InvariantCulture, symbolTable: recalcEngine.EngineSymbols, allowSideEffects: true);
+            var func = recalcEngine.Functions.WithName("A").First() as UserDefinedFunction;
 
-            if (func is UserDefinedFunction udf)
-            {
-                Assert.True(udf.IsAsync);
-                Assert.True(udf.IsDelegatable);
-            }
+            Assert.True(func.IsAsync);
+            Assert.True(func.IsDelegatable);
 
-            func = recalcEngine.Functions.WithName("C");
+            func = recalcEngine.Functions.WithName("C").First() as UserDefinedFunction;
 
-            if (func is UserDefinedFunction udf2)
-            {
-                Assert.True(udf2.IsAsync);
-                Assert.True(udf2.IsDelegatable);
-            }
+            Assert.True(func.IsAsync);
+            Assert.True(func.IsDelegatable);
+
+            func = recalcEngine.Functions.WithName("B").First() as UserDefinedFunction;
+
+            Assert.True(func.IsAsync);
+            Assert.True(func.IsDelegatable);
+
+            func = recalcEngine.Functions.WithName("D").First() as UserDefinedFunction;
+
+            // Imperative function is not delegable
+            Assert.True(func.IsAsync);
+            Assert.True(!func.IsDelegatable);
         }
 
         // Binding to inner functions does not impact outer functions.