revert use of collect with referencing UDFS

Author: hamnis

core/src/main/scala/no/nrk/bigquery/BQSqlFrag.scala

@@ -106,8 +106,35 @@ sealed trait BQSqlFrag {
       case BQSqlFrag.FilledTableRef(fill) => fill.tableDef.unpartitioned.assertPartition
     }.distinct
 
-  final def allReferencedUDFs: Seq[UDF[UDF.UDFId]] =
-    this.collect { case BQSqlFrag.Call(udf, _) => udf }.distinct
+  final def allReferencedUDFs: Seq[UDF[UDF.UDFId]] = {
+    def fromBody(body: UDF.Body) =
+      body match {
+        case UDF.Body.Sql(body) => body.allReferencedUDFs
+        case _: UDF.Body.Js => Nil
+      }
+
+    this match {
+      case BQSqlFrag.Frag(_) => Nil
+      case BQSqlFrag.Call(udf @ UDF.Temporary(_, _, body, _), args) =>
+        (fromBody(body) ++ args.flatMap(
+          _.allReferencedUDFs
+        ) ++ List(udf)).distinct
+      case BQSqlFrag.Call(udf @ UDF.Persistent(_, _, body, _), args) =>
+        (fromBody(body) ++ args.flatMap(
+          _.allReferencedUDFs
+        ) ++ List(udf)).distinct
+      case BQSqlFrag.Call(udf: UDF.Reference, args) =>
+        (args.flatMap(
+          _.allReferencedUDFs
+        ) ++ List(udf)).distinct
+
+      case BQSqlFrag.Combined(values) =>
+        values.flatMap(_.allReferencedUDFs).distinct
+      case BQSqlFrag.PartitionRef(_) => Nil
+      case BQSqlFrag.FillRef(_) => Nil
+      case BQSqlFrag.FilledTableRef(_) => Nil
+    }
+  }
 
   override def toString: String = asString
 }

generated/bq-query-static/query_with_UDF.sql

@@ -1,4 +1,4 @@
-CREATE TEMP FUNCTION half_TMP(input INT64) RETURNS FLOAT64 AS ((double_TMP(input) / 2));
-CREATE TEMP FUNCTION double_TMP(input INT64) RETURNS INT64 AS ((input + input));
+CREATE TEMP FUNCTION xxdouble_TMP(input INT64) RETURNS INT64 AS ((input + input));
+CREATE TEMP FUNCTION half_TMP(input INT64) RETURNS FLOAT64 AS ((xxdouble_TMP(input) / 2));
 
 select half_TMP(1)

generated/bq-query/query_with_UDF.sql

@@ -1,5 +1,5 @@
-CREATE TEMP FUNCTION half_TMP(input INT64) RETURNS FLOAT64 AS ((double_TMP(input) / 2));
+CREATE TEMP FUNCTION xxdouble_TMP(input INT64) RETURNS INT64 AS ((input + input));
 
-CREATE TEMP FUNCTION double_TMP(input INT64) RETURNS INT64 AS ((input + input));
+CREATE TEMP FUNCTION half_TMP(input INT64) RETURNS FLOAT64 AS ((xxdouble_TMP(input) / 2));
 
 select half_TMP(1)

testing/src/test/scala/no/nrk/bigquery/LiveTempUdfTest.scala

@@ -5,7 +5,7 @@ import no.nrk.bigquery.testing.{BQSmokeTest, BigQueryTestClient}
 
 class LiveTempUdfTest extends BQSmokeTest(BigQueryTestClient.testClient) {
   val udf1 = UDF.temporary(
-    Ident("double_TMP"),
+    Ident("xxdouble_TMP"),
     List(UDF.Param("input", BQType.INT64)),
     UDF.Body.Sql(bqsql"(input + input)"),
     Some(BQType.INT64))