From 79c34582538f0229c9e6be7f8079e6cc68f41838 Mon Sep 17 00:00:00 2001
From: JingShang Lu <lujingshang@apache.org>
Date: Thu, 24 Sep 2020 17:46:28 +0800
Subject: [PATCH] ref support of subquery (#7583)

* ref issue#6497

* delete useless import

* fix
---
 .../impl/ProjectionsTokenGenerator.java       | 30 ++++++++++++-------
 .../src/test/resources/sharding/select.xml    |  6 ++++
 .../statement/dml/SelectStatementContext.java |  2 +-
 .../parser/mysql/visitor/MySQLVisitor.java    |  2 +-
 4 files changed, 27 insertions(+), 13 deletions(-)

diff --git a/shardingsphere-features/shardingsphere-sharding/shardingsphere-sharding-rewrite/src/main/java/org/apache/shardingsphere/sharding/rewrite/token/generator/impl/ProjectionsTokenGenerator.java b/shardingsphere-features/shardingsphere-sharding/shardingsphere-sharding-rewrite/src/main/java/org/apache/shardingsphere/sharding/rewrite/token/generator/impl/ProjectionsTokenGenerator.java
index e635d26e7c..426e32cfce 100644
--- a/shardingsphere-features/shardingsphere-sharding/shardingsphere-sharding-rewrite/src/main/java/org/apache/shardingsphere/sharding/rewrite/token/generator/impl/ProjectionsTokenGenerator.java
+++ b/shardingsphere-features/shardingsphere-sharding/shardingsphere-sharding-rewrite/src/main/java/org/apache/shardingsphere/sharding/rewrite/token/generator/impl/ProjectionsTokenGenerator.java
@@ -67,17 +67,25 @@ public final class ProjectionsTokenGenerator implements OptionalSQLTokenGenerato
     private Map<RouteUnit, Collection<String>> getDerivedProjectionTexts(final SelectStatementContext selectStatementContext) {
         Map<RouteUnit, Collection<String>> result = new HashMap<>();
         for (RouteUnit routeUnit : routeContext.getRouteResult().getRouteUnits()) {
-            result.put(routeUnit, new LinkedList<>());
-            for (Projection each : selectStatementContext.getProjectionsContext().getProjections()) {
-                if (each instanceof AggregationProjection && !((AggregationProjection) each).getDerivedAggregationProjections().isEmpty()) {
-                    result.get(routeUnit).addAll(((AggregationProjection) each).getDerivedAggregationProjections().stream().map(this::getDerivedProjectionText).collect(Collectors.toList()));
-                } else if (each instanceof DerivedProjection && ((DerivedProjection) each).getDerivedProjection() instanceof ColumnOrderByItemSegment) {
-                    TableExtractor tableExtractor = new TableExtractor();
-                    tableExtractor.extractTablesFromSelect(selectStatementContext.getSqlStatement());
-                    result.get(routeUnit).add(getDerivedProjectionTextFromColumnOrderByItemSegment((DerivedProjection) each, tableExtractor, routeUnit));
-                } else if (each instanceof DerivedProjection) {
-                    result.get(routeUnit).add(getDerivedProjectionText(each));
-                }
+            Collection<String> projectionTexts = getDerivedProjectionTextsByRouteUnit(selectStatementContext, routeUnit);
+            if (!projectionTexts.isEmpty()) {
+                result.put(routeUnit, projectionTexts);
+            }
+        }
+        return result;
+    }
+    
+    private Collection<String> getDerivedProjectionTextsByRouteUnit(final SelectStatementContext selectStatementContext, final RouteUnit routeUnit) {
+        Collection<String> result = new LinkedList<>();
+        for (Projection each : selectStatementContext.getProjectionsContext().getProjections()) {
+            if (each instanceof AggregationProjection && !((AggregationProjection) each).getDerivedAggregationProjections().isEmpty()) {
+                result.addAll(((AggregationProjection) each).getDerivedAggregationProjections().stream().map(this::getDerivedProjectionText).collect(Collectors.toList()));
+            } else if (each instanceof DerivedProjection && ((DerivedProjection) each).getDerivedProjection() instanceof ColumnOrderByItemSegment) {
+                TableExtractor tableExtractor = new TableExtractor();
+                tableExtractor.extractTablesFromSelect(selectStatementContext.getSqlStatement());
+                result.add(getDerivedProjectionTextFromColumnOrderByItemSegment((DerivedProjection) each, tableExtractor, routeUnit));
+            } else if (each instanceof DerivedProjection) {
+                result.add(getDerivedProjectionText(each));
             }
         }
         return result;
diff --git a/shardingsphere-features/shardingsphere-sharding/shardingsphere-sharding-rewrite/src/test/resources/sharding/select.xml b/shardingsphere-features/shardingsphere-sharding/shardingsphere-sharding-rewrite/src/test/resources/sharding/select.xml
index b617724f9c..c01a58d0b4 100644
--- a/shardingsphere-features/shardingsphere-sharding/shardingsphere-sharding-rewrite/src/test/resources/sharding/select.xml
+++ b/shardingsphere-features/shardingsphere-sharding/shardingsphere-sharding-rewrite/src/test/resources/sharding/select.xml
@@ -64,6 +64,12 @@
         <output sql="SELECT (select id from t_account_0 limit 1) as myid FROM (select b.account_id from (select t_account_0.account_id from t_account_0) b where b.account_id=?) a WHERE account_id >= (select account_id from t_account_0 limit 1)"  parameters="100"/>
     </rewrite-assertion>
 
+    <rewrite-assertion id="select_with_subquery_only_in_projection" db-type="MySQL">
+        <input sql="SELECT (select id from t_account)"/>
+        <output sql="SELECT (select id from t_account_0)"/>
+        <output sql="SELECT (select id from t_account_1)"/>
+    </rewrite-assertion>
+
     <rewrite-assertion id="select_with_subquery_for_where_in_predicate" db-type="MySQL">
         <input sql="SELECT * FROM t_account WHERE account_id = ? AND amount IN (SELECT amount FROM t_account WHERE account_id = ?)"  parameters="100, 100"/>
         <output sql="SELECT * FROM t_account_0 WHERE account_id = ? AND amount IN (SELECT amount FROM t_account_0 WHERE account_id = ?)"  parameters="100, 100"/>
diff --git a/shardingsphere-sql-parser/shardingsphere-sql-parser-binder/src/main/java/org/apache/shardingsphere/sql/parser/binder/statement/dml/SelectStatementContext.java b/shardingsphere-sql-parser/shardingsphere-sql-parser-binder/src/main/java/org/apache/shardingsphere/sql/parser/binder/statement/dml/SelectStatementContext.java
index 157509aad7..90497400ae 100644
--- a/shardingsphere-sql-parser/shardingsphere-sql-parser-binder/src/main/java/org/apache/shardingsphere/sql/parser/binder/statement/dml/SelectStatementContext.java
+++ b/shardingsphere-sql-parser/shardingsphere-sql-parser-binder/src/main/java/org/apache/shardingsphere/sql/parser/binder/statement/dml/SelectStatementContext.java
@@ -36,6 +36,7 @@ import org.apache.shardingsphere.sql.parser.binder.segment.table.TablesContext;
 import org.apache.shardingsphere.sql.parser.binder.statement.CommonSQLStatementContext;
 import org.apache.shardingsphere.sql.parser.binder.type.TableAvailable;
 import org.apache.shardingsphere.sql.parser.binder.type.WhereAvailable;
+import org.apache.shardingsphere.sql.parser.sql.common.extractor.TableExtractor;
 import org.apache.shardingsphere.sql.parser.sql.common.segment.dml.order.item.ColumnOrderByItemSegment;
 import org.apache.shardingsphere.sql.parser.sql.common.segment.dml.order.item.ExpressionOrderByItemSegment;
 import org.apache.shardingsphere.sql.parser.sql.common.segment.dml.order.item.IndexOrderByItemSegment;
@@ -44,7 +45,6 @@ import org.apache.shardingsphere.sql.parser.sql.common.segment.dml.predicate.Whe
 import org.apache.shardingsphere.sql.parser.sql.common.segment.generic.table.SimpleTableSegment;
 import org.apache.shardingsphere.sql.parser.sql.common.statement.dml.SelectStatement;
 import org.apache.shardingsphere.sql.parser.sql.common.util.SQLUtil;
-import org.apache.shardingsphere.sql.parser.sql.common.extractor.TableExtractor;
 import org.apache.shardingsphere.sql.parser.sql.common.util.WhereSegmentExtractUtils;
 
 import java.util.Collection;
diff --git a/shardingsphere-sql-parser/shardingsphere-sql-parser-dialect/shardingsphere-sql-parser-mysql/src/main/java/org/apache/shardingsphere/sql/parser/mysql/visitor/MySQLVisitor.java b/shardingsphere-sql-parser/shardingsphere-sql-parser-dialect/shardingsphere-sql-parser-mysql/src/main/java/org/apache/shardingsphere/sql/parser/mysql/visitor/MySQLVisitor.java
index 73054929df..c6e3d0c18c 100644
--- a/shardingsphere-sql-parser/shardingsphere-sql-parser-dialect/shardingsphere-sql-parser-mysql/src/main/java/org/apache/shardingsphere/sql/parser/mysql/visitor/MySQLVisitor.java
+++ b/shardingsphere-sql-parser/shardingsphere-sql-parser-dialect/shardingsphere-sql-parser-mysql/src/main/java/org/apache/shardingsphere/sql/parser/mysql/visitor/MySQLVisitor.java
@@ -303,7 +303,7 @@ public abstract class MySQLVisitor extends MySQLStatementBaseVisitor<ASTNode> {
         if (null != ctx.predicate()) {
             right = (ExpressionSegment) visit(ctx.predicate());
         } else {
-            right = (ExpressionSegment) visit(ctx.subquery());
+            right = new SubqueryExpressionSegment(new SubquerySegment(ctx.subquery().start.getStartIndex(), ctx.subquery().stop.getStopIndex(), (MySQLSelectStatement) visit(ctx.subquery())));
         }
         String operator = null != ctx.SAFE_EQ_() ? ctx.SAFE_EQ_().getText() : ctx.comparisonOperator().getText();
         String text = ctx.start.getInputStream().getText(new Interval(ctx.start.getStartIndex(), ctx.stop.getStopIndex()));
-- 
GitLab