content: /*
   * Copyright (c) 2024, Alibaba Cloud;
   * Licensed under the Apache License, Version 2.0 (the "License");
   * you may not use this file except in compliance with the License.
   * You may obtain a copy of the License at
   *
   *    http://www.apache.org/licenses/LICENSE-2.0
   *
   * Unless required by applicable law or agreed to in writing, software
   * distributed under the License is distributed on an "AS IS" BASIS,
   * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
   * See the License for the specific language governing permissions and
   * limitations under the License.
   */
  
  package com.aliyun.dataworks.migrationx.domain.dataworks.dolphinscheduler.utils;
  
  import java.util.ArrayList;
  import java.util.List;
  
  import com.aliyun.dataworks.migrationx.domain.dataworks.dolphinscheduler.v1.v139.entity.ResourceInfo;
  import com.aliyun.dataworks.migrationx.domain.dataworks.dolphinscheduler.v1.v139.enums.ProgramType;
  import com.aliyun.dataworks.migrationx.domain.dataworks.dolphinscheduler.v1.v139.task.spark.SparkParameters;
  import com.aliyun.dataworks.migrationx.domain.dataworks.dolphinscheduler.v3.model.SparkConstants;
  import org.apache.commons.lang3.StringUtils;
  
  /**
   * spark args utils
   */
  public class SparkArgsUtils {
  
      private static final String SPARK_CLUSTER = "cluster";
  
      private static final String SPARK_LOCAL = "local";
  
      private static final String SPARK_ON_YARN = "yarn";
  
      /**
       * spark params constant
       */
      public static final String MASTER = "--master";
  
      public static final String DEPLOY_MODE = "--deploy-mode";
  
      /**
       * --class CLASS_NAME
       */
      public static final String MAIN_CLASS = "--class";
  
      /**
       * --driver-cores NUM
       */
      public static final String DRIVER_CORES = "--driver-cores";
  
      /**
       * --driver-memory MEM
       */
      public static final String DRIVER_MEMORY = "--driver-memory";
  
      /**
       * --num-executors NUM
       */
      public static final String NUM_EXECUTORS = "--num-executors";
  
      /**
       * --executor-cores NUM
       */
      public static final String EXECUTOR_CORES = "--executor-cores";
  
      /**
       * --executor-memory MEM
       */
      public static final String EXECUTOR_MEMORY = "--executor-memory";
  
      private static final String DEFAULT_QUALITY_CLASS =
          "org.apache.dolphinscheduler.data.quality.DataQualityApplication";
  
      /**
       * --name NAME
       */
      public static final String SPARK_NAME = "--name";
  
      /**
       * --queue QUEUE
       */
      public static final String SPARK_QUEUE = "--queue";
  
      private SparkArgsUtils() {
          throw new IllegalStateException("Utility class");
      }
  
      /**
       * build args
       *
       * @param param param
       * @return argument list
       */
      public static List<String> buildArgs(SparkParameters param) {
          List<String> args = new ArrayList<>();
          args.add(MASTER);
  
          String deployMode = StringUtils.isNotEmpty(param.getDeployMode()) ? param.getDeployMode() : SPARK_CLUSTER;
          if (!SPARK_LOCAL.equals(deployMode)) {
              args.add(SPARK_ON_YARN);
              args.add(DEPLOY_MODE);
          }
          args.add(deployMode);
  
          ProgramType programType = param.getProgramType();
          String mainClass = param.getMainClass();
          if (programType != null && programType != ProgramType.PYTHON && StringUtils.isNotEmpty(mainClass)) {
              args.add(MAIN_CLASS);
              args.add(mainClass);
          }
  
          int driverCores = param.getDriverCores();
          if (driverCores > 0) {
              args.add(DRIVER_CORES);
              args.add(String.format("%d", driverCores));
          }
  
          String driverMemory = param.getDriverMemory();
          if (StringUtils.isNotEmpty(driverMemory)) {
              args.add(DRIVER_MEMORY);
              args.add(driverMemory);
          }
  
          int numExecutors = param.getNumExecutors();
          if (numExecutors > 0) {
              args.add(NUM_EXECUTORS);
              args.add(String.format("%d", numExecutors));
          }
  
          int executorCores = param.getExecutorCores();
          if (executorCores > 0) {
              args.add(EXECUTOR_CORES);
              args.add(String.format("%d", executorCores));
          }
  
          String executorMemory = param.getExecutorMemory();
          if (StringUtils.isNotEmpty(executorMemory)) {
              args.add(EXECUTOR_MEMORY);
              args.add(executorMemory);
          }
  
          String appName = param.getAppName();
          if (StringUtils.isNotEmpty(appName)) {
              args.add(SPARK_NAME);
              args.add(ArgsUtils.escape(appName));
          }
  
          String others = param.getOthers();
          if (!SPARK_LOCAL.equals(deployMode)) {
              if (StringUtils.isEmpty(others) || !others.contains(SPARK_QUEUE)) {
                  String queue = param.getQueue();
                  if (StringUtils.isNotEmpty(queue)) {
                      args.add(SPARK_QUEUE);
                      args.add(queue);
                  }
              }
          }
  
          // --conf --files --jars --packages
          if (StringUtils.isNotEmpty(others)) {
              args.add(others);
          }
  
          ResourceInfo mainJar = param.getMainJar();
          if (mainJar != null) {
              args.add(StringUtils.defaultString(mainJar.getRes(), mainJar.getResourceName()));
          }
  
          String mainArgs = param.getMainArgs();
          if (StringUtils.isNotEmpty(mainArgs)) {
              args.add(mainArgs);
          }
  
          return args;
      }
  
      public static List<String> buildArgs(com.aliyun.dataworks.migrationx.domain.dataworks.dolphinscheduler.v3.task.spark.SparkParameters param) {
          List<String> args = new ArrayList<>();
          args.add(SparkConstants.MASTER);
  
          String deployMode = StringUtils.isNotEmpty(param.getDeployMode()) ? param.getDeployMode() : SPARK_CLUSTER;
          if (!SPARK_LOCAL.equals(deployMode)) {
              args.add(SPARK_ON_YARN);
              args.add(SparkConstants.DEPLOY_MODE);
          }
          args.add(deployMode);
  
          com.aliyun.dataworks.migrationx.domain.dataworks.dolphinscheduler.v3.model.ProgramType programType = param.getProgramType();
          String mainClass = param.getMainClass();
          if (programType != null && programType != com.aliyun.dataworks.migrationx.domain.dataworks.dolphinscheduler.v3.model.ProgramType.PYTHON) {
              args.add(SparkConstants.MAIN_CLASS);
              args.add(StringUtils.isNotEmpty(mainClass) ? mainClass : DEFAULT_QUALITY_CLASS);
          }
  
          int driverCores = param.getDriverCores();
          if (driverCores > 0) {
              args.add(SparkConstants.DRIVER_CORES);
              args.add(String.format("%d", driverCores));
          }
  
          String driverMemory = param.getDriverMemory();
          if (StringUtils.isNotEmpty(driverMemory)) {
              args.add(SparkConstants.DRIVER_MEMORY);
              args.add(driverMemory);
          }
  
          int numExecutors = param.getNumExecutors();
          if (numExecutors > 0) {
              args.add(SparkConstants.NUM_EXECUTORS);
              args.add(String.format("%d", numExecutors));
          }
  
          int executorCores = param.getExecutorCores();
          if (executorCores > 0) {
              args.add(SparkConstants.EXECUTOR_CORES);
              args.add(String.format("%d", executorCores));
          }
  
          String executorMemory = param.getExecutorMemory();
          if (StringUtils.isNotEmpty(executorMemory)) {
              args.add(SparkConstants.EXECUTOR_MEMORY);
              args.add(executorMemory);
          }
  
          String appName = param.getAppName();
          if (StringUtils.isNotEmpty(appName)) {
              args.add(SparkConstants.SPARK_NAME);
              args.add(ArgsUtils.escape(appName));
          }
  
          String others = param.getOthers();
          if (!SPARK_LOCAL.equals(deployMode)
              && (StringUtils.isEmpty(others) || !others.contains(SparkConstants.SPARK_QUEUE))) {
              String queue = param.getYarnQueue();
              if (StringUtils.isNotEmpty(queue)) {
                  args.add(SparkConstants.SPARK_QUEUE);
                  args.add(queue);
              }
          }
  
          // --conf --files --jars --packages
          if (StringUtils.isNotEmpty(others)) {
              args.add(others);
          }
  
          com.aliyun.dataworks.migrationx.domain.dataworks.dolphinscheduler.v3.model.ResourceInfo mainJar = param.getMainJar();
          if (mainJar != null) {
              args.add(mainJar.getResourceName());
          }
  
          String mainArgs = param.getMainArgs();
          if (StringUtils.isNotEmpty(mainArgs)) {
              args.add(mainArgs);
          }
  
          return args;
      }
  }