Fix all referenced tables (#348)

* Refactor allReferencedTables so it no longer uses allReferencedAsPartitions. It wipes partition type from tables if there is no outerRef.

* Refactor allReferencedAsPartitions to also expand views as TableRefs

* Don't expand and exclude views in ZetaSql

Author: HenningKoller

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

@@ -7,8 +7,8 @@
 package no.nrk.bigquery
 
 import cats.syntax.all.*
-import no.nrk.bigquery.syntax.*
 import no.nrk.bigquery.BQSqlFrag.asSubQuery
+import no.nrk.bigquery.syntax.*
 
 import scala.annotation.tailrec
 
@@ -127,43 +127,49 @@ sealed trait BQSqlFrag {
   final def allReferencedAsPartitions: Seq[BQPartitionId[Any]] =
     allReferencedAsPartitions(expandAndExcludeViews = true)
   final def allReferencedAsPartitions(expandAndExcludeViews: Boolean): Seq[BQPartitionId[Any]] = {
-    def pf(outerRef: Option[BQPartitionId[Any]]): PartialFunction[BQSqlFrag, List[BQPartitionId[Any]]] = {
-      case BQSqlFrag.PartitionRef(partitionRef) =>
-        partitionRef.wholeTable match {
-          case tableDef: BQTableDef.View[?] if expandAndExcludeViews =>
-            tableDef.query.collect(pf(Some(partitionRef))).flatten
-          case tvf: BQAppliedTableValuedFunction[?] if expandAndExcludeViews =>
-            tvf.query.collect(pf(Some(partitionRef))).flatten
-          case _ => List(partitionRef)
-        }
-
-      case BQSqlFrag.TableRef(table) =>
-        (table.partitionType, outerRef) match {
-          case (partitionType: BQPartitionType.DatePartitioned, Some(partitionRef: BQPartitionId.DatePartitioned)) =>
-            List(table.withTableType(partitionType).assertPartition(partitionRef.partition))
-          case (partitionType: BQPartitionType.HourPartitioned, Some(partitionRef: BQPartitionId.HourPartitioned)) =>
-            List(table.withTableType(partitionType).assertPartition(partitionRef.partition))
-          case (_, _) => List(table.unpartitioned.assertPartition)
-        }
+    def expandAndApplyOuterRef(
+        table: BQTableLike[Any],
+        outerRef: Option[BQPartitionId[Any]]): List[BQPartitionId[Any]] =
+      table match {
+        case tableDef: BQTableDef.View[?] if expandAndExcludeViews => tableDef.query.collect(pf(outerRef)).flatten
+        case tvf: BQAppliedTableValuedFunction[?] if expandAndExcludeViews => tvf.query.collect(pf(outerRef)).flatten
+        case _ =>
+          (table.partitionType, outerRef) match {
+            case (partitionType: BQPartitionType.DatePartitioned, Some(partitionRef: BQPartitionId.DatePartitioned)) =>
+              List(table.withTableType(partitionType).assertPartition(partitionRef.partition))
+            case (partitionType: BQPartitionType.HourPartitioned, Some(partitionRef: BQPartitionId.HourPartitioned)) =>
+              List(table.withTableType(partitionType).assertPartition(partitionRef.partition))
+            case (_, _) => List(table.unpartitioned.assertPartition)
+          }
+      }
 
+    def pf(outerRef: Option[BQPartitionId[Any]]): PartialFunction[BQSqlFrag, List[BQPartitionId[Any]]] = {
+      case BQSqlFrag.PartitionRef(partitionRef) => expandAndApplyOuterRef(partitionRef.wholeTable, Some(partitionRef))
+      case BQSqlFrag.TableRef(table) => expandAndApplyOuterRef(table, outerRef)
       case BQSqlFrag.FillRef(fill) => List(fill.destination)
-      case BQSqlFrag.FilledTableRef(fill) =>
-        (fill.tableDef.partitionType, outerRef) match {
-          case (partitionType: BQPartitionType.DatePartitioned, Some(partitionRef: BQPartitionId.DatePartitioned)) =>
-            List(fill.tableDef.withTableType(partitionType).assertPartition(partitionRef.partition))
-          case (partitionType: BQPartitionType.HourPartitioned, Some(partitionRef: BQPartitionId.HourPartitioned)) =>
-            List(fill.tableDef.withTableType(partitionType).assertPartition(partitionRef.partition))
-          case (_, _) => List(fill.tableDef.unpartitioned.assertPartition)
-        }
+      case BQSqlFrag.FilledTableRef(fill) => expandAndApplyOuterRef(fill.tableDef, outerRef)
     }
 
     this.collect(pf(None)).flatten.distinct
   }
 
-  final def allReferencedTables: Seq[BQTableLike[Any]] =
-    allReferencedAsPartitions
-      .map(_.wholeTable)
-      .filterNot(tableLike => tableLike.isInstanceOf[BQTableDef.View[?]])
+  final def allReferencedTables(expandAndExcludeViews: Boolean): Seq[BQTableLike[Any]] = {
+    def expand(table: BQTableLike[Any]): List[BQTableLike[Any]] =
+      table match {
+        case tableDef: BQTableDef.View[?] if expandAndExcludeViews => tableDef.query.collect(pf).flatten
+        case tvf: BQAppliedTableValuedFunction[?] if expandAndExcludeViews => tvf.query.collect(pf).flatten
+        case _ => List(table)
+      }
+
+    def pf: PartialFunction[BQSqlFrag, List[BQTableLike[Any]]] = {
+      case BQSqlFrag.PartitionRef(partitionRef) => expand(partitionRef.wholeTable)
+      case BQSqlFrag.TableRef(table) => expand(table)
+      case BQSqlFrag.FillRef(fill) => List(fill.destination.wholeTable)
+      case BQSqlFrag.FilledTableRef(fill) => List(fill.tableDef)
+    }
+    this.collect(pf).flatten.distinct
+  }
+  final def allReferencedTables: Seq[BQTableLike[Any]] = allReferencedTables(expandAndExcludeViews = true)
 
   final def allReferencedTablesAsPartitions: Seq[BQPartitionId[Any]] =
     allReferencedAsPartitions(expandAndExcludeViews = true)

zetasql/src/main/scala/no/nrk/bigquery/ZetaSql.scala

@@ -170,7 +170,7 @@ class ZetaSql[F[_]](implicit F: Sync[F]) {
 
   def analyzeFirst(frag: BQSqlFrag): F[Either[AnalysisException, AnalyzedStatement]] =
     F.interruptible {
-      val tables = frag.allReferencedTables
+      val tables = frag.allReferencedTables(expandAndExcludeViews = false)
       val catalog = toCatalog(tables*)
       val rendered = frag.asString
 

zetasql/src/test/scala/no/nrk/bigquery/ZetaTest.scala

@@ -144,7 +144,10 @@ class ZetaTest extends munit.CatsEffectSuite {
 
     zetaSql
       .parseAndBuildAnalysableFragment(query, List(table), List(tvf))
-      .flatMap(fragment => zetaSql.queryFields(fragment).tupleRight(fragment.allReferencedTables.map(_.tableId)))
+      .flatMap(fragment =>
+        zetaSql
+          .queryFields(fragment)
+          .tupleRight(fragment.allReferencedTables(expandAndExcludeViews = false).map(_.tableId)))
       .assertEquals(expected -> List(BQTableId(tvf.name.dataset, tvf.name.name.value)))
   }