change process_definition usage of Subprocess task

devosend · devosend · commit e54457f120d8 · 2021-12-02T22:27:49.000+08:00
diff --git a/dolphinscheduler-python/pydolphinscheduler/src/pydolphinscheduler/exceptions.py b/dolphinscheduler-python/pydolphinscheduler/src/pydolphinscheduler/exceptions.py
@@ -40,3 +40,7 @@ class PyDSJavaGatewayException(PyDSBaseException):
     """Exception for pydolphinscheduler Java gateway error."""
 
     pass
+
+
+class PyDSProcessDefinitionNotAssignException(PyDSBaseException):
+    """Exception for pydolphinscheduler process definition not assign error."""
diff --git a/dolphinscheduler-python/pydolphinscheduler/src/pydolphinscheduler/tasks/sub_process.py b/dolphinscheduler-python/pydolphinscheduler/src/pydolphinscheduler/tasks/sub_process.py
@@ -22,6 +22,7 @@
 from pydolphinscheduler.constants import TaskType
 from pydolphinscheduler.core.process_definition import ProcessDefinitionContext
 from pydolphinscheduler.core.task import Task, TaskParams
+from pydolphinscheduler.exceptions import PyDSProcessDefinitionNotAssignException
 from pydolphinscheduler.java_gateway import launch_gateway
 
 
@@ -38,7 +39,16 @@ class SubProcess(Task):
 
     def __init__(self, name: str, process_definition_name: str, *args, **kwargs):
         self._process_definition_name = process_definition_name
-        self._process_definition = {}
+        self._process_definition_info = {}
+        # TODO: Optimize the way of obtaining process_definition
+        self.process_definition = kwargs.get(
+            "process_definition", ProcessDefinitionContext.get()
+        )
+        if not self.process_definition:
+            raise PyDSProcessDefinitionNotAssignException(
+                "ProcessDefinition must be provider when SubProcess initialization."
+            )
+
         task_params = SubProcessTaskParams(
             process_definition_code=self.get_process_definition_code(),
         )
@@ -52,14 +62,15 @@ def get_process_definition_code(self) -> str:
 
     def get_process_definition_info(self, process_definition_name: str) -> Dict:
         """Get process definition info from java gateway, contains process definition id, name, code."""
-        if self._process_definition:
-            return self._process_definition
+        if self._process_definition_info:
+            return self._process_definition_info
         else:
             gateway = launch_gateway()
-            process_definition = ProcessDefinitionContext.get()
-            self._process_definition = gateway.entry_point.getProcessDefinitionInfo(
-                process_definition.user.name,
-                process_definition.project.name,
-                process_definition_name,
+            self._process_definition_info = (
+                gateway.entry_point.getProcessDefinitionInfo(
+                    self.process_definition.user.name,
+                    self.process_definition.project.name,
+                    process_definition_name,
+                )
             )
-            return self._process_definition
+            return self._process_definition_info
diff --git a/dolphinscheduler-python/pydolphinscheduler/tests/tasks/test_sub_process.py b/dolphinscheduler-python/pydolphinscheduler/tests/tasks/test_sub_process.py
@@ -22,10 +22,12 @@
 
 import pytest
 
+from pydolphinscheduler.core.process_definition import ProcessDefinition
 from pydolphinscheduler.tasks.sub_process import SubProcess, SubProcessTaskParams
 
-PROCESS_DEFINITION_NAME = "test-process-definition"
-PROCESS_DEFINITION_CODE = "3643589832320"
+TEST_SUB_PROCESS_DEFINITION_NAME = "sub-test-process-definition"
+TEST_SUB_PROCESS_DEFINITION_CODE = "3643589832320"
+TEST_PROCESS_DEFINITION_NAME = "simple-test-process-definition"
 
 
 @pytest.mark.parametrize(
@@ -50,7 +52,11 @@ def test_sub_process_task_params_attr_setter(name, value):
 @patch(
     "pydolphinscheduler.tasks.sub_process.SubProcess.get_process_definition_info",
     return_value=(
-        {"id": 1, "name": PROCESS_DEFINITION_NAME, "code": PROCESS_DEFINITION_CODE}
+        {
+            "id": 1,
+            "name": TEST_SUB_PROCESS_DEFINITION_NAME,
+            "code": TEST_SUB_PROCESS_DEFINITION_CODE,
+        }
     ),
 )
 def test_sub_process_to_dict(mock_process_definition):
@@ -68,7 +74,7 @@ def test_sub_process_to_dict(mock_process_definition):
         "taskParams": {
             "resourceList": [],
             "localParams": [],
-            "processDefinitionCode": PROCESS_DEFINITION_CODE,
+            "processDefinitionCode": TEST_SUB_PROCESS_DEFINITION_CODE,
             "dependence": {},
             "conditionResult": {"successNode": [""], "failedNode": [""]},
             "waitStartTimeout": {},
@@ -86,5 +92,6 @@ def test_sub_process_to_dict(mock_process_definition):
         "pydolphinscheduler.core.task.Task.gen_code_and_version",
         return_value=(code, version),
     ):
-        sub_process = SubProcess(name, PROCESS_DEFINITION_NAME)
-        assert sub_process.to_dict() == expect
+        with ProcessDefinition(TEST_PROCESS_DEFINITION_NAME):
+            sub_process = SubProcess(name, TEST_SUB_PROCESS_DEFINITION_NAME)
+            assert sub_process.to_dict() == expect
