add test for source analyze

Canner · May 23, 2024 · b75e0a9 · b75e0a9
1 parent a1777e9
commit b75e0a9
Show file tree

Hide file tree

Showing 2 changed files with 51 additions and 0 deletions.
diff --git a/wren-base/src/test/java/io/wren/base/sqlrewrite/analyzer/TestDecisionPointAnalyzer.java b/wren-base/src/test/java/io/wren/base/sqlrewrite/analyzer/TestDecisionPointAnalyzer.java
@@ -33,6 +33,7 @@
 import org.testng.annotations.Test;
 
 import java.util.List;
+import java.util.Set;
 
 import static io.wren.base.dto.Model.onTableReference;
 import static io.wren.base.dto.TableReference.tableReference;
@@ -180,6 +181,7 @@ public void testRelation()
             assertThat(joinRelation.getRight().getAlias()).isNull();
             assertThat(joinRelation.getRight().getType()).isEqualTo(RelationAnalysis.Type.TABLE);
             assertThat(((TableRelation) joinRelation.getRight()).getTableName()).isEqualTo("orders");
+            assertThat(joinRelation.getExprSources()).isEmpty();
         }
         else {
             throw new AssertionError("wrong type");
@@ -218,6 +220,23 @@ public void testRelation()
             assertThat(joinRelation.getRight().getType()).isEqualTo(RelationAnalysis.Type.TABLE);
             assertThat(((TableRelation) joinRelation.getRight()).getTableName()).isEqualTo("orders");
             assertThat(joinRelation.getCriteria()).isEqualTo("ON (customer.custkey = orders.custkey)");
+            assertThat(joinRelation.getExprSources().size()).isEqualTo(2);
+            assertThat(Set.copyOf(joinRelation.getExprSources())).isEqualTo(Set.of(
+                    new RelationAnalysis.ExprSource("customer.custkey", "customer"),
+                    new RelationAnalysis.ExprSource("orders.custkey", "orders")));
+        }
+        else {
+            throw new AssertionError("wrong type");
+        }
+
+        statement = parseSql("SELECT * FROM (customer c JOIN orders o ON c.custkey = o.custkey) join_relation");
+        result = DecisionPointAnalyzer.analyze(statement, DEFAULT_SESSION_CONTEXT, mdl);
+        if (result.get(0).getRelation() instanceof JoinRelation) {
+            JoinRelation joinRelation = (JoinRelation) result.get(0).getRelation();
+            assertThat(joinRelation.getExprSources().size()).isEqualTo(2);
+            assertThat(Set.copyOf(joinRelation.getExprSources())).isEqualTo(Set.of(
+                    new RelationAnalysis.ExprSource("c.custkey", "customer"),
+                    new RelationAnalysis.ExprSource("o.custkey", "orders")));
         }
         else {
             throw new AssertionError("wrong type");
@@ -238,6 +257,9 @@ public void testRelation()
             assertThat(joinRelation.getRight().getType()).isEqualTo(RelationAnalysis.Type.TABLE);
             assertThat(((TableRelation) joinRelation.getRight()).getTableName()).isEqualTo("lineitem");
             assertThat(joinRelation.getCriteria()).isEqualTo("ON (orders.orderkey = lineitem.orderkey)");
+            assertThat(Set.copyOf(joinRelation.getExprSources())).isEqualTo(Set.of(
+                    new RelationAnalysis.ExprSource("lineitem.orderkey", "lineitem"),
+                    new RelationAnalysis.ExprSource("orders.orderkey", "orders")));
         }
         else {
             throw new AssertionError("wrong type");
@@ -257,6 +279,9 @@ public void testRelation()
             assertThat(joinRelation.getRight().getType()).isEqualTo(RelationAnalysis.Type.TABLE);
             assertThat(((TableRelation) joinRelation.getRight()).getTableName()).isEqualTo("orders");
             assertThat(joinRelation.getCriteria()).isEqualTo("USING (custkey)");
+            assertThat(Set.copyOf(joinRelation.getExprSources())).isEqualTo(Set.of(
+                    new RelationAnalysis.ExprSource("custkey", "customer"),
+                    new RelationAnalysis.ExprSource("custkey", "orders")));
         }
         else {
             throw new AssertionError("wrong type");
@@ -280,6 +305,28 @@ public void testRelation()
         else {
             throw new AssertionError("wrong type");
         }
+
+        statement = parseSql("SELECT * FROM (customer JOIN (SELECT 1 as custkey, 'xxx' as name) orders(custkey, name) ON customer.custkey = orders.custkey) join_relation");
+        result = DecisionPointAnalyzer.analyze(statement, DEFAULT_SESSION_CONTEXT, mdl);
+        assertThat(result.size()).isEqualTo(1);
+        assertThat(result.get(0).getRelation().getAlias()).isEqualTo("join_relation");
+        assertThat(result.get(0).getRelation().getType()).isEqualTo(RelationAnalysis.Type.INNER_JOIN);
+        if (result.get(0).getRelation() instanceof JoinRelation) {
+            JoinRelation joinRelation = (JoinRelation) result.get(0).getRelation();
+            assertThat(joinRelation.getLeft().getAlias()).isNull();
+            assertThat(joinRelation.getLeft().getType()).isEqualTo(RelationAnalysis.Type.TABLE);
+            assertThat(((TableRelation) joinRelation.getLeft()).getTableName()).isEqualTo("customer");
+            assertThat(joinRelation.getRight().getAlias()).isEqualTo("orders");
+            assertThat(joinRelation.getRight().getType()).isEqualTo(RelationAnalysis.Type.SUBQUERY);
+            assertThat(((SubqueryRelation) joinRelation.getRight()).getBody().size()).isEqualTo(1);
+            assertThat(joinRelation.getCriteria()).isEqualTo("ON (customer.custkey = orders.custkey)");
+            assertThat(joinRelation.getExprSources().size()).isEqualTo(1);
+            assertThat(Set.copyOf(joinRelation.getExprSources())).isEqualTo(Set.of(
+                    new RelationAnalysis.ExprSource("customer.custkey", "customer")));
+        }
+        else {
+            throw new AssertionError("wrong type");
+        }
     }
 
     @Test

diff --git a/wren-tests/src/test/java/io/wren/testing/TestAnalysisResource.java b/wren-tests/src/test/java/io/wren/testing/TestAnalysisResource.java
@@ -29,6 +29,7 @@
 import java.nio.file.Files;
 import java.nio.file.Path;
 import java.util.List;
+import java.util.Set;
 
 import static io.wren.base.dto.Manifest.MANIFEST_JSON_CODEC;
 import static io.wren.base.dto.Model.onTableReference;
@@ -140,6 +141,9 @@ public void testBasic()
         assertThat(result.get(0).getRelation().getLeft().getType()).isEqualTo(RelationAnalysis.Type.TABLE.name());
         assertThat(result.get(0).getRelation().getRight().getType()).isEqualTo(RelationAnalysis.Type.TABLE.name());
         assertThat(result.get(0).getRelation().getCriteria()).isEqualTo("ON (c.custkey = o.custkey)");
+        assertThat(Set.copyOf(result.get(0).getRelation().getExprSources()))
+                .isEqualTo(Set.of(new QueryAnalysisDto.ExprSourceDto("c.custkey", "customer"),
+                        new QueryAnalysisDto.ExprSourceDto("o.custkey", "orders")));
 
         result = getSqlAnalysis(new SqlAnalysisInputDto(null, "SELECT * FROM customer WHERE custkey = 1 OR (name = 'test' AND address = 'test')"));
         assertThat(result.size()).isEqualTo(1);