From 25f1cd87c5dbab1e82bf19df27daa4420064261c Mon Sep 17 00:00:00 2001
From: Fernando Velasquez <fevelasquez@google.com>
Date: Tue, 12 Apr 2022 13:55:59 -0400
Subject: [PATCH] Added logic to qualify stage names for Dedup and Group By
 operations

---
 .../bigquery/relational/BigQueryRelation.java    |  4 ++--
 .../relational/BigQueryRelationTest.java         | 16 ++++++++--------
 2 files changed, 10 insertions(+), 10 deletions(-)

diff --git a/src/main/java/io/cdap/plugin/gcp/bigquery/relational/BigQueryRelation.java b/src/main/java/io/cdap/plugin/gcp/bigquery/relational/BigQueryRelation.java
index 3e08723a6b..41997f74c8 100644
--- a/src/main/java/io/cdap/plugin/gcp/bigquery/relational/BigQueryRelation.java
+++ b/src/main/java/io/cdap/plugin/gcp/bigquery/relational/BigQueryRelation.java
@@ -290,7 +290,7 @@ private static String buildGroupBy(GroupByAggregationDefinition definition,
     // Instantiate query builder and generate select expression
     BigQueryGroupBySQLBuilder builder = new BigQueryGroupBySQLBuilder(qualify(definition),
                                                                       sourceExpression,
-                                                                      datasetName);
+                                                                      qualify(datasetName));
     return builder.getQuery();
   }
 
@@ -300,7 +300,7 @@ private static String buildDeduplicate(DeduplicateAggregationDefinition definiti
     // Instantiate query builder and generate select expression
     BigQueryDeduplicateSQLBuilder builder = new BigQueryDeduplicateSQLBuilder(qualify(definition),
                                                                               sourceExpression,
-                                                                              datasetName);
+                                                                              qualify(datasetName));
     return builder.getQuery();
   }
 
diff --git a/src/test/java/io/cdap/plugin/gcp/bigquery/relational/BigQueryRelationTest.java b/src/test/java/io/cdap/plugin/gcp/bigquery/relational/BigQueryRelationTest.java
index 8a0240a100..7834a3d546 100644
--- a/src/test/java/io/cdap/plugin/gcp/bigquery/relational/BigQueryRelationTest.java
+++ b/src/test/java/io/cdap/plugin/gcp/bigquery/relational/BigQueryRelationTest.java
@@ -57,11 +57,11 @@ public void setUp() {
     columns.add("a");
     columns.add("b");
 
-    baseRelation = new BigQueryRelation("ds",
+    baseRelation = new BigQueryRelation("d s",
                                         columns,
                                         null,
                                         () -> "select * from tbl");
-    baseRelation.setInputDatasets(Collections.singletonMap("ds", ds));
+    baseRelation.setInputDatasets(Collections.singletonMap("d s", ds));
   }
 
   @Test
@@ -78,7 +78,7 @@ public void testSetColumn() {
     Assert.assertTrue(columns.contains("a"));
     Assert.assertTrue(columns.contains("b"));
     Assert.assertTrue(columns.contains("c"));
-    Assert.assertEquals("SELECT `a` AS `a` , `b` AS `b` , a+b AS `c` FROM (select * from tbl) AS `ds`",
+    Assert.assertEquals("SELECT `a` AS `a` , `b` AS `b` , a+b AS `c` FROM (select * from tbl) AS `d s`",
                         bqRelation.getSQLStatement());
   }
 
@@ -105,7 +105,7 @@ public void testDropColumn() {
     Set<String> columns = bqRelation.getColumns();
     Assert.assertEquals(1, columns.size());
     Assert.assertTrue(columns.contains("a"));
-    Assert.assertEquals("SELECT `a` AS `a` FROM (select * from tbl) AS `ds`",
+    Assert.assertEquals("SELECT `a` AS `a` FROM (select * from tbl) AS `d s`",
                         bqRelation.getSQLStatement());
   }
 
@@ -132,7 +132,7 @@ public void testSelect() {
     Assert.assertEquals(2, columns.size());
     Assert.assertTrue(columns.contains("new_a"));
     Assert.assertTrue(columns.contains("new_b"));
-    Assert.assertEquals("SELECT a AS `new_a` , b AS `new_b` FROM (select * from tbl) AS `ds`",
+    Assert.assertEquals("SELECT a AS `new_a` , b AS `new_b` FROM (select * from tbl) AS `d s`",
                         bqRelation.getSQLStatement());
   }
 
@@ -162,7 +162,7 @@ public void testFilter() {
     Assert.assertEquals(2, columns.size());
     Assert.assertTrue(columns.contains("a"));
     Assert.assertTrue(columns.contains("b"));
-    Assert.assertEquals("SELECT `a` AS `a` , `b` AS `b` FROM (select * from tbl) AS `ds` WHERE a > 2",
+    Assert.assertEquals("SELECT `a` AS `a` , `b` AS `b` FROM (select * from tbl) AS `d s` WHERE a > 2",
                         bqRelation.getSQLStatement());
   }
 
@@ -214,7 +214,7 @@ public void testGroupBy() {
     Assert.assertTrue(columns.contains("c"));
     Assert.assertTrue(columns.contains("d"));
     Assert.assertEquals("SELECT a AS `a` , MAX(a) AS `b` , MIN(b) AS `c` , d AS `d` "
-                          + "FROM ( select * from tbl ) AS ds "
+                          + "FROM ( select * from tbl ) AS `d s` "
                           + "GROUP BY a , d",
                         bqRelation.getSQLStatement());
   }
@@ -288,7 +288,7 @@ public void testDeduplicate() {
     Assert.assertTrue(transformExpression.contains("`) FROM (SELECT a AS `a` , b AS `b` , c AS `c` , d AS `d` , " +
                                                      "ROW_NUMBER() OVER ( PARTITION BY a , d ORDER BY a DESC " +
                                                       "NULLS LAST ) AS `"));
-    Assert.assertTrue(transformExpression.contains("` FROM ( select * from tbl ) AS ds) WHERE `rn_"));
+    Assert.assertTrue(transformExpression.contains("` FROM ( select * from tbl ) AS `d s`) WHERE `rn_"));
     Assert.assertTrue(transformExpression.endsWith("` = 1"));
   }