Merge pull request #114 from nrkno/more-helpers-bqsqlfrag

More helpers bqsqlfrag

Author: HenningKoller

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

@@ -113,11 +113,29 @@ sealed trait BQSqlFrag {
   }
 
   final def allReferencedAsPartitions: Seq[BQPartitionId[Any]] =
-    this.collect {
-      case BQSqlFrag.PartitionRef(ref) => ref
-      case BQSqlFrag.FillRef(fill) => fill.destination
-      case BQSqlFrag.FilledTableRef(fill) => fill.tableDef.unpartitioned.assertPartition
-    }.distinct
+    allReferencedAsPartitions(expandAndExcludeViews = true)
+  final def allReferencedAsPartitions(expandAndExcludeViews: Boolean): Seq[BQPartitionId[Any]] = {
+    def pf: PartialFunction[BQSqlFrag, List[BQPartitionId[Any]]] = {
+      case BQSqlFrag.PartitionRef(ref) =>
+        ref.wholeTable match {
+          case tableDef: BQTableDef.View[_] if expandAndExcludeViews => tableDef.query.collect(pf).flatten
+          case _ => List(ref)
+        }
+      case BQSqlFrag.FillRef(fill) => List(fill.destination)
+      case BQSqlFrag.FilledTableRef(fill) => List(fill.tableDef.unpartitioned.assertPartition)
+    }
+
+    this.collect(pf).flatten.distinct
+  }
+
+  final def allReferencedTables: Seq[BQTableLike[Any]] =
+    allReferencedAsPartitions(expandAndExcludeViews = true)
+      .map(_.wholeTable)
+      .filterNot(tableLike => tableLike.isInstanceOf[BQTableDef.View[_]])
+
+  final def allReferencedTablesAsPartitions: Seq[BQPartitionId[Any]] =
+    allReferencedAsPartitions(expandAndExcludeViews = true)
+      .filterNot(pid => pid.wholeTable.isInstanceOf[BQTableDef.View[_]])
 
   final def allReferencedUDFs: Seq[UDF[UDF.UDFId]] =
     this.collect { case BQSqlFrag.Call(udf, _) => udf }.distinct

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

@@ -73,7 +73,11 @@ object BQTableLike {
   * a non-partitioned and a date-partitioned table into a list (or otherwise lose precise types), you also lose the
   * ability to construct legal [[BQPartitionId]] s with `assertPartition`
   */
-case class BQTableRef[+P](tableId: BQTableId, partitionType: BQPartitionType[P]) extends BQTableLike[P] {
+case class BQTableRef[+P](
+    tableId: BQTableId,
+    partitionType: BQPartitionType[P],
+    labels: TableLabels = TableLabels.Empty)
+    extends BQTableLike[P] {
   override def unpartitioned: BQTableRef[Unit] =
     withTableType(BQPartitionType.ignoredPartitioning(partitionType))
 

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

@@ -48,6 +48,13 @@ case class TableLabels(values: SortedMap[String, String]) {
   def ++(other: TableLabels): TableLabels =
     withAll(other.values)
 
+  def contains(tableLabels: TableLabels): Boolean =
+    if (tableLabels.values.nonEmpty) {
+      tableLabels.values.forall { case (key, value) =>
+        values.get(key).contains(value)
+      }
+    } else false
+
   /** This method is needed for the case where we delete a label. It is deleted by setting it to `null`.
     *
     * As such, we need to know the labels of a table in production before we compute the new set of labels to use when

testing/src/main/scala/no/nrk/bigquery/testing/BQSmokeTest.scala

@@ -424,7 +424,7 @@ object BQSmokeTest {
       case p @ BQSqlFrag.PartitionRef(pid) =>
         val schemaOpt: Option[BQSchema] =
           pid.wholeTable match {
-            case BQTableRef(_, _) => None
+            case BQTableRef(_, _, _) => None
             case x: BQTableDef[Any] => Some(x.schema)
           }