Merge pull request #74 from nrkno/udf-live-test

Fix allReferencedUDFs

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

@@ -0,0 +1,4 @@
+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

@@ -0,0 +1,5 @@
+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)

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

@@ -0,0 +1,21 @@
+package no.nrk.bigquery
+
+import no.nrk.bigquery.syntax._
+import no.nrk.bigquery.testing.{BQSmokeTest, BigQueryTestClient}
+
+class LiveTempUdfTest extends BQSmokeTest(BigQueryTestClient.testClient) {
+  val udf1 = UDF.temporary(
+    Ident("xxdouble_TMP"),
+    List(UDF.Param("input", BQType.INT64)),
+    UDF.Body.Sql(bqsql"(input + input)"),
+    Some(BQType.INT64))
+  val udf2 = UDF.temporary(
+    Ident("half_TMP"),
+    List(UDF.Param("input", BQType.INT64)),
+    UDF.Body.Sql(bqsql"(${udf1(Ident("input"))} / 2)"),
+    Some(BQType.FLOAT64))
+
+  bqTypeCheckTest("query with UDF") {
+    BQQuery[Double](bqsql"select ${udf2(1)}")
+  }
+}