Ver Fonte

[Improvement-14371][Hivecli] Add support for hivevar setting in hivecli task (#14423)

Rick Cheng há 1 ano atrás
pai
commit
13dfa859b5

+ 63 - 8
dolphinscheduler-task-plugin/dolphinscheduler-task-hivecli/src/main/java/org/apache/dolphinscheduler/plugin/task/hivecli/HiveCliTask.java

@@ -18,8 +18,10 @@
 package org.apache.dolphinscheduler.plugin.task.hivecli;
 
 import static org.apache.dolphinscheduler.plugin.task.api.TaskConstants.EXIT_CODE_FAILURE;
+import static org.apache.dolphinscheduler.plugin.task.api.TaskConstants.RWXR_XR_X;
 
 import org.apache.dolphinscheduler.common.utils.JSONUtils;
+import org.apache.dolphinscheduler.common.utils.OSUtils;
 import org.apache.dolphinscheduler.plugin.task.api.AbstractRemoteTask;
 import org.apache.dolphinscheduler.plugin.task.api.ShellCommandExecutor;
 import org.apache.dolphinscheduler.plugin.task.api.TaskCallBack;
@@ -31,12 +33,23 @@ import org.apache.dolphinscheduler.plugin.task.api.model.TaskResponse;
 import org.apache.dolphinscheduler.plugin.task.api.parameters.AbstractParameters;
 import org.apache.dolphinscheduler.plugin.task.api.utils.ParameterUtils;
 
+import org.apache.commons.io.FileUtils;
 import org.apache.commons.lang3.StringUtils;
 
+import java.io.File;
+import java.io.IOException;
+import java.nio.charset.StandardCharsets;
+import java.nio.file.Files;
+import java.nio.file.Path;
+import java.nio.file.StandardOpenOption;
+import java.nio.file.attribute.FileAttribute;
+import java.nio.file.attribute.PosixFilePermission;
+import java.nio.file.attribute.PosixFilePermissions;
 import java.util.ArrayList;
 import java.util.Collections;
 import java.util.List;
 import java.util.Map;
+import java.util.Set;
 
 public class HiveCliTask extends AbstractRemoteTask {
 
@@ -108,29 +121,41 @@ public class HiveCliTask extends AbstractRemoteTask {
 
         final String type = hiveCliParameters.getHiveCliTaskExecutionType();
 
+        String sqlContent = "";
+        String resourceFileName = "";
         // TODO: make sure type is not unknown
         if (HiveCliConstants.TYPE_FILE.equals(type)) {
-            args.add(HiveCliConstants.HIVE_CLI_EXECUTE_FILE);
             final List<ResourceInfo> resourceInfos = hiveCliParameters.getResourceList();
             if (resourceInfos.size() > 1) {
                 log.warn("more than 1 files detected, use the first one by default");
             }
 
-            args.add(StringUtils.stripStart(resourceInfos.get(0).getResourceName(), "/"));
+            try {
+                resourceFileName = resourceInfos.get(0).getResourceName();
+                sqlContent = FileUtils.readFileToString(
+                        new File(String.format("%s/%s", taskExecutionContext.getExecutePath(), resourceFileName)),
+                        StandardCharsets.UTF_8);
+            } catch (IOException e) {
+                log.error("read hive sql content from file {} error ", resourceFileName, e);
+                throw new TaskException("read hive sql content error", e);
+            }
         } else {
-            final String script = hiveCliParameters.getHiveSqlScript();
-            args.add(String.format(HiveCliConstants.HIVE_CLI_EXECUTE_SCRIPT, script));
+            sqlContent = hiveCliParameters.getHiveSqlScript();
         }
 
+        final Map<String, Property> paramsMap = taskExecutionContext.getPrepareParamsMap();
+        sqlContent = ParameterUtils.convertParameterPlaceholders(sqlContent, ParameterUtils.convert(paramsMap));
+        log.info("HiveCli sql content: {}", sqlContent);
+        String sqlFilePath = generateSqlScriptFile(sqlContent);
+
+        args.add(HiveCliConstants.HIVE_CLI_EXECUTE_FILE);
+        args.add(new File(sqlFilePath).getName());
         final String hiveCliOptions = hiveCliParameters.getHiveCliOptions();
         if (StringUtils.isNotEmpty(hiveCliOptions)) {
             args.add(hiveCliOptions);
         }
 
-        final Map<String, Property> paramsMap = taskExecutionContext.getPrepareParamsMap();
-        final String command =
-                ParameterUtils.convertParameterPlaceholders(String.join(" ", args), ParameterUtils.convert(paramsMap));
-
+        String command = String.join(" ", args);
         log.info("hiveCli task command: {}", command);
 
         return command;
@@ -151,4 +176,34 @@ public class HiveCliTask extends AbstractRemoteTask {
         }
     }
 
+    protected String generateSqlScriptFile(String rawScript) {
+        String scriptFileName = String.format("%s/%s_node.sql", taskExecutionContext.getExecutePath(),
+                taskExecutionContext.getTaskAppId());
+
+        File file = new File(scriptFileName);
+        Path path = file.toPath();
+
+        if (!Files.exists(path)) {
+            String script = rawScript.replaceAll("\\r\\n", "\n");
+
+            Set<PosixFilePermission> perms = PosixFilePermissions.fromString(RWXR_XR_X);
+            FileAttribute<Set<PosixFilePermission>> attr = PosixFilePermissions.asFileAttribute(perms);
+            try {
+                if (OSUtils.isWindows()) {
+                    Files.createFile(path);
+                } else {
+                    if (!file.getParentFile().exists()) {
+                        file.getParentFile().mkdirs();
+                    }
+                    Files.createFile(path, attr);
+                }
+                Files.write(path, script.getBytes(), StandardOpenOption.APPEND);
+            } catch (IOException e) {
+                log.error("generate hivecli sql script error", e);
+                throw new TaskException("generate hivecli sql script error", e);
+            }
+        }
+        return scriptFileName;
+    }
+
 }

+ 22 - 3
dolphinscheduler-task-plugin/dolphinscheduler-task-hivecli/src/test/java/org/apache/dolphinscheduler/plugin/task/hivecli/HiveCliTaskTest.java

@@ -17,6 +17,7 @@
 
 package org.apache.dolphinscheduler.plugin.task.hivecli;
 
+import static org.mockito.Mockito.doReturn;
 import static org.mockito.Mockito.spy;
 import static org.mockito.Mockito.when;
 
@@ -24,12 +25,17 @@ import org.apache.dolphinscheduler.common.utils.JSONUtils;
 import org.apache.dolphinscheduler.plugin.task.api.TaskExecutionContext;
 import org.apache.dolphinscheduler.plugin.task.api.model.ResourceInfo;
 
+import org.apache.commons.io.FileUtils;
+
 import java.util.ArrayList;
 import java.util.List;
 
+import org.junit.jupiter.api.AfterEach;
 import org.junit.jupiter.api.Assertions;
+import org.junit.jupiter.api.BeforeEach;
 import org.junit.jupiter.api.Test;
 import org.junit.jupiter.api.extension.ExtendWith;
+import org.mockito.MockedStatic;
 import org.mockito.Mockito;
 import org.mockito.junit.jupiter.MockitoExtension;
 
@@ -37,13 +43,25 @@ import org.mockito.junit.jupiter.MockitoExtension;
 public class HiveCliTaskTest {
 
     public static final String EXPECTED_HIVE_CLI_TASK_EXECUTE_FROM_SCRIPT_COMMAND =
-            "hive -e \"SHOW DATABASES;\"";
+            "hive -f 123_node.sql";
 
     public static final String EXPECTED_HIVE_CLI_TASK_EXECUTE_FROM_FILE_COMMAND =
-            "hive -f sql_tasks/hive_task.sql";
+            "hive -f 123_node.sql";
 
     public static final String EXPECTED_HIVE_CLI_TASK_EXECUTE_WITH_OPTIONS =
-            "hive -e \"SHOW DATABASES;\" --verbose";
+            "hive -f 123_node.sql --verbose";
+
+    private MockedStatic<FileUtils> mockedStaticFileUtils;
+
+    @BeforeEach
+    public void setUp() {
+        mockedStaticFileUtils = Mockito.mockStatic(FileUtils.class);
+    }
+
+    @AfterEach
+    public void after() {
+        mockedStaticFileUtils.close();
+    }
 
     @Test
     public void hiveCliTaskExecuteSqlFromScript() throws Exception {
@@ -73,6 +91,7 @@ public class HiveCliTaskTest {
         TaskExecutionContext taskExecutionContext = Mockito.mock(TaskExecutionContext.class);
         when(taskExecutionContext.getTaskParams()).thenReturn(hiveCliTaskParameters);
         HiveCliTask hiveCliTask = spy(new HiveCliTask(taskExecutionContext));
+        doReturn("123_node.sql").when(hiveCliTask).generateSqlScriptFile(Mockito.any());
         return hiveCliTask;
     }