Repository: flink-extended/dl-on-flink (language: python)

Issue:
[Flink AI Flow][Plugins] Flink job plugin supports SQL run mode
### Describe the feature

Currently Flink job plugin supports Python and Java run mode, but doesn't support SQL run mode. The SQL run mode should be supported for the users which only want to define the job through the SQL.

### Describe the solution you'd like

Introduce the `FlinkSqlProcessor`, which can implement the interface `FlinkPythonProcessor`.

### Describe alternatives you've considered

### Additional context



Discussion / hints:

Fix patch:
diff --git a/flink-ai-flow/ai_flow/scheduler_service/service/workflow_event_processor.py b/flink-ai-flow/ai_flow/scheduler_service/service/workflow_event_processor.py
index 5503d88dc..5b32fce5c 100644
--- a/flink-ai-flow/ai_flow/scheduler_service/service/workflow_event_processor.py
+++ b/flink-ai-flow/ai_flow/scheduler_service/service/workflow_event_processor.py
@@ -70,6 +70,8 @@ def run(self) -> None:
     def _process_event(self, event: BaseEvent):
         subscribed_workflow = []
         projects = self.store.list_project(sys.maxsize, 0)
+        if projects is None:
+            return
         for project in projects:
             project_name = project.name
             workflows = self._get_subscribed_workflow(event, project_name)
diff --git a/flink-ai-flow/ai_flow_plugins/job_plugins/flink/__init__.py b/flink-ai-flow/ai_flow_plugins/job_plugins/flink/__init__.py
index 7370d907d..4fd4fa115 100644
--- a/flink-ai-flow/ai_flow_plugins/job_plugins/flink/__init__.py
+++ b/flink-ai-flow/ai_flow_plugins/job_plugins/flink/__init__.py
@@ -16,7 +16,8 @@
 # under the License.
 from ai_flow.plugin_interface import register_job_plugin_factory
 from ai_flow_plugins.job_plugins.flink.flink_job_plugin import FlinkJobPluginFactory
-from ai_flow_plugins.job_plugins.flink.flink_processor import FlinkPythonProcessor, FlinkJavaProcessor, ExecutionContext
+from ai_flow_plugins.job_plugins.flink.flink_processor import FlinkPythonProcessor, FlinkJavaProcessor, \
+    ExecutionContext, FlinkSqlProcessor, UDFWrapper
 from ai_flow_plugins.job_plugins.flink.flink_job_config import FlinkJobConfig
 from ai_flow_plugins.job_plugins.flink.flink_env import set_flink_env, FlinkEnv, \
     FlinkBatchEnv, FlinkStreamEnv, WrappedBatchTableEnvironment, WrappedStreamTableEnvironment
diff --git a/flink-ai-flow/ai_flow_plugins/job_plugins/flink/flink_processor.py b/flink-ai-flow/ai_flow_plugins/job_plugins/flink/flink_processor.py
index 3e3188212..f741e623c 100644
--- a/flink-ai-flow/ai_flow_plugins/job_plugins/flink/flink_processor.py
+++ b/flink-ai-flow/ai_flow_plugins/job_plugins/flink/flink_processor.py
@@ -15,7 +15,11 @@
 # specific language governing permissions and limitations
 # under the License.
 from abc import abstractmethod
-from typing import List, Dict, Text, Optional
+from typing import List, Dict, Text, Optional, Union
+
+from pyflink.table.udf import UserDefinedScalarFunctionWrapper
+
+from ai_flow import DatasetMeta
 from ai_flow.plugin_interface.scheduler_interface import JobExecutionInfo
 from ai_flow.util import json_utils
 from pyflink.dataset import ExecutionEnvironment
@@ -28,6 +32,7 @@ class ExecutionContext(json_utils.Jsonable):
     AINode node configuration parameters(config), ExecutionEnvironment, TableEnvironment and StatementSet.
     ExecutionContext is passed as a parameter to the process function of FlinkPythonProcessor.
     """
+
     def __init__(self,
                  job_execution_info: JobExecutionInfo,
                  config: Dict,
@@ -101,6 +106,7 @@ class FlinkJavaProcessor(object):
     """
     FlinkJavaProcessor is the processor of flink java jobs.
     """
+
     def __init__(self,
                  entry_class: Optional[Text],
                  main_jar_file: Text,
@@ -114,3 +120,113 @@ def __init__(self,
         self.entry_class: Text = entry_class
         self.main_jar_file: Text = main_jar_file
         self.args = args
+
+
+class UDFWrapper(object):
+    """
+    A wrapper for flink udf, which will be utilized by FlinkSqlProcessor to pass user-defined functions(udf and udtf).
+    """
+    def __init__(self, name, func: Union[str, UserDefinedScalarFunctionWrapper]):
+        """
+       :param name: The name of user-defined functions which will be registered in the table env.
+       :param func: Thr user-defined function. For python, it is the python user-defined function to register. For java,
+        it is the java full qualified class name of the function to register.
+       """
+        self.name = name
+        self.func = func
+
+    def is_java_udf(self):
+        return type(self.func) == str
+
+
+class FlinkSqlProcessor(FlinkPythonProcessor):
+    """
+    FlinkSqlProcessor is the processor of flink sql jobs based on pyflink.
+    """
+    @abstractmethod
+    def sql_statements(self, execution_context: ExecutionContext) -> List[str]:
+        """
+        The user should override this method to define their own sql statements. Multiple insertions in a job will be
+        added into one same statement set.
+        
+        For the processor in :py:func:`~ai_flow.api.ops.read_dataset` or :py:func:`~ai_flow.api.ops.write_dataset`, 
+        this method should return at most one DDL(i.e. CREATE statement) for table source or table sink and the DDL
+        should be consistent with the dataset_info in the :py:func:`~ai_flow.api.ops.read_dataset` or 
+        :py:func:`~ai_flow.api.ops.write_dataset` (i.e. dataset_info.data_format should be equal to the 'format' option 
+        in DDL if it is defined.). 
+        
+        If multiple table sinks are required, users can call multiple :py:func:`~ai_flow.api.ops.write_dataset` or use 
+        single :py:func:`~ai_flow.api.ops.user_define_operation` with multiple DDL. 
+        
+
+        :param execution_context: The :class:`~ai_flow_plugins.job_plugins.flink.flink_processor.ExecutionContext` of
+        the processor
+        :rtype List[str]: A list of sql statement. Each element should be a DDL/DML/DQL statement.
+        """
+        pass
+
+    @abstractmethod
+    def udf_list(self, execution_context: ExecutionContext) -> List[UDFWrapper]:
+        """
+        The user should override this method to register user-defined functions(udf or udtf). 
+
+        :param execution_context: The :class:`~ai_flow_plugins.job_plugins.flink.flink_processor.ExecutionContext` of
+        the processor
+        :rtype List[UDFWrapper]: A list of :classs:`ai_flow_plugins.job_plugins.flink.flink_processor.UDFWrapper`. 
+        """
+        pass
+
+    def process(self, execution_context: ExecutionContext, input_list: List[Table] = None) -> List[Table]:
+        """
+        The method will register udfs defined in :py:func:`~ai_flow_plugins.job_plugins.flink.flink_processor.FlinkSqlProcessor.udf_list`
+        and execute sql statements in :py:func:`~ai_flow_plugins.job_plugins.flink.flink_processor.FlinkSqlProcessor.sql_statements`.
+        """
+        _sql_statements = self.sql_statements(execution_context)
+        if _sql_statements is None or len(_sql_statements) == 0:
+            raise Exception("The sql_statements() cannot be None or empty.")
+        sql_statements = [statement.strip() for statement in _sql_statements]
+        table_env = execution_context.table_env
+        statement_set = execution_context.statement_set
+        _udfs = self.udf_list(execution_context)
+        if _udfs is not None:
+            for udf in _udfs:
+                if udf.is_java_udf():
+                    table_env.register_java_function(udf.name, udf.func)
+                else:
+                    table_env.register_function(udf.name, udf.func)
+
+        a = execution_context.node_type
+        for sql_statement in sql_statements:
+            if sql_statement.lower().startswith('insert'):
+                statement_set.add_insert_sql(sql_statement)
+            if sql_statement.lower().startswith('create'):
+                # Check if users' DDL are consistent with info they provide in the
+                # :py:class:`ai_flow.meta.dataset_meta.DatasetMeta`
+                if execution_context.node_type == 'read_dataset' or execution_context.node_type == 'write_dataset':
+                    dataset_meta: DatasetMeta = execution_context.config['dataset']
+                    if not _validate_create_statement(sql_statement, dataset_meta):
+                        raise Exception("'format' option in CREATE statement is inconsistent with the attached "
+                                        "dataset! The registered format of the dataset is {}", dataset_meta.data_format)
+
+                table_env.execute_sql(sql_statement)
+            else:
+                table_env.execute_sql(sql_statement)
+        return []
+
+
+def _check_options_in_sql(stmt, option_name, target_value):
+    if target_value is None:
+        return True
+    words = stmt.split()
+    for i in range(len(words)):
+        if words[i] == option_name and i + 1 < len(words) and words[i + 1] == '=':
+            if i + 2 < len(words) and target_value.lower() != words[i + 2].replace("'", '').lower():
+                return False
+    return True
+
+
+def _validate_create_statement(stmt, dataset_meta: DatasetMeta):
+    required_format = dataset_meta.data_format
+    required_connector = None if dataset_meta.properties is None else dataset_meta.properties.get('connector')
+    return _check_options_in_sql(stmt, "'format'", required_format) and \
+           _check_options_in_sql(stmt, "'connector'", required_connector)
diff --git a/flink-ai-flow/ai_flow_plugins/job_plugins/flink/flink_wrapped_env.py b/flink-ai-flow/ai_flow_plugins/job_plugins/flink/flink_wrapped_env.py
index 6d1a4e353..d9e5306f6 100644
--- a/flink-ai-flow/ai_flow_plugins/job_plugins/flink/flink_wrapped_env.py
+++ b/flink-ai-flow/ai_flow_plugins/job_plugins/flink/flink_wrapped_env.py
@@ -28,10 +28,11 @@ class WrappedTableEnvironmentContext:
     - Lists the ids of the Flink jobs.
     - Waits for the execution result of the Flink job.
     """
+
     def __init__(self):
-        self.execute_sql_results:List[TableResult] = []
-    
-    def add_execute_sql_result(self, result:TableResult):
+        self.execute_sql_results: List[TableResult] = []
+
+    def add_execute_sql_result(self, result: TableResult):
         self.execute_sql_results.append(result)
 
     def get_job_ids(self) -> List[Text]:
@@ -71,6 +72,7 @@ def create_statement_set(self) -> 'WrappedStatementSet':
     def wait_execution_results(self):
         self.wrapped_context.wait_execution_results()
 
+
 class WrappedBatchTableEnvironment(BatchTableEnvironment, WrappedTableEnvironment):
 
     def __init__(self, j_tenv):
@@ -100,6 +102,7 @@ class WrappedStatementSetContext:
     - Waits for the execution result of the Flink job.
     - Decides whether 'StatementSet' needs to call 'execute()' method based on whether it contains 'execute_sql'.
     """
+
     def __init__(self):
         self.execute_results: List[TableResult] = []
         self.need_execute = False
@@ -137,16 +140,16 @@ def execute(self) -> TableResult:
         self.wrapped_context.add_execute_result(result)
         self.wrapped_context.need_execute = False
         return result
-    
+
     def add_insert_sql(self, stmt: str) -> 'StatementSet':
         result = super().add_insert_sql(stmt)
         self.wrapped_context.need_execute = True
         return result
 
-    def add_insert(self, target_path, table, overwrite:bool=False) -> 'StatementSet':
+    def add_insert(self, target_path, table, overwrite: bool = False) -> 'StatementSet':
         result = super().add_insert(target_path, table, overwrite=overwrite)
         self.wrapped_context.need_execute = True
         return result
 
     def wait_execution_results(self):
-        self.wrapped_context.wait_execution_results()
\ No newline at end of file
+        self.wrapped_context.wait_execution_results()


Test patch:
diff --git a/flink-ai-flow/ai_flow_plugins/tests/job_plugins/ut_workflows/workflows/test_flink/test_flink.py b/flink-ai-flow/ai_flow_plugins/tests/job_plugins/ut_workflows/workflows/test_flink/test_flink.py
index 800f98ba5..30f086628 100644
--- a/flink-ai-flow/ai_flow_plugins/tests/job_plugins/ut_workflows/workflows/test_flink/test_flink.py
+++ b/flink-ai-flow/ai_flow_plugins/tests/job_plugins/ut_workflows/workflows/test_flink/test_flink.py
@@ -21,7 +21,8 @@
 from ai_flow import AIFlowServerRunner, init_ai_flow_context
 from ai_flow.workflow.status import Status
 from ai_flow_plugins.job_plugins import flink
-from test_flink_processor import SinkWithAddInsertSql, SinkWithExecuteSql, Source, Sink, Transformer, Transformer2
+from test_flink_processor import SinkWithAddInsertSql, SinkWithExecuteSql, Source, Sink, Transformer, Transformer2, \
+    SourceSql, SinkSql
 import ai_flow as af
 
 project_path = os.path.dirname(os.path.dirname(os.path.dirname(__file__)))
@@ -64,6 +65,17 @@ def test_local_flink_task(self):
         je = af.workflow_operation.get_job_execution(job_name='task_1', execution_id='1')
         self.assertEqual(Status.FINISHED, je.status)
 
+    def test_local_flink_task_with_flink_sql_processor(self):
+        flink.set_flink_env(flink.FlinkStreamEnv())
+        with af.job_config('task_1'):
+            af.user_define_operation(processor=SourceSql())
+            af.user_define_operation(processor=SinkSql())
+
+        workflow = af.workflow_operation.submit_workflow(workflow_name=af.current_workflow_config().workflow_name)
+        job_execution = af.workflow_operation.start_job_execution(job_name='task_1', execution_id='1')
+        job_execution = af.workflow_operation.get_job_execution(job_name='task_1', execution_id='1')
+        self.assertEqual(Status.FINISHED, job_execution.status)
+
     def test_stop_local_flink_task(self):
         with af.job_config('task_1'):
             input_example = af.user_define_operation(processor=Source())
diff --git a/flink-ai-flow/ai_flow_plugins/tests/job_plugins/ut_workflows/workflows/test_flink/test_flink_processor.py b/flink-ai-flow/ai_flow_plugins/tests/job_plugins/ut_workflows/workflows/test_flink/test_flink_processor.py
index 7f3d111b9..d5ce1c1c3 100644
--- a/flink-ai-flow/ai_flow_plugins/tests/job_plugins/ut_workflows/workflows/test_flink/test_flink_processor.py
+++ b/flink-ai-flow/ai_flow_plugins/tests/job_plugins/ut_workflows/workflows/test_flink/test_flink_processor.py
@@ -15,9 +15,11 @@
 # specific language governing permissions and limitations
 # under the License.
 import os
+import shutil
 import time
 from typing import List
 
+from ai_flow_plugins.job_plugins.flink.flink_processor import UDFWrapper
 from pyflink.table import Table, DataTypes, ScalarFunction
 from pyflink.table.descriptors import Schema, OldCsv, FileSystem
 from pyflink.table.udf import udf
@@ -87,6 +89,7 @@ def process(self, execution_context: flink.ExecutionContext, input_list: List[Ta
         statement_set.add_insert('mySink', input_list[0])
         return []
 
+
 class SinkWithExecuteSql(flink.FlinkPythonProcessor):
     def process(self,
                 execution_context: flink.ExecutionContext,
@@ -137,4 +140,59 @@ def process(self,
         SELECT *
         FROM {input_list[0]}
         """)
-        return []
\ No newline at end of file
+        return []
+
+
+class SourceSql(flink.FlinkSqlProcessor):
+    def __init__(self):
+        super().__init__()
+        self.input_file = None
+
+    def open(self, execution_context: flink.ExecutionContext):
+        self.input_file = os.path.join(os.getcwd(), 'resources', 'word_count.txt')
+
+    def sql_statements(self, execution_context: flink.ExecutionContext) -> List[str]:
+        table_source_stmt = '''
+                           CREATE TABLE mySource (
+                               word STRING 
+                           ) WITH (
+                               'connector' = 'filesystem',
+                               'path' = '{uri}',
+                               'format' = 'csv',
+                               'csv.ignore-parse-errors' = 'true'
+                           )
+                            '''.format(uri=self.input_file)
+        return [table_source_stmt]
+
+
+class SinkSql(flink.FlinkSqlProcessor):
+    def __init__(self):
+        super().__init__()
+        self.output_file = None
+
+    def open(self, execution_context: flink.ExecutionContext):
+        self.output_file = os.path.join(os.getcwd(), 'output')
+
+        if os.path.exists(self.output_file):
+            os.remove(self.output_file)
+        if os.path.isdir(self.output_file):
+            shutil.rmtree(self.output_file)
+
+    def udf_list(self, execution_context: flink.ExecutionContext) -> List[UDFWrapper]:
+        return [UDFWrapper(name='pass_func', func=udf(PassUDF(),
+                                                      input_types=[DataTypes.STRING()],
+                                                      result_type=DataTypes.STRING()))]
+
+    def sql_statements(self, execution_context: flink.ExecutionContext) -> List[str]:
+        table_sink_stmt = '''
+                           CREATE TABLE mySink (
+                               newword STRING
+                           ) WITH (
+                               'connector' = 'filesystem',
+                               'path' = '{uri}',
+                               'format' = 'csv',
+                               'csv.ignore-parse-errors' = 'true'
+                           )
+        '''.format(uri=self.output_file)
+        insert_stmt = 'INSERT INTO mySink SELECT pass_func(word) FROM mySource'
+        return [table_sink_stmt, insert_stmt]