mage-ai 0.9.66__py3-none-any.whl → 0.9.67__py3-none-any.whl

mage_ai/shared/logger.py

@@ -5,6 +5,7 @@ from contextlib import contextmanager, redirect_stdout
 from enum import Enum
 from typing import Callable, List
 
+from mage_ai.settings import SERVER_LOGGING_TEMPLATE
 from mage_ai.shared.hash import merge_dict
 
 logger = logging.getLogger(__name__)
@@ -84,16 +85,21 @@ class JSONFormatter(logging.Formatter):
 
 
 def set_logging_format(logging_format: str = None, level: str = None) -> None:
-    handler = logging.StreamHandler()
-    handler.setFormatter(JSONFormatter())
-
-    root_logger = logging.getLogger()
     if isinstance(logging_format, str):
         logging_format = logging_format.lower()
+
+    root_logger = logging.getLogger()
+    if len(root_logger.handlers) > 0:
+        root_logger.removeHandler(root_logger.handlers[0])
+
+    handler = logging.StreamHandler()
     if logging_format == 'json':
-        if len(root_logger.handlers) > 0:
-            root_logger.removeHandler(root_logger.handlers[0])
+        handler.setFormatter(JSONFormatter())
         root_logger.addHandler(handler)
+    else:
+        handler.setFormatter(logging.Formatter(SERVER_LOGGING_TEMPLATE))
+        root_logger.addHandler(handler)
+
     if level:
         try:
             root_logger.setLevel(level.upper())

mage_ai/streaming/sources/base_python.py

@@ -0,0 +1,30 @@
+from typing import Callable
+
+from mage_ai.streaming.sources.base import BaseSource
+
+
+class BasePythonSource(BaseSource):
+    def __init__(self, **kwargs):
+        """
+        Not require config in the python source
+        """
+        super().__init__(None, **kwargs)
+
+    def init_client(self):
+        """
+        Intialize the client for the source.
+        """
+
+    def read(self, handler: Callable):
+        """
+        Read the message from the source and use handler to process the message.
+
+        This method only needs to be implemented when consume_method is 'READ'.
+        """
+
+    def batch_read(self, handler: Callable):
+        """
+        Batch read the messages from the source and use handler to process the messages.
+
+        This method only needs to be implemented when consume_method is 'BATCH_READ'.
+        """

mage_ai/streaming/sources/source_factory.py

@@ -1,5 +1,6 @@
 from typing import Dict
 
+from mage_ai.data_preparation.decorators import collect_decorated_objs
 from mage_ai.streaming.constants import SourceType
 
 
@@ -51,3 +52,27 @@ class SourceFactory:
         raise Exception(
             f'Consuming data from {connector_type} is not supported in streaming pipelines yet.',
         )
+
+    @classmethod
+    def get_python_source(self, content: str, **kwargs):
+        """
+        Find the class that's decorated with streaming_source from the source code.
+
+        Args:
+            content (str): The python code that contains the streaming source implementation.
+            **kwargs: {'global_vars': {...}}
+
+        Returns:
+            The initialized class object.
+
+        Raises:
+            Exception: Description
+        """
+        decorated_sources = []
+
+        exec(content, {'streaming_source': collect_decorated_objs(decorated_sources)})
+
+        if not decorated_sources:
+            raise Exception('Not find the class that has streaming_source decorator.')
+
+        return decorated_sources[0](**kwargs)

mage_ai/tests/data_preparation/models/test_block.py

@@ -7,14 +7,17 @@ from faker import Faker
 from pandas.testing import assert_frame_equal
 
 # from mage_ai.data_cleaner.column_types.constants import ColumnType
-from mage_ai.data_preparation.models.block import Block, BlockType
+from mage_ai.data_preparation.models.block import Block, BlockType, CallbackBlock
 from mage_ai.data_preparation.models.block.errors import HasDownstreamDependencies
 from mage_ai.data_preparation.models.pipeline import Pipeline
 from mage_ai.data_preparation.repo_manager import get_repo_config
 from mage_ai.data_preparation.variable_manager import VariableManager
 from mage_ai.shared.path_fixer import add_root_repo_path_to_relative_path
 from mage_ai.tests.base_test import DBTestCase
-from mage_ai.tests.factory import create_integration_pipeline_with_blocks
+from mage_ai.tests.factory import (
+    create_integration_pipeline_with_blocks,
+    create_pipeline,
+)
 from mage_ai.tests.shared.mixins import ProjectPlatformMixin
 
 
@@ -763,3 +766,37 @@ class BlockProjectPlatformTests(ProjectPlatformMixin):
                 )
 
                 self.assertEqual(block.file_path, add_root_repo_path_to_relative_path(path))
+
+
+class CallbackBlockTest(DBTestCase):
+    def setUp(self):
+        self.pipeline = create_pipeline('callback_pipeline', self.repo_path)
+
+    def tearDown(self):
+        self.pipeline.delete()
+
+    def test_create_global_vars_from_parent_block(self):
+        parent_block = Block.create(
+            'test_data_loader',
+            'data_loader',
+            self.repo_path,
+            pipeline=self.pipeline,
+        )
+        callback_block = CallbackBlock.create(parent_block.name)
+        self.pipeline.add_block(callback_block)
+        parent_block = parent_block.update(
+            dict(callback_blocks=[callback_block.uuid])
+        )
+        parent_block.global_vars = dict(
+            configuration=dict(table_name='load_data_table')
+        )
+
+        global_vars = dict(
+            random_var=1,
+        )
+        new_vars = callback_block._create_global_vars(
+            global_vars,
+            parent_block=parent_block,
+        )
+
+        self.assertIsNotNone(new_vars.get('configuration'))

{mage_ai-0.9.66.dist-info → mage_ai-0.9.67.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: mage-ai
-Version: 0.9.66
+Version: 0.9.67
 Summary: Mage is a tool for building and deploying data pipelines.
 Home-page: https://github.com/mage-ai/mage-ai
 Author: Mage