[multistage][bugfix] fix tenant detection issues

pinot-common/src/main/java/org/apache/pinot/sql/parsers/CalciteSqlParser.java

@@ -146,12 +146,26 @@ public static List<String> extractTableNamesFromNode(SqlNode sqlNode) {
     if (sqlNode instanceof SqlSelect) {
       SqlNode fromNode = ((SqlSelect) sqlNode).getFrom();
       if (fromNode instanceof SqlJoin) {
-        tableNames.addAll(extractTableNamesFromNode(((SqlJoin) fromNode).getLeft()));
-        tableNames.addAll(extractTableNamesFromNode(((SqlJoin) fromNode).getRight()));
+        SqlNode left = ((SqlJoin) fromNode).getLeft();
+        SqlNode right = ((SqlJoin) fromNode).getRight();
+        if (left instanceof SqlIdentifier) {
+          tableNames.addAll(((SqlIdentifier) left).names);
+        } else {
+          tableNames.addAll(extractTableNamesFromNode(left));
+        }
+        if (right instanceof SqlIdentifier) {
+          tableNames.addAll(((SqlIdentifier) right).names);
+        } else {
+          tableNames.addAll(extractTableNamesFromNode(right));
+        }
       } else {
         tableNames.addAll(((SqlIdentifier) fromNode).names);
         tableNames.addAll(extractTableNamesFromNode(((SqlSelect) sqlNode).getWhere()));
       }
+    } else if (sqlNode instanceof SqlOrderBy) {
+      for (SqlNode node : ((SqlOrderBy) sqlNode).getOperandList()) {
+        tableNames.addAll(extractTableNamesFromNode(node));
+      }
     } else if (sqlNode instanceof SqlBasicCall) {
       if (((SqlBasicCall) sqlNode).getOperator() instanceof SqlAsOperator) {
         SqlNode firstOperand = ((SqlBasicCall) sqlNode).getOperandList().get(0);
@@ -161,10 +175,6 @@ public static List<String> extractTableNamesFromNode(SqlNode sqlNode) {
           tableNames.addAll(extractTableNamesFromNode(node));
         }
       }
-    } else if (sqlNode instanceof SqlOrderBy) {
-      for (SqlNode node : ((SqlOrderBy) sqlNode).getOperandList()) {
-        tableNames.addAll(extractTableNamesFromNode(node));
-      }
     } else if (sqlNode instanceof SqlWith) {
       List<SqlNode> withList = ((SqlWith) sqlNode).withList;
       Set<String> aliases = new HashSet<>();

pinot-common/src/test/java/org/apache/pinot/sql/parsers/CalciteSqlCompilerTest.java

@@ -3046,6 +3046,24 @@ public void testExtractTableNamesFromNode() {
     Assert.assertEquals(tableNames.get(2), "tbl3");
 
     // query with JOIN clause
+    query = "SELECT tbl1.col1, tbl2.col2 FROM tbl1 JOIN tbl2 ON tbl1.key = tbl2.key WHERE tbl1.col1 = value1";
+    sqlNodeAndOptions = RequestUtils.parseQuery(query);
+    tableNames = CalciteSqlParser.extractTableNamesFromNode(sqlNodeAndOptions.getSqlNode());
+    Assert.assertEquals(tableNames.size(), 2);
+    Collections.sort(tableNames);
+    Assert.assertEquals(tableNames.get(0), "tbl1");
+    Assert.assertEquals(tableNames.get(1), "tbl2");
+
+    // query with WHERE clause JOIN
+    query = "SELECT tbl1.col1, tbl2.col2 FROM tbl1, tbl2 WHERE tbl1.key = tbl2.key AND tbl1.col1 = value1";
+    sqlNodeAndOptions = RequestUtils.parseQuery(query);
+    tableNames = CalciteSqlParser.extractTableNamesFromNode(sqlNodeAndOptions.getSqlNode());
+    Assert.assertEquals(tableNames.size(), 2);
+    Collections.sort(tableNames);
+    Assert.assertEquals(tableNames.get(0), "tbl1");
+    Assert.assertEquals(tableNames.get(1), "tbl2");
+
+    // query with JOIN clause and table alias
     query = "SELECT A.col1, B.col2 FROM tbl1 AS A JOIN tbl2 AS B ON A.key = B.key WHERE A.col1 = value1";
     sqlNodeAndOptions = RequestUtils.parseQuery(query);
     tableNames = CalciteSqlParser.extractTableNamesFromNode(sqlNodeAndOptions.getSqlNode());
@@ -3066,6 +3084,19 @@ public void testExtractTableNamesFromNode() {
     Assert.assertEquals(tableNames.get(2), "tbl3");
     Assert.assertEquals(tableNames.get(3), "tbl4");
 
+    // query with aliases, JOIN, IN/NOT-IN, group-by
+    query = "with tmp as (select col1, count(*) from tbl1 where condition1 = filter1 group by col1 order by col2), "
+        + "tmp2 as (select A.col1, B.col2 from tbl2 as A JOIN tbl3 AS B on A.key = B.key) "
+        + "select sum(col1) from tmp where col1 in (select col1 from tmp2) and col1 not in (select col1 from tbl4) "
+        + "order by A.col1";
+    sqlNodeAndOptions = RequestUtils.parseQuery(query);
+    tableNames = CalciteSqlParser.extractTableNamesFromNode(sqlNodeAndOptions.getSqlNode());
+    Assert.assertEquals(tableNames.size(), 4);
+    Assert.assertEquals(tableNames.get(0), "tbl1");
+    Assert.assertEquals(tableNames.get(1), "tbl2");
+    Assert.assertEquals(tableNames.get(2), "tbl3");
+    Assert.assertEquals(tableNames.get(3), "tbl4");
+
     // query with aliases, JOIN, IN/NOT-IN, group-by and explain
     query = "explain plan for with tmp as (select col1, count(*) from tbl1 where condition1 = filter1 group by col1), "
         + "tmp2 as (select A.col1, B.col2 from tbl2 as A JOIN tbl3 AS B on A.key = B.key) "

pinot-controller/src/main/java/org/apache/pinot/controller/api/resources/PinotQueryResource.java

@@ -174,7 +174,8 @@ private String getMultiStageQueryResponse(String query, String queryOptions, Htt
     SqlNodeAndOptions sqlNodeAndOptions = RequestUtils.parseQuery(query);
     List<String> tableNames = CalciteSqlParser.extractTableNamesFromNode(sqlNodeAndOptions.getSqlNode());
     if (tableNames.size() == 0) {
-      return QueryException.SQL_PARSING_ERROR.toString();
+      return QueryException.getException(QueryException.SQL_PARSING_ERROR,
+          new Exception("Unable to find table name from SQL thus cannot dispatch to broker.")).toString();
     }
 
     String brokerTenant = getCommonBrokerTenant(tableNames);