[KYUUBI apache#4320] Support GetPrimaryKeys for Trino Fe

### _Why are the changes needed?_

Support GetPrimaryKeys for Trino Fe, close apache#4320

### _How was this patch tested?_
- [x] Add some test cases that check the changes thoroughly including negative and positive cases if possible

- [ ] Add screenshots for manual tests if appropriate

- [x] [Run test](https://kyuubi.readthedocs.io/en/master/develop_tools/testing.html#running-tests) locally before make a pull request

Closes apache#4321 from Yikf/primaryKeys.

Closes apache#4320

3690a2c [Yikf] Support GetPrimaryKeys for Trino Fe

Authored-by: Yikf <[email protected]>
Signed-off-by: ulyssesyou <[email protected]>

Author: yikf

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: '\'' '\\' '\'';
 

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
     ;
 

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)
     }

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
   }

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"
+}

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])
+  }
 }