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

mage_ai/shared/path_fixer.py

@@ -29,6 +29,9 @@ def add_absolute_path(file_path: str, add_base_repo_path: bool = True) -> str:
                 full_path = os.path.join(repo_path, parts[0])
                 if os.path.exists(full_path):
                     full_path = os.path.join(repo_path, file_path)
+                elif os.path.exists(os.path.join(repo_path, 'dbt', parts[0])):
+                    # DBT v1 paths
+                    full_path = os.path.join(repo_path, 'dbt', file_path)
                 else:
                     full_path = find_directory(repo_path, lambda x: x.endswith(file_path))
 

mage_ai/streaming/sources/base.py

@@ -3,6 +3,8 @@ from abc import ABC, abstractmethod
 from enum import Enum
 from typing import Callable, Dict
 
+from mage_ai.shared.environments import is_test
+
 
 class SourceConsumeMethod(str, Enum):
     BATCH_READ = 'BATCH_READ'
@@ -22,6 +24,9 @@ class BaseSource(ABC):
         self.checkpoint_path = kwargs.get('checkpoint_path')
         self.checkpoint = self.read_checkpoint()
         self.init_client()
+        if not is_test():
+            # Not test the connection in unit tests
+            self.test_connection()
 
     @abstractmethod
     def init_client(self):

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/kafka.py

@@ -256,7 +256,8 @@ class KafkaSource(BaseSource):
                 handler(message_values)
 
     def test_connection(self):
-        return True
+        self.consumer._client.check_version(timeout=5)
+        self._print('Test connectino successfully.')
 
     def __deserialize_message(self, message):
         if self.config.serde_config is None:

mage_ai/streaming/sources/mongodb.py

@@ -18,6 +18,7 @@ class MongoDBConfig(BaseConfig):
     pipeline: Optional[_Pipeline] = None
     operation_time: Optional[Timestamp] = None
     start_after: Optional[Mapping[str, Any]] = None
+    full_document: Optional[str] = None
 
 
 class MongoSource(BaseSource):
@@ -46,6 +47,9 @@ class MongoSource(BaseSource):
         if hasattr(self.config, 'start_after') and self.config.start_after is not None:
             watch_args['start_after'] = self.config.start_after
 
+        if hasattr(self.config, 'full_document'):
+            watch_args['full_document'] = self.config.full_document
+
         return watch_args
 
     def batch_read(self, handler: Callable):

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/block/dbt/test_block.py

@@ -38,7 +38,7 @@ class DBTBlockTest(AsyncDBTestCase):
         self.block_sql.language = BlockLanguage.SQL
         self.block_sql.pipeline = pipeline
 
-        self.dbt_block_sql = DBTBlock(
+        self.dbt_block_sql = DBTBlock.create(
             name='test_dbt_block_sql',
             uuid='test_dbt_block_sql',
             block_type=BlockType.DBT,
@@ -50,7 +50,7 @@ class DBTBlockTest(AsyncDBTestCase):
             pipeline=pipeline,
         )
 
-        self.dbt_block_yaml = DBTBlock(
+        self.dbt_block_yaml = DBTBlock.create(
             name='test_dbt_block_yaml',
             uuid='test_dbt_block_yaml',
             block_type=BlockType.DBT,

mage_ai/tests/data_preparation/models/block/dbt/test_block_sql.py

@@ -23,7 +23,7 @@
 #     pipeline.repo_path = 'test_repo_path'
 #     pipeline.get_block.return_value = None
 
-#     return DBTBlock(
+#     return DBTBlock.create(
 #         name='test_dbt_block_sql',
 #         uuid='test_dbt_block_sql',
 #         block_type=BlockType.DBT,

mage_ai/tests/data_preparation/models/block/dbt/test_block_yaml.py

@@ -18,7 +18,7 @@
 
 
 # def build_block(pipeline, content: str) -> DBTBlock:
-#     return DBTBlock(
+#     return DBTBlock.create(
 #         name='test_dbt_block_yaml',
 #         uuid='test_dbt_block_yaml',
 #         block_type=BlockType.DBT,

mage_ai/tests/data_preparation/models/test_block.py

@@ -7,14 +7,18 @@ 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.block_factory import BlockFactory
 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
 
 
@@ -603,7 +607,7 @@ SELECT 1 AS id;
             upstream_block_uuids=['test_data_loader'],
         )
 
-        test_block = Block.get_block('test_data_loader', 'test_data_loader', 'data_loader')
+        test_block = BlockFactory.get_block('test_data_loader', 'test_data_loader', 'data_loader')
 
         self.assertRaises(HasDownstreamDependencies, block1.delete)
         self.assertRaises(HasDownstreamDependencies, test_block.delete)
@@ -763,3 +767,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/tests/orchestration/test_pipeline_scheduler.py

@@ -396,6 +396,7 @@ class PipelineSchedulerTests(DBTestCase):
                 error=f'Failed blocks: {block_runs[0].block_uuid}.',
                 pipeline=scheduler.pipeline,
                 pipeline_run=pipeline_run,
+                stacktrace=None,
             )
 
     def test_schedule_with_block_failures(self):
@@ -419,6 +420,7 @@ class PipelineSchedulerTests(DBTestCase):
                 error=f'Failed blocks: {block_runs[0].block_uuid}.',
                 pipeline=scheduler.pipeline,
                 pipeline_run=pipeline_run,
+                stacktrace=None,
             )
 
     @patch('mage_ai.orchestration.pipeline_scheduler_original.run_pipeline')

mage_ai/tests/settings/models/test_configuration_option.py

@@ -100,7 +100,7 @@ class ConfigurationOptionTest(AsyncDBTestCase):
             self.faker.unique.name(),
             os.path.join(self.repo_path, 'mage_platform'),
         )
-        block = DBTBlock(
+        block = DBTBlock.create(
             block_type='dbt',
             configuration=dict(file_path='dir2/mage.sql'),
             language='sql',
@@ -139,7 +139,7 @@ class ConfigurationOptionTest(AsyncDBTestCase):
                         self.faker.unique.name(),
                         os.path.join(self.repo_path, 'mage_platform'),
                     )
-                    block = DBTBlock(
+                    block = DBTBlock.create(
                         block_type='dbt',
                         configuration=dict(file_path='dir2/mage.sql'),
                         language='sql',

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

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: mage-ai
-Version: 0.9.66
+Version: 0.9.68
 Summary: Mage is a tool for building and deploying data pipelines.
 Home-page: https://github.com/mage-ai/mage-ai
 Author: Mage
@@ -128,7 +128,7 @@ Requires-Dist: pyodbc ==4.0.35 ; extra == 'all'
 Requires-Dist: redshift-connector ==2.0.915 ; extra == 'all'
 Requires-Dist: lxml ==4.9.4 ; extra == 'all'
 Requires-Dist: requests-aws4auth ==1.1.2 ; extra == 'all'
-Requires-Dist: snowflake-connector-python ==3.4.0 ; extra == 'all'
+Requires-Dist: snowflake-connector-python ==3.5.0 ; extra == 'all'
 Requires-Dist: sshtunnel ==0.4.0 ; extra == 'all'
 Requires-Dist: stomp.py ==8.1.0 ; extra == 'all'
 Requires-Dist: thefuzz[speedup] ==0.19.0 ; extra == 'all'
@@ -183,7 +183,7 @@ Provides-Extra: s3
 Requires-Dist: boto3 ==1.26.60 ; extra == 's3'
 Requires-Dist: botocore ==1.29.60 ; extra == 's3'
 Provides-Extra: snowflake
-Requires-Dist: snowflake-connector-python ==3.4.0 ; extra == 'snowflake'
+Requires-Dist: snowflake-connector-python ==3.5.0 ; extra == 'snowflake'
 Provides-Extra: spark
 Requires-Dist: boto3 ==1.26.60 ; extra == 'spark'
 Requires-Dist: botocore ==1.29.60 ; extra == 'spark'