[xy] Add AzureEventHub source to streaming pipeline (mage-ai#1260)

* [xy] Add AzureEventHub source.

* [xy] Add unit tests to streaming module.

* [xy] Add unit test for AzureEventHub.

* [xy] Update event handling logic.

Author: wangxiaoyou1993

mage_ai/data_preparation/executors/streaming_pipeline_executor.py

@@ -1,9 +1,9 @@
-from contextlib import redirect_stdout
+from contextlib import redirect_stderr, redirect_stdout
 from mage_ai.data_preparation.executors.pipeline_executor import PipelineExecutor
 from mage_ai.data_preparation.models.constants import BlockType
 from mage_ai.data_preparation.models.pipeline import Pipeline
 from mage_ai.data_preparation.shared.stream import StreamToLogger
-from typing import Callable, Dict
+from typing import Callable, Dict, List, Union
 import yaml
 
 
@@ -74,7 +74,8 @@ def execute(
             stdout = StreamToLogger(self.logger)
         try:
             with redirect_stdout(stdout):
-                self.__execute_in_python()
+                with redirect_stderr(stdout):
+                    self.__execute_in_python()
         except Exception as e:
             if not build_block_output_stdout:
                 self.logger.exception(
@@ -90,13 +91,17 @@ def __execute_in_python(self):
         sink_config = yaml.safe_load(self.sink_block.content)
         source = SourceFactory.get_source(source_config)
         sink = SinkFactory.get_sink(sink_config)
-        for messages in source.batch_read():
+
+        def handle_batch_events(messages: List[Union[Dict, str]]):
             if self.transformer_block is not None:
                 messages = self.transformer_block.execute_block(
                     input_args=[messages],
                 )['output']
             sink.batch_write(messages)
 
+        # Long running method
+        source.batch_read(handler=handle_batch_events)
+
     def __excute_in_flink(self):
         """
         TODO: Implement this method

mage_ai/streaming/constants.py

@@ -2,6 +2,7 @@
 
 
 class SourceType(str, Enum):
+    AZURE_EVENT_HUB = 'azure_event_hub'
     KAFKA = 'kafka'
 
 

mage_ai/streaming/sinks/base.py

@@ -1,7 +1,20 @@
 from abc import ABC, abstractmethod
+from typing import Dict
 
 
 class BaseSink(ABC):
+    config_class = None
+
+    def __init__(self, config: Dict):
+        if self.config_class is not None:
+            if 'connector_type' in config:
+                config.pop('connector_type')
+            self.config = self.config_class.load(config=config)
+        self.init_client()
+
+    def init_client():
+        pass
+
     @abstractmethod
     def write(self, data):
         pass

mage_ai/streaming/sinks/opensearch.py

@@ -1,4 +1,6 @@
 from dataclasses import dataclass
+from mage_ai.shared.config import BaseConfig
+from mage_ai.streaming.sinks.base import BaseSink
 from opensearchpy import OpenSearch, RequestsHttpConnection
 from opensearchpy.helpers import bulk
 from requests_aws4auth import AWS4Auth
@@ -8,19 +10,17 @@
 
 
 @dataclass
-class OpensearchSinkConfig:
+class OpensearchConfig(BaseConfig):
     host: str
     index_name: str
     verify_certs: bool = True
     http_auth: str = '@awsauth'
 
 
-class OpenSearchSink():
-    def __init__(self, config: Dict):
-        if 'connector_type' in config:
-            config.pop('connector_type')
-        self.config = OpensearchSinkConfig(**config)
+class OpenSearchSink(BaseSink):
+    config_class = OpensearchConfig
 
+    def init_client(self):
         # Initialize opensearch client
         if self.config.http_auth == '@awsauth':
             session = boto3.Session()

mage_ai/streaming/sources/azure_event_hub.py

@@ -0,0 +1,90 @@
+from azure.eventhub import EventHubConsumerClient
+from dataclasses import dataclass
+from mage_ai.shared.config import BaseConfig
+from mage_ai.streaming.sources.base import BaseSource
+from typing import Callable, List
+import traceback
+
+
+@dataclass
+class AzureEventHubConfig(BaseConfig):
+    connection_str: str
+    eventhub_name: str
+    consumer_group: str = '$Default'
+
+
+class AzureEventHubSource(BaseSource):
+    config_class = AzureEventHubConfig
+
+    def init_client(self):
+        self.consumer_client = EventHubConsumerClient.from_connection_string(
+            conn_str=self.config.connection_str,
+            consumer_group=self.config.consumer_group,
+            eventhub_name=self.config.eventhub_name,
+        )
+
+    def read(self, handler: Callable):
+        try:
+            def on_event(partition_context, event):
+                print(f'Received event from partition: {partition_context.partition_id}.')
+                print(f'Event: {event}')
+
+                handler(dict(data=event.body_as_str()))
+
+            with self.consumer_client:
+                self.consumer_client.receive(
+                    on_event=on_event,
+                    on_partition_initialize=self.on_partition_initialize,
+                    on_partition_close=self.on_partition_close,
+                    on_error=self.on_error,
+                    starting_position='-1',  # '-1' is from the beginning of the partition.
+                )
+        except KeyboardInterrupt:
+            print('Stopped receiving.')
+
+    def batch_read(self, handler: Callable):
+        try:
+            def on_event_batch(partition_context, event_batch: List):
+                if len(event_batch) == 0:
+                    return
+                print(f'Partition {partition_context.partition_id},'
+                      f'Received count: {len(event_batch)}')
+                print(f'Sample event: {event_batch[0]}')
+
+                # Handle events
+                try:
+                    handler([dict(data=e.body_as_str()) for e in event_batch])
+                except Exception as e:
+                    traceback.print_exc()
+                    raise e
+
+                partition_context.update_checkpoint()
+
+            with self.consumer_client:
+                self.consumer_client.receive_batch(
+                    on_event_batch=on_event_batch,
+                    max_batch_size=100,
+                    on_partition_initialize=self.on_partition_initialize,
+                    on_partition_close=self.on_partition_close,
+                    on_error=self.on_error,
+                    starting_position='-1',  # '-1' is from the beginning of the partition.
+                )
+        except KeyboardInterrupt:
+            print('Stopped receiving.')
+
+    def test_connection(self):
+        return True
+
+    def on_partition_initialize(partition_context):
+        print(f'Partition: {partition_context.partition_id} has been initialized.')
+
+    def on_partition_close(partition_context, reason):
+        print(f'Partition: {partition_context.partition_id} has been closed, '
+              f'reason for closing: {reason}.')
+
+    def on_error(partition_context, error):
+        if partition_context:
+            print(f'An exception: {partition_context.partition_id} occurred during'
+                  f' receiving from Partition: {error}.')
+        else:
+            print(f'An exception: {error} occurred during the load balance process.')

mage_ai/streaming/sources/base.py

@@ -1,7 +1,20 @@
 from abc import ABC, abstractmethod
+from typing import Dict
 
 
 class BaseSource(ABC):
+    config_class = None
+
+    def __init__(self, config: Dict):
+        if self.config_class is not None:
+            if 'connector_type' in config:
+                config.pop('connector_type')
+            self.config = self.config_class.load(config=config)
+        self.init_client()
+
+    def init_client():
+        pass
+
     @abstractmethod
     def read(self):
         pass

mage_ai/streaming/sources/kafka.py

@@ -1,7 +1,8 @@
 from dataclasses import dataclass
 from kafka import KafkaConsumer
 from mage_ai.shared.config import BaseConfig
-from typing import Dict
+from mage_ai.streaming.sources.base import BaseSource
+from typing import Callable, Dict
 import json
 import time
 
@@ -18,7 +19,7 @@ class SSLConfig:
 
 
 @dataclass
-class KafkaSourceConfig(BaseConfig):
+class KafkaConfig(BaseConfig):
     bootstrap_server: str
     consumer_group: str
     topic: str
@@ -34,12 +35,10 @@ def parse_config(self, config: Dict) -> Dict:
         return config
 
 
-class KafkaSource:
-    def __init__(self, config: Dict):
-        if 'connector_type' in config:
-            config.pop('connector_type')
-        self.config = KafkaSourceConfig.load(config=config)
+class KafkaSource(BaseSource):
+    config_class = KafkaConfig
 
+    def init_client(self):
         print('Start initializing kafka consumer.')
         # Initialize kafka consumer
         consumer_kwargs = dict(
@@ -61,14 +60,14 @@ def __init__(self, config: Dict):
         )
         print('Finish initializing kafka consumer.')
 
-    def read(self):
+    def read(self, handler: Callable):
         print('Start consuming messages from kafka.')
         for message in self.consumer:
             self.__print_message(message)
             data = json.loads(message.value.decode('utf-8'))
-            yield data
+            handler(data)
 
-    def batch_read(self):
+    def batch_read(self, handler: Callable):
         print('Start consuming messages from kafka.')
         if self.config.batch_size > 0:
             batch_size = self.config.batch_size
@@ -87,7 +86,7 @@ def batch_read(self):
                     self.__print_message(message)
                     message_values.append(json.loads(message.value.decode('utf-8')))
             if len(message_values) > 0:
-                yield message_values
+                handler(message_values)
 
     def test_connection(self):
         return True

mage_ai/streaming/sources/source_factory.py

@@ -9,3 +9,9 @@ def get_source(self, config: Dict):
         if connector_type == SourceType.KAFKA:
             from mage_ai.streaming.sources.kafka import KafkaSource
             return KafkaSource(config)
+        elif connector_type == SourceType.AZURE_EVENT_HUB:
+            from mage_ai.streaming.sources.azure_event_hub import AzureEventHubSource
+            return AzureEventHubSource(config)
+        raise Exception(
+            f'Consuming data from {connector_type} is not supported in streaming pipelines yet.',
+        )

mage_ai/tests/base_test.py

@@ -6,7 +6,7 @@
 import unittest
 
 
-class TestCase(unittest.TestCase):
+class DBTestCase(unittest.TestCase):
     def setUp(self):
         pass
 
@@ -28,3 +28,11 @@ def tearDownClass(self):
         db_connection.close_session()
         os.remove(TEST_DB)
         super().tearDownClass()
+
+
+class TestCase(unittest.TestCase):
+    def setUp(self):
+        pass
+
+    def tearDown(self):
+        pass

mage_ai/tests/data_preparation/models/test_block.py

@@ -3,13 +3,13 @@
 from mage_ai.data_preparation.models.block import Block, BlockType
 from mage_ai.data_preparation.models.pipeline import Pipeline
 from mage_ai.data_preparation.variable_manager import VariableManager
-from mage_ai.tests.base_test import TestCase
+from mage_ai.tests.base_test import DBTestCase
 from pandas.util.testing import assert_frame_equal
 import os
 import pandas as pd
 
 
-class BlockTest(TestCase):
+class BlockTest(DBTestCase):
     def test_create(self):
         block1 = Block.create('test_transformer', 'transformer', self.repo_path)
         block2 = Block.create('test data loader', BlockType.DATA_LOADER, self.repo_path)

mage_ai/tests/data_preparation/models/test_pipeline.py

@@ -1,11 +1,11 @@
 from mage_ai.data_preparation.models.block import Block
 from mage_ai.data_preparation.models.pipeline import InvalidPipelineError, Pipeline
 from mage_ai.data_preparation.models.widget import Widget
-from mage_ai.tests.base_test import TestCase
+from mage_ai.tests.base_test import DBTestCase
 import os
 
 
-class PipelineTest(TestCase):
+class PipelineTest(DBTestCase):
     def test_create(self):
         pipeline = Pipeline.create(
             'test pipeline',

mage_ai/tests/data_preparation/models/test_variable.py

@@ -1,13 +1,13 @@
 from mage_ai.data_preparation.models.block import Block
 from mage_ai.data_preparation.models.pipeline import Pipeline
 from mage_ai.data_preparation.models.variable import Variable, VariableType
-from mage_ai.tests.base_test import TestCase
+from mage_ai.tests.base_test import DBTestCase
 from pandas.util.testing import assert_frame_equal
 import os
 import pandas as pd
 
 
-class VariableTest(TestCase):
+class VariableTest(DBTestCase):
     def test_write_and_read_data(self):
         pipeline = self.__create_pipeline('test pipeline 1')
         variable1 = Variable('var1', pipeline.dir_path, 'block1')

mage_ai/tests/data_preparation/test_variable_manager.py

@@ -7,12 +7,12 @@
     get_global_variable,
     set_global_variable,
 )
-from mage_ai.tests.base_test import TestCase
+from mage_ai.tests.base_test import DBTestCase
 from pandas.util.testing import assert_frame_equal
 import pandas as pd
 
 
-class VariableManagerTest(TestCase):
+class VariableManagerTest(DBTestCase):
     def test_add_and_get_variable(self):
         self.__create_pipeline('test pipeline 1')
         variable_manager = VariableManager(self.repo_path)

mage_ai/tests/io/test_config.py

@@ -1,10 +1,10 @@
 from mage_ai.io.config import ConfigFileLoader, ConfigKey, EnvironmentVariableLoader
-from mage_ai.tests.base_test import TestCase
+from mage_ai.tests.base_test import DBTestCase
 from pathlib import Path
 from unittest import mock
 
 
-class ConfigLoaderTests(TestCase):
+class ConfigLoaderTests(DBTestCase):
     def setUp(self):
         super().setUp()
         self.test_path = Path('./test')