diff --git a/kyuubi-server/src/main/antlr4/org/apache/kyuubi/sql/KyuubiTrinoFeBaseLexer.g4 b/kyuubi-server/src/main/antlr4/org/apache/kyuubi/sql/KyuubiTrinoFeBaseLexer.g4 index 0b9543a430c..0cc02de6435 100644 --- a/kyuubi-server/src/main/antlr4/org/apache/kyuubi/sql/KyuubiTrinoFeBaseLexer.g4 +++ b/kyuubi-server/src/main/antlr4/org/apache/kyuubi/sql/KyuubiTrinoFeBaseLexer.g4 @@ -97,6 +97,12 @@ SCOPE_TABLE: 'SCOPE_TABLE'; SOURCE_DATA_TYPE: 'SOURCE_DATA_TYPE'; IS_AUTOINCREMENT: 'IS_AUTOINCREMENT'; IS_GENERATEDCOLUMN: 'IS_GENERATEDCOLUMN'; +VARCHAR: 'VARCHAR'; +SMALLINT: 'SMALLINT'; +CAST: 'CAST'; +AS: 'AS'; +KEY_SEQ: 'KEY_SEQ'; +PK_NAME: 'PK_NAME'; fragment SEARCH_STRING_ESCAPE: '\'' '\\' '\''; diff --git a/kyuubi-server/src/main/antlr4/org/apache/kyuubi/sql/KyuubiTrinoFeBaseParser.g4 b/kyuubi-server/src/main/antlr4/org/apache/kyuubi/sql/KyuubiTrinoFeBaseParser.g4 index 590c4378d52..6af00af5de5 100644 --- a/kyuubi-server/src/main/antlr4/org/apache/kyuubi/sql/KyuubiTrinoFeBaseParser.g4 +++ b/kyuubi-server/src/main/antlr4/org/apache/kyuubi/sql/KyuubiTrinoFeBaseParser.g4 @@ -47,6 +47,13 @@ statement SOURCE_DATA_TYPE COMMA IS_AUTOINCREMENT COMMA IS_GENERATEDCOLUMN FROM SYSTEM_JDBC_COLUMNS (WHERE tableCatalogFilter? AND? tableSchemaFilter? AND? tableNameFilter? AND? colNameFilter?)? ORDER BY TABLE_CAT COMMA TABLE_SCHEM COMMA TABLE_NAME COMMA ORDINAL_POSITION #getColumns + | SELECT CAST LEFT_PAREN NULL AS VARCHAR RIGHT_PAREN TABLE_CAT COMMA + CAST LEFT_PAREN NULL AS VARCHAR RIGHT_PAREN TABLE_SCHEM COMMA + CAST LEFT_PAREN NULL AS VARCHAR RIGHT_PAREN TABLE_NAME COMMA + CAST LEFT_PAREN NULL AS VARCHAR RIGHT_PAREN COLUMN_NAME COMMA + CAST LEFT_PAREN NULL AS SMALLINT RIGHT_PAREN KEY_SEQ COMMA + CAST LEFT_PAREN NULL AS VARCHAR RIGHT_PAREN PK_NAME + WHERE FALSE #getPrimaryKeys | .*? #passThrough ; diff --git a/kyuubi-server/src/main/scala/org/apache/kyuubi/server/trino/api/KyuubiTrinoOperationTranslator.scala b/kyuubi-server/src/main/scala/org/apache/kyuubi/server/trino/api/KyuubiTrinoOperationTranslator.scala index 6ec9fc1c80e..84686f6b035 100644 --- a/kyuubi-server/src/main/scala/org/apache/kyuubi/server/trino/api/KyuubiTrinoOperationTranslator.scala +++ b/kyuubi-server/src/main/scala/org/apache/kyuubi/server/trino/api/KyuubiTrinoOperationTranslator.scala @@ -25,7 +25,7 @@ import org.apache.kyuubi.operation.OperationHandle import org.apache.kyuubi.service.BackendService import org.apache.kyuubi.sql.parser.trino.KyuubiTrinoFeParser import org.apache.kyuubi.sql.plan.PassThroughNode -import org.apache.kyuubi.sql.plan.trino.{GetCatalogs, GetColumns, GetSchemas, GetTables, GetTableTypes, GetTypeInfo} +import org.apache.kyuubi.sql.plan.trino.{GetCatalogs, GetColumns, GetPrimaryKeys, GetSchemas, GetTables, GetTableTypes, GetTypeInfo} class KyuubiTrinoOperationTranslator(backendService: BackendService) { lazy val parser = new KyuubiTrinoFeParser() @@ -68,6 +68,11 @@ class KyuubiTrinoOperationTranslator(backendService: BackendService) { schemaPattern, tableNamePattern, colNamePattern) + case GetPrimaryKeys() => + val operationHandle = backendService.getPrimaryKeys(sessionHandle, null, null, null) + // The trino implementation always returns empty. + operationHandle.setHasResultSet(false) + operationHandle case PassThroughNode() => backendService.executeStatement(sessionHandle, statement, configs, runAsync, queryTimeout) } diff --git a/kyuubi-server/src/main/scala/org/apache/kyuubi/sql/parser/trino/KyuubiTrinoFeAstBuilder.scala b/kyuubi-server/src/main/scala/org/apache/kyuubi/sql/parser/trino/KyuubiTrinoFeAstBuilder.scala index c5ae9719947..061985c1caf 100644 --- a/kyuubi-server/src/main/scala/org/apache/kyuubi/sql/parser/trino/KyuubiTrinoFeAstBuilder.scala +++ b/kyuubi-server/src/main/scala/org/apache/kyuubi/sql/parser/trino/KyuubiTrinoFeAstBuilder.scala @@ -25,7 +25,7 @@ import org.apache.kyuubi.sql.KyuubiTrinoFeBaseParser._ import org.apache.kyuubi.sql.KyuubiTrinoFeBaseParserBaseVisitor import org.apache.kyuubi.sql.parser.KyuubiParser.unescapeSQLString import org.apache.kyuubi.sql.plan.{KyuubiTreeNode, PassThroughNode} -import org.apache.kyuubi.sql.plan.trino.{GetCatalogs, GetColumns, GetSchemas, GetTables, GetTableTypes, GetTypeInfo} +import org.apache.kyuubi.sql.plan.trino.{GetCatalogs, GetColumns, GetPrimaryKeys, GetSchemas, GetTables, GetTableTypes, GetTypeInfo} class KyuubiTrinoFeAstBuilder extends KyuubiTrinoFeBaseParserBaseVisitor[AnyRef] { @@ -92,6 +92,10 @@ class KyuubiTrinoFeAstBuilder extends KyuubiTrinoFeBaseParserBaseVisitor[AnyRef] GetColumns(catalog, schemaPattern, tableNamePattern, colNamePattern) } + override def visitGetPrimaryKeys(ctx: GetPrimaryKeysContext): KyuubiTreeNode = { + GetPrimaryKeys() + } + override def visitNullCatalog(ctx: NullCatalogContext): AnyRef = { null } diff --git a/kyuubi-server/src/main/scala/org/apache/kyuubi/sql/plan/trino/TrinoFeOperations.scala b/kyuubi-server/src/main/scala/org/apache/kyuubi/sql/plan/trino/TrinoFeOperations.scala index 85e6f168bcb..6136995ab10 100644 --- a/kyuubi-server/src/main/scala/org/apache/kyuubi/sql/plan/trino/TrinoFeOperations.scala +++ b/kyuubi-server/src/main/scala/org/apache/kyuubi/sql/plan/trino/TrinoFeOperations.scala @@ -55,3 +55,7 @@ case class GetColumns( colNamePattern: String) extends KyuubiTreeNode { override def name(): String = "Get Columns" } + +case class GetPrimaryKeys() extends KyuubiTreeNode { + override def name(): String = "Get Primary Keys" +} diff --git a/kyuubi-server/src/test/scala/org/apache/kyuubi/parser/trino/KyuubiTrinoFeParserSuite.scala b/kyuubi-server/src/test/scala/org/apache/kyuubi/parser/trino/KyuubiTrinoFeParserSuite.scala index 3f5cf70b559..bbced0b61ad 100644 --- a/kyuubi-server/src/test/scala/org/apache/kyuubi/parser/trino/KyuubiTrinoFeParserSuite.scala +++ b/kyuubi-server/src/test/scala/org/apache/kyuubi/parser/trino/KyuubiTrinoFeParserSuite.scala @@ -20,7 +20,7 @@ package org.apache.kyuubi.parser.trino import org.apache.kyuubi.KyuubiFunSuite import org.apache.kyuubi.sql.parser.trino.KyuubiTrinoFeParser import org.apache.kyuubi.sql.plan.{KyuubiTreeNode, PassThroughNode} -import org.apache.kyuubi.sql.plan.trino.{GetCatalogs, GetColumns, GetSchemas, GetTables, GetTableTypes, GetTypeInfo} +import org.apache.kyuubi.sql.plan.trino.{GetCatalogs, GetColumns, GetPrimaryKeys, GetSchemas, GetTables, GetTableTypes, GetTypeInfo} class KyuubiTrinoFeParserSuite extends KyuubiFunSuite { val parser = new KyuubiTrinoFeParser() @@ -354,4 +354,19 @@ class KyuubiTrinoFeParserSuite extends KyuubiFunSuite { tableName = "%aa", colName = "%bb") } + + test("Support GetPrimaryKeys for Trino Fe") { + val kyuubiTreeNode = parse( + """ + | SELECT CAST(NULL AS varchar) TABLE_CAT, + | CAST(NULL AS varchar) TABLE_SCHEM, + | CAST(NULL AS varchar) TABLE_NAME, + | CAST(NULL AS varchar) COLUMN_NAME, + | CAST(NULL AS smallint) KEY_SEQ, + | CAST(NULL AS varchar) PK_NAME + | WHERE false + |""".stripMargin) + + assert(kyuubiTreeNode.isInstanceOf[GetPrimaryKeys]) + } }