[KYUUBI #6995] Support Flink 2.0, 2.1, 2.2 and 2.3

.github/workflows/master.yml

@@ -334,6 +334,22 @@ jobs:
             flink: '1.20'
             flink-archive: '-Dflink.archive.mirror=https://www.apache.org/dyn/closer.lua/flink/flink-1.19.3 -Dflink.archive.name=flink-1.19.3-bin-scala_2.12.tgz'
             comment: 'verify-on-flink-1.19-binary'
+          - java: 17
+            flink: '1.20'
+            flink-archive: '-Dflink.archive.mirror=https://www.apache.org/dyn/closer.lua/flink/flink-2.0.2 -Dflink.archive.name=flink-2.0.2-bin-scala_2.12.tgz'
+            comment: 'verify-on-flink-2.0-binary'
+          - java: 17
+            flink: '1.20'
+            flink-archive: '-Dflink.archive.mirror=https://www.apache.org/dyn/closer.lua/flink/flink-2.1.3 -Dflink.archive.name=flink-2.1.3-bin-scala_2.12.tgz'
+            comment: 'verify-on-flink-2.1-binary'
+          - java: 17
+            flink: '1.20'
+            flink-archive: '-Dflink.archive.mirror=https://www.apache.org/dyn/closer.lua/flink/flink-2.2.1 -Dflink.archive.name=flink-2.2.1-bin-scala_2.12.tgz'
+            comment: 'verify-on-flink-2.2-binary'
+          - java: 17
+            flink: '1.20'
+            flink-archive: '-Dflink.archive.mirror=https://www.apache.org/dyn/closer.lua/flink/flink-2.3.0 -Dflink.archive.name=flink-2.3.0-bin-scala_2.12.tgz'
+            comment: 'verify-on-flink-2.3-binary'
     steps:
       - uses: actions/checkout@v7
       - name: Free up disk space

docs/quick_start/quick_start.rst

@@ -44,7 +44,7 @@ pre-installed and the ``JAVA_HOME`` is correctly set to each component.
                    Engine lib                            - Kyuubi Engine
                    Beeline                               - Kyuubi Beeline
   **Spark**        Engine       3.3 to 3.5, 4.0 to 4.2   A Spark distribution
-  **Flink**        Engine       1.17 to 1.20             A Flink distribution
+  **Flink**        Engine       1.17 to 1.20, 2.0 to 2.3 A Flink distribution
   **Trino**        Engine       N/A                      A Trino cluster allows to access via trino-client v411
   **Doris**        Engine       N/A                      A Doris cluster
   **Hive**         Engine       - 2.1-cdh6, 2.3, 3.1     - A Hive distribution

externals/kyuubi-flink-sql-engine/src/main/java/org/apache/flink/client/deployment/application/executors/EmbeddedExecutorFactory.java

@@ -20,30 +20,120 @@
 import static org.apache.flink.util.Preconditions.checkNotNull;
 import static org.apache.flink.util.Preconditions.checkState;
 
+import java.time.Duration;
 import java.util.Collection;
+import java.util.Optional;
 import java.util.concurrent.ConcurrentLinkedQueue;
 import org.apache.flink.annotation.Internal;
 import org.apache.flink.api.common.JobID;
-import org.apache.flink.api.common.time.Time;
 import org.apache.flink.client.cli.ClientOptions;
 import org.apache.flink.client.deployment.application.EmbeddedJobClient;
 import org.apache.flink.configuration.Configuration;
 import org.apache.flink.configuration.DeploymentOptions;
+import org.apache.flink.core.execution.JobClient;
 import org.apache.flink.core.execution.PipelineExecutor;
 import org.apache.flink.core.execution.PipelineExecutorFactory;
 import org.apache.flink.runtime.dispatcher.DispatcherGateway;
+import org.apache.flink.streaming.api.graph.StreamGraph;
 import org.apache.flink.util.concurrent.ScheduledExecutor;
+import org.apache.kyuubi.util.reflect.DynClasses;
+import org.apache.kyuubi.util.reflect.DynConstructors;
+import org.apache.kyuubi.util.reflect.DynMethods;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
 
-/** Copied from Apache Flink to exposed the DispatcherGateway for Kyuubi statements. */
+/**
+ * Copied from Apache Flink to expose the DispatcherGateway for Kyuubi statements and stamp the
+ * application id on submitted StreamGraphs, which is required since FLINK-38974 (2.3.0).
+ */
 @Internal
 public class EmbeddedExecutorFactory implements PipelineExecutorFactory {
 
+  /** FLINK-14068 (2.0.0) removed {@code Time} in favor of {@link Duration}. */
+  private static final Class<?> LEGACY_TIME_CLASS =
+      DynClasses.builder().impl("org.apache.flink.api.common.time.Time").orNull().build();
+
+  /** FLINK-38974 (2.3.0) introduced {@code ApplicationID}, absent in older Flink versions. */
+  private static final Class<?> APPLICATION_ID_CLASS =
+      DynClasses.builder().impl("org.apache.flink.api.common.ApplicationID").orNull().build();
+
+  private static final boolean IS_FLINK_1 = LEGACY_TIME_CLASS != null;
+
+  /**
+   * EmbeddedExecutor's constructor changed twice, FLINK-33212 (2.0.0) added the {@link
+   * Configuration} parameter, FLINK-38974 (2.3.0) split the job ids into the application, suspended
+   * and terminal ones. Bind all variants reflectively so that a single engine jar runs on all
+   * supported Flink versions.
+   */
+  private static final DynConstructors.Ctor<EmbeddedExecutor> EMBEDDED_EXECUTOR_CTOR =
+      DynConstructors.builder()
+          .impl(
+              EmbeddedExecutor.class,
+              Collection.class,
+              Collection.class,
+              Collection.class,
+              DispatcherGateway.class,
+              Configuration.class,
+              EmbeddedJobClientCreator.class)
+          .impl(
+              EmbeddedExecutor.class,
+              Collection.class,
+              DispatcherGateway.class,
+              Configuration.class,
+              EmbeddedJobClientCreator.class)
+          .impl(
+              EmbeddedExecutor.class,
+              Collection.class,
+              DispatcherGateway.class,
+              EmbeddedJobClientCreator.class)
+          .build();
+
+  /** FLINK-14068 (2.0.0) replaced the {@code Time} timeout parameter with {@link Duration}. */
+  private static final DynConstructors.Ctor<EmbeddedJobClient> EMBEDDED_JOB_CLIENT_CTOR =
+      DynConstructors.builder()
+          .impl(
+              EmbeddedJobClient.class,
+              JobID.class,
+              DispatcherGateway.class,
+              ScheduledExecutor.class,
+              Duration.class,
+              ClassLoader.class)
+          .impl(
+              EmbeddedJobClient.class,
+              JobID.class,
+              DispatcherGateway.class,
+              ScheduledExecutor.class,
+              LEGACY_TIME_CLASS,
+              ClassLoader.class)
+          .build();
+
+  private static final DynMethods.StaticMethod LEGACY_TIME_OF_MILLIS =
+      IS_FLINK_1
+          ? DynMethods.builder("milliseconds").impl(LEGACY_TIME_CLASS, long.class).buildStatic()
+          : null;
+
   private static Collection<JobID> bootstrapJobIds;
 
   private static Collection<JobID> submittedJobIds;
 
+  /**
+   * Both are null before FLINK-38974 (2.3.0), which is also how the Flink version is told apart.
+   */
+  private static Collection<JobID> suspendedJobIds;
+
+  private static Collection<JobID> terminalJobIds;
+
+  /**
+   * FLINK-38974 (2.3.0): the dispatcher rejects jobs that carry no application registered in it.
+   * Kyuubi statements run through the SQL gateway, which builds plain StreamExecutionEnvironments,
+   * so capture the id from the first StreamGraph that carries it and stamp it on the rest.
+   */
+  private static volatile Object applicationId;
+
+  private static volatile DynMethods.UnboundMethod streamGraphGetApplicationId;
+
+  private static volatile DynMethods.UnboundMethod streamGraphSetApplicationId;
+
   private static DispatcherGateway dispatcherGateway;
 
   private static ScheduledExecutor retryExecutor;
@@ -65,7 +155,7 @@ public EmbeddedExecutorFactory() {
   }
 
   /**
-   * Creates an {@link EmbeddedExecutorFactory}.
+   * Creates an {@link EmbeddedExecutorFactory}, invoked by Flink before FLINK-38974 (2.3.0).
    *
    * @param submittedJobIds a list that is going to be filled with the job ids of the new jobs that
    *     will be submitted. This is essentially used to return the submitted job ids to the caller.
@@ -76,6 +166,28 @@ public EmbeddedExecutorFactory(
       final Collection<JobID> submittedJobIds,
       final DispatcherGateway dispatcherGateway,
       final ScheduledExecutor retryExecutor) {
+    this(submittedJobIds, null, null, dispatcherGateway, retryExecutor);
+  }
+
+  /**
+   * Creates an {@link EmbeddedExecutorFactory}, invoked by Flink since FLINK-38974 (2.3.0). Flink
+   * resolves this constructor at compile time, so it has to exist even though Kyuubi never calls it
+   * directly.
+   *
+   * @param applicationJobIds a list that is going to be filled with the job ids of the new jobs
+   *     that will be submitted. This is essentially used to return the submitted job ids to the
+   *     caller.
+   * @param suspendedJobIds ids of jobs suspended by a previous application execution.
+   * @param terminalJobIds ids of jobs already terminated by a previous application execution.
+   * @param dispatcherGateway the dispatcher of the cluster which is going to be used to submit
+   *     jobs.
+   */
+  public EmbeddedExecutorFactory(
+      final Collection<JobID> applicationJobIds,
+      final Collection<JobID> suspendedJobIds,
+      final Collection<JobID> terminalJobIds,
+      final DispatcherGateway dispatcherGateway,
+      final ScheduledExecutor retryExecutor) {
     // there should be only one instance of EmbeddedExecutorFactory
     LOGGER.debug(
         "{} initiated in thread {} with classloader {}.",
@@ -90,8 +202,10 @@ public EmbeddedExecutorFactory(
       // issues
       LOGGER.debug("Bootstrapping EmbeddedExecutorFactory.");
       EmbeddedExecutorFactory.submittedJobIds =
-          new ConcurrentLinkedQueue<>(checkNotNull(submittedJobIds));
-      EmbeddedExecutorFactory.bootstrapJobIds = submittedJobIds;
+          new ConcurrentLinkedQueue<>(checkNotNull(applicationJobIds));
+      EmbeddedExecutorFactory.bootstrapJobIds = applicationJobIds;
+      EmbeddedExecutorFactory.suspendedJobIds = suspendedJobIds;
+      EmbeddedExecutorFactory.terminalJobIds = terminalJobIds;
       EmbeddedExecutorFactory.dispatcherGateway = checkNotNull(dispatcherGateway);
       EmbeddedExecutorFactory.retryExecutor = checkNotNull(retryExecutor);
       bootstrapLock.notifyAll();
@@ -141,14 +255,94 @@ public PipelineExecutor getExecutor(final Configuration configuration) {
       LOGGER.info("Bootstrapping Flink SQL engine with the initial SQL.");
       executorJobIDs = bootstrapJobIds;
     }
-    return new EmbeddedExecutor(
-        executorJobIDs,
-        dispatcherGateway,
-        (jobId, userCodeClassloader) -> {
-          final Time timeout =
-              Time.milliseconds(configuration.get(ClientOptions.CLIENT_TIMEOUT).toMillis());
-          return new EmbeddedJobClient(
-              jobId, dispatcherGateway, retryExecutor, timeout, userCodeClassloader);
-        });
+    final EmbeddedJobClientCreator jobClientCreator =
+        (jobId, userCodeClassloader) ->
+            newEmbeddedJobClient(
+                jobId, configuration.get(ClientOptions.CLIENT_TIMEOUT), userCodeClassloader);
+    return stampApplicationId(newEmbeddedExecutor(executorJobIDs, configuration, jobClientCreator));
+  }
+
+  private static PipelineExecutor newEmbeddedExecutor(
+      final Collection<JobID> jobIds,
+      final Configuration configuration,
+      final EmbeddedJobClientCreator jobClientCreator) {
+    if (suspendedJobIds != null) {
+      return EMBEDDED_EXECUTOR_CTOR.newInstance(
+          jobIds,
+          suspendedJobIds,
+          terminalJobIds,
+          dispatcherGateway,
+          configuration,
+          jobClientCreator);
+    }
+    return IS_FLINK_1
+        ? EMBEDDED_EXECUTOR_CTOR.newInstance(jobIds, dispatcherGateway, jobClientCreator)
+        : EMBEDDED_EXECUTOR_CTOR.newInstance(
+            jobIds, dispatcherGateway, configuration, jobClientCreator);
+  }
+
+  /**
+   * FLINK-38974 (2.3.0) requires submitted jobs to carry the id of the application registered in
+   * the dispatcher. The bootstrap SQL goes through StreamContextEnvironment which stamps the id,
+   * while Kyuubi statements go through plain environments, so cache the id and stamp it on every
+   * StreamGraph before delegating to the {@link EmbeddedExecutor}.
+   */
+  private static PipelineExecutor stampApplicationId(final PipelineExecutor executor) {
+    if (suspendedJobIds == null) {
+      return executor;
+    }
+    return (pipeline, configuration, userCodeClassloader) -> {
+      if (pipeline instanceof StreamGraph) {
+        final StreamGraph streamGraph = (StreamGraph) pipeline;
+        final Object appId = captureApplicationId(streamGraph);
+        if (appId != null) {
+          streamGraphSetApplicationId(streamGraph, appId);
+        }
+      }
+      return executor.execute(pipeline, configuration, userCodeClassloader);
+    };
+  }
+
+  private static Object captureApplicationId(final StreamGraph streamGraph) {
+    Object appId = applicationId;
+    if (appId == null) {
+      synchronized (bootstrapLock) {
+        appId = applicationId;
+        if (appId == null) {
+          initStreamGraphApplicationIdMethods();
+          applicationId = appId = streamGraphGetApplicationId(streamGraph);
+        }
+      }
+    }
+    return appId;
+  }
+
+  @SuppressWarnings("unchecked")
+  private static Object streamGraphGetApplicationId(final StreamGraph streamGraph) {
+    return ((Optional<Object>) streamGraphGetApplicationId.invoke(streamGraph)).orElse(null);
+  }
+
+  private static void streamGraphSetApplicationId(
+      final StreamGraph streamGraph, final Object applicationId) {
+    streamGraphSetApplicationId.invoke(streamGraph, applicationId);
+  }
+
+  private static void initStreamGraphApplicationIdMethods() {
+    if (streamGraphGetApplicationId == null) {
+      streamGraphGetApplicationId =
+          DynMethods.builder("getApplicationId").impl(StreamGraph.class).build();
+      streamGraphSetApplicationId =
+          DynMethods.builder("setApplicationId")
+              .impl(StreamGraph.class, APPLICATION_ID_CLASS)
+              .build();
+    }
+  }
+
+  private static JobClient newEmbeddedJobClient(
+      final JobID jobId, final Duration timeout, final ClassLoader userCodeClassloader) {
+    final Object rpcTimeout =
+        IS_FLINK_1 ? LEGACY_TIME_OF_MILLIS.invoke(timeout.toMillis()) : timeout;
+    return EMBEDDED_JOB_CLIENT_CTOR.newInstance(
+        jobId, dispatcherGateway, retryExecutor, rpcTimeout, userCodeClassloader);
   }
 }

externals/kyuubi-flink-sql-engine/src/main/scala/org/apache/kyuubi/engine/flink/FlinkEngineUtils.scala

@@ -47,7 +47,7 @@ object FlinkEngineUtils extends Logging {
   val EMBEDDED_MODE_CLIENT_OPTIONS: Options = getEmbeddedModeClientOptions(new Options)
 
   private def SUPPORTED_FLINK_VERSIONS =
-    Set("1.17", "1.18", "1.19", "1.20").map(SemanticVersion.apply)
+    Set("1.17", "1.18", "1.19", "1.20", "2.0", "2.1", "2.2", "2.3").map(SemanticVersion.apply)
 
   val FLINK_RUNTIME_VERSION: SemanticVersion = SemanticVersion(EnvironmentInformation.getVersion)
 
@@ -172,6 +172,17 @@ object FlinkEngineUtils extends Logging {
     } else null
   }
 
+  /**
+   * Copied from [[org.apache.flink.table.client.cli.CliOptionsParser]], which dropped the method
+   * in Flink 2.0.
+   */
+  private def checkFilePath(filePath: String): Unit = {
+    val scheme = new Path(filePath).toUri.getScheme
+    if (scheme != null && scheme != "file") {
+      throw new SqlClientException("SQL Client only supports to load files in local.")
+    }
+  }
+
   def renewDelegationToken(delegationToken: String): Unit = {
     val newCreds = KyuubiHadoopUtils.decodeCredentials(delegationToken)
     val newTokens = KyuubiHadoopUtils.getTokenMap(newCreds)

externals/kyuubi-flink-sql-engine/src/main/scala/org/apache/kyuubi/engine/flink/FlinkSQLEngine.scala

@@ -94,7 +94,7 @@ object FlinkSQLEngine extends Logging {
           .toMap
       flinkConf.addAll(Configuration.fromMap(flinkConfFromArgs.asJava))
 
-      val executionTarget = flinkConf.getString(DeploymentOptions.TARGET)
+      val executionTarget = flinkConf.get(DeploymentOptions.TARGET)
       setDeploymentConf(executionTarget, flinkConf)
 
       kyuubiConf.setIfMissing(KyuubiConf.FRONTEND_THRIFT_BINARY_BIND_PORT, 0)

externals/kyuubi-flink-sql-engine/src/test/scala/org/apache/kyuubi/engine/flink/WithFlinkSQLEngineLocal.scala

@@ -187,8 +187,8 @@ trait WithFlinkSQLEngineLocal extends KyuubiFunSuite with WithFlinkTestResources
       .build
     miniCluster = new MiniCluster(cfg)
     miniCluster.start()
-    flinkConfig.setString(RestOptions.ADDRESS, miniCluster.getRestAddress.get().getHost)
-    flinkConfig.setInteger(RestOptions.PORT, miniCluster.getRestAddress.get().getPort)
+    flinkConfig.set(RestOptions.ADDRESS, miniCluster.getRestAddress.get().getHost)
+    flinkConfig.set(RestOptions.PORT, Int.box(miniCluster.getRestAddress.get().getPort))
   }
 
   protected def getJdbcUrl: String = s"jdbc:hive2://$connectionUrl/;"

externals/kyuubi-flink-sql-engine/src/test/scala/org/apache/kyuubi/engine/flink/WithFlinkSQLEngineOnYarn.scala

@@ -33,6 +33,7 @@ import org.apache.hadoop.yarn.server.MiniYARNCluster
 import org.apache.kyuubi.{KYUUBI_VERSION, KyuubiFunSuite, SCALA_COMPILE_VERSION, Utils}
 import org.apache.kyuubi.config.KyuubiConf
 import org.apache.kyuubi.config.KyuubiConf.{ENGINE_FLINK_APPLICATION_JARS, KYUUBI_HOME_ENV_VAR_NAME}
+import org.apache.kyuubi.engine.flink.FlinkEngineUtils.FLINK_RUNTIME_VERSION
 import org.apache.kyuubi.ha.HighAvailabilityConf.HA_ADDRESSES
 import org.apache.kyuubi.util.JavaUtils
 import org.apache.kyuubi.util.command.CommandLineUtils._
@@ -167,7 +168,9 @@ trait WithFlinkSQLEngineOnYarn extends KyuubiFunSuite with WithFlinkTestResource
     val command = new ArrayBuffer[String]()
 
     command += s"${envs("FLINK_HOME")}${File.separator}bin/flink"
-    command += "run-application"
+    // Flink 2.0 merged `run-application` into `run` (FLINK-35625) and removed the former
+    // (FLINK-36310)
+    command += (if (FLINK_RUNTIME_VERSION < "2.0") "run-application" else "run")
     command += "-t"
     command += "yarn-application"
     command += s"-Dyarn.ship-files=${flinkExtraJars.mkString(";")}"