[Enhancement] support Trino's show schemas function

Signed-off-by: Rohit Satardekar <[email protected]>
StarRocks · Sep 14, 2024 · 16de7d2 · 16de7d2
1 parent 4515b29
commit 16de7d2
Show file tree

Hide file tree

Showing 5 changed files with 33 additions and 3 deletions.
diff --git a/fe/fe-core/src/main/java/com/starrocks/connector/parser/trino/AstBuilder.java b/fe/fe-core/src/main/java/com/starrocks/connector/parser/trino/AstBuilder.java
@@ -83,6 +83,7 @@
 import com.starrocks.sql.ast.SelectListItem;
 import com.starrocks.sql.ast.SelectRelation;
 import com.starrocks.sql.ast.SetQualifier;
+import com.starrocks.sql.ast.ShowDbStmt;
 import com.starrocks.sql.ast.StatementBase;
 import com.starrocks.sql.ast.SubqueryRelation;
 import com.starrocks.sql.ast.TableFunctionRelation;
@@ -162,6 +163,7 @@
 import io.trino.sql.tree.RowDataType;
 import io.trino.sql.tree.SearchedCaseExpression;
 import io.trino.sql.tree.SetOperation;
+import io.trino.sql.tree.ShowSchemas;
 import io.trino.sql.tree.SimpleCaseExpression;
 import io.trino.sql.tree.SimpleGroupBy;
 import io.trino.sql.tree.SingleColumn;
@@ -739,6 +741,21 @@ protected ParseNode visitFunctionCall(FunctionCall node, ParseTreeContext contex
 
     }
 
+    @Override
+    protected ParseNode visitShowSchemas(ShowSchemas node, ParseTreeContext context) {
+        String catalog = null;
+        if (!node.getCatalog().isEmpty()) {
+            catalog = node.getCatalog().map(Identifier::getValue).get();
+        }
+
+        if (!node.getLikePattern().isEmpty()) {
+            String likePattern = node.getLikePattern().get();
+            return new ShowDbStmt(likePattern, null, catalog, null);
+        } else {
+            return new ShowDbStmt(null, null, catalog, null);
+        }
+    }
+
     private static AnalyticWindow.Type getFrameType(WindowFrame.Type type) {
         switch (type) {
             case RANGE:

diff --git a/fe/fe-core/src/main/java/com/starrocks/connector/parser/trino/TrinoParserUtils.java b/fe/fe-core/src/main/java/com/starrocks/connector/parser/trino/TrinoParserUtils.java
@@ -27,6 +27,7 @@
 import io.trino.sql.tree.Explain;
 import io.trino.sql.tree.ExplainAnalyze;
 import io.trino.sql.tree.Query;
+import io.trino.sql.tree.ShowSchemas;
 import io.trino.sql.tree.Statement;
 
 import java.time.format.DateTimeParseException;
@@ -40,7 +41,7 @@ public static StatementBase toStatement(String query, long sqlMode) {
         String trimmedQuery = query.trim();
         Statement statement = TrinoParser.parse(trimmedQuery);
         if (statement instanceof Query || statement instanceof Explain || statement instanceof ExplainAnalyze
-                || statement instanceof CreateTableAsSelect) {
+                || statement instanceof CreateTableAsSelect || statement instanceof ShowSchemas) {
             return (StatementBase) statement.accept(new AstBuilder(sqlMode), new ParseTreeContext());
         } else {
             throw trinoParserUnsupportedException("Unsupported statement type: " + statement.getClass().getName());

diff --git a/fe/fe-core/src/main/java/com/starrocks/sql/parser/StarRocks.g4 b/fe/fe-core/src/main/java/com/starrocks/sql/parser/StarRocks.g4
@@ -321,7 +321,7 @@ setCatalogStatement
 
 showDatabasesStatement
     : SHOW DATABASES ((FROM | IN) catalog=qualifiedName)? ((LIKE pattern=string) | (WHERE expression))?
-    | SHOW SCHEMAS ((LIKE pattern=string) | (WHERE expression))?
+    | SHOW SCHEMAS (FROM catalog=qualifiedName)? ((LIKE pattern=string) | (WHERE expression))?
     ;
 
 alterDbQuotaStatement

diff --git a/fe/fe-core/src/test/java/com/starrocks/analysis/ShowDbStmtTest.java b/fe/fe-core/src/test/java/com/starrocks/analysis/ShowDbStmtTest.java
@@ -111,4 +111,4 @@ public void testShowSchemas() throws Exception {
         ShowDbStmt showDbStmt = (ShowDbStmt) UtFrameUtils.parseStmtWithNewParser(showSQL, ctx);
     }
 
-}
+}
diff --git a/fe/fe-core/src/test/java/com/starrocks/connector/parser/trino/TrinoQueryTest.java b/fe/fe-core/src/test/java/com/starrocks/connector/parser/trino/TrinoQueryTest.java
@@ -14,6 +14,9 @@
 
 package com.starrocks.connector.parser.trino;
 
+import com.starrocks.common.util.UUIDUtil;
+import com.starrocks.sql.ast.ShowDbStmt;
+import com.starrocks.utframe.UtFrameUtils;
 import org.apache.commons.lang3.StringUtils;
 import org.junit.Assert;
 import org.junit.BeforeClass;
@@ -1228,4 +1231,13 @@ public void testCastArrayDataType() throws Exception {
         String sql = "select cast(ARRAY[1] as array(int))";
         assertPlanContains(sql, "CAST([1] AS ARRAY<INT>)");
     }
+
+    @Test
+    public void testShowSchemasFroCatalog() throws Exception {
+        connectContext.setExecutionId(UUIDUtil.toTUniqueId(UUIDUtil.genUUID()));
+        String showSQL = "show schemas from default_catalog";
+        ShowDbStmt showDbStmt = (ShowDbStmt) UtFrameUtils.parseStmtWithNewParser(showSQL, connectContext);
+        Assert.assertNotNull(showDbStmt);
+        Assert.assertEquals(1, showDbStmt.getMetaData().getColumnCount());
+    }
 }