Skip to content

Commit 543e32c

Browse files
committed
Refactor SQL table extraction to use SqlBasicVisitor pattern
Replace instanceof checks with SqlTableNameExtractor visitor, consistent with how PPL uses AbstractNodeVisitor for index name extraction. Signed-off-by: Kai Huang <ahkcs@amazon.com>
1 parent 617cbed commit 543e32c

1 file changed

Lines changed: 25 additions & 17 deletions

File tree

plugin/src/main/java/org/opensearch/sql/plugin/rest/RestUnifiedQueryAction.java

Lines changed: 25 additions & 17 deletions
Original file line numberDiff line numberDiff line change
@@ -13,6 +13,7 @@
1313
import java.util.Map;
1414
import java.util.Optional;
1515
import org.apache.calcite.rel.RelNode;
16+
import org.apache.calcite.sql.SqlCall;
1617
import org.apache.calcite.sql.SqlIdentifier;
1718
import org.apache.calcite.sql.SqlJoin;
1819
import org.apache.calcite.sql.SqlNode;
@@ -179,31 +180,38 @@ private static Optional<String> extractIndexName(
179180
return Optional.ofNullable(extractTableNameFromSqlNode(sqlNode));
180181
}
181182

182-
/** Extracts the table name from a Calcite SqlNode parse tree. */
183-
private static String extractTableNameFromSqlNode(SqlNode sqlNode) {
184-
if (sqlNode instanceof SqlSelect select) {
185-
SqlNode from = select.getFrom();
186-
if (from instanceof SqlIdentifier id) {
187-
return id.toString();
183+
/** AST visitor that extracts the source index name from a Relation node (PPL path). */
184+
private static class IndexNameExtractor extends AbstractNodeVisitor<String, Void> {
185+
@Override
186+
public String visitRelation(Relation node, Void context) {
187+
return node.getTableQualifiedName().toString();
188+
}
189+
}
190+
191+
/** SqlNode visitor that extracts the source table name from a SQL parse tree. */
192+
private static class SqlTableNameExtractor
193+
extends org.apache.calcite.sql.util.SqlBasicVisitor<String> {
194+
@Override
195+
public String visit(SqlCall call) {
196+
if (call instanceof SqlSelect select) {
197+
return select.getFrom().accept(this);
188198
}
189-
if (from instanceof SqlJoin join) {
190-
// For joins, extract from the left table
191-
if (join.getLeft() instanceof SqlIdentifier leftId) {
192-
return leftId.toString();
193-
}
199+
if (call instanceof SqlJoin join) {
200+
return join.getLeft().accept(this);
194201
}
202+
return null;
195203
}
196-
return null;
197-
}
198204

199-
/** AST visitor that extracts the source index name from a Relation node. */
200-
private static class IndexNameExtractor extends AbstractNodeVisitor<String, Void> {
201205
@Override
202-
public String visitRelation(Relation node, Void context) {
203-
return node.getTableQualifiedName().toString();
206+
public String visit(SqlIdentifier id) {
207+
return id.toString();
204208
}
205209
}
206210

211+
private static String extractTableNameFromSqlNode(SqlNode sqlNode) {
212+
return sqlNode.accept(new SqlTableNameExtractor());
213+
}
214+
207215
private static RelNode addQuerySizeLimit(RelNode plan, CalcitePlanContext context) {
208216
return LogicalSystemLimit.create(
209217
LogicalSystemLimit.SystemLimitType.QUERY_SIZE_LIMIT,

0 commit comments

Comments
 (0)