mage-ai 0.9.45__py3-none-any.whl → 0.9.47__py3-none-any.whl

mage_ai/tests/shared/mixins.py

@@ -0,0 +1,328 @@
+import json
+import os
+from typing import Callable, Dict, List
+
+from mage_ai.authentication.permissions.constants import EntityName
+from mage_ai.data_preparation.models.constants import PipelineType
+from mage_ai.data_preparation.models.global_hooks.constants import (
+    HookOutputKey,
+    PredicateAndOrOperator,
+    PredicateOperator,
+    PredicateValueDataType,
+)
+from mage_ai.data_preparation.models.global_hooks.models import (
+    GlobalHooks,
+    Hook,
+    HookCondition,
+    HookOperation,
+    HookOutputBlock,
+    HookOutputSettings,
+    HookStage,
+    HookStrategy,
+)
+from mage_ai.data_preparation.models.global_hooks.predicates import (
+    HookPredicate,
+    PredicateValueType,
+)
+from mage_ai.shared.hash import merge_dict
+from mage_ai.tests.api.operations.test_base import BaseApiTestCase
+from mage_ai.tests.factory import build_pipeline_with_blocks_and_content
+
+
+def build_content(query: Dict) -> str:
+    return f"""
+if 'data_loader' not in globals():
+    from mage_ai.data_preparation.decorators import data_loader
+
+
+if 'transformer' not in globals():
+    from mage_ai.data_preparation.decorators import transformer
+
+
+if 'data_exporter' not in globals():
+    from mage_ai.data_preparation.decorators import data_exporter
+
+
+@data_loader
+def load_data(*args, **kwargs):
+    return {json.dumps(query)}
+
+
+@transformer
+def transform(*args, **kwargs):
+    return {json.dumps(query)}
+
+
+@data_exporter
+def export_data(*args, **kwargs):
+    return {json.dumps(query)}
+"""
+
+
+def build_hooks(
+    test_case,
+    pipeline: PipelineType,
+    global_hooks: GlobalHooks = None,
+    hook_settings: Dict = None,
+    matching_predicate_resource: Dict = None,
+    operation_type: HookOperation = None,
+    resource_type: EntityName = None,
+    test_predicate_after: bool = False,
+    test_predicate_before: bool = False,
+    predicate_match: HookPredicate = None,
+    predicate_match_after: HookPredicate = None,
+    predicate_match_before: HookPredicate = None,
+    predicate_miss: HookPredicate = None,
+    snapshot: bool = True,
+) -> List[Hook]:
+    if not global_hooks:
+        global_hooks = GlobalHooks.load_from_file()
+
+    if matching_predicate_resource is None:
+        matching_predicate_resource = dict(
+            name=pipeline.name,
+            type=pipeline.type,
+        )
+
+    if predicate_match is None and predicate_match_after is None and predicate_match_before is None:
+        predicate_match = HookPredicate.load(
+            and_or_operator=PredicateAndOrOperator.AND,
+            predicates=[HookPredicate.load(
+                left_value=v,
+                left_value_type=PredicateValueType.load(
+                    value_data_type=PredicateValueDataType.STRING,
+                ),
+                operator=PredicateOperator.EQUALS,
+                right_value=v,
+                right_value_type=PredicateValueType.load(
+                        value_data_type=PredicateValueDataType.STRING,
+                ),
+            ) for v in matching_predicate_resource.values()],
+        )
+
+    if predicate_miss is None:
+        predicate_miss = HookPredicate.load(
+            and_or_operator=PredicateAndOrOperator.AND,
+            predicates=[HookPredicate.load(
+                left_value=v,
+                left_value_type=PredicateValueType.load(
+                    value_data_type=PredicateValueDataType.STRING,
+                ),
+                operator=PredicateOperator.EQUALS,
+                right_value=None,
+                right_value_type=PredicateValueType.load(
+                        value_data_type=PredicateValueDataType.STRING,
+                ),
+            ) for v in matching_predicate_resource.values()],
+        )
+
+    if not operation_type:
+        operation_type = HookOperation.DETAIL
+
+    if not resource_type:
+        resource_type = EntityName.Pipeline
+
+    hooks = []
+    hooks_match = []
+    hooks_miss = []
+    for idx, hook_dict in enumerate([
+        dict(
+            stages=[HookStage.BEFORE],
+        ),
+        dict(
+            conditions=[HookCondition.FAILURE],
+            stages=[HookStage.BEFORE],
+        ),
+        dict(
+            conditions=[HookCondition.SUCCESS],
+            stages=[HookStage.AFTER],
+        ),
+        dict(
+            conditions=[HookCondition.SUCCESS],
+            stages=[HookStage.AFTER],
+        ),
+    ]):
+        hook_payload = hook_dict.copy()
+        if hook_settings and idx in hook_settings:
+            hook_setting = hook_settings[idx]
+            hook_payload.update(hook_setting)
+
+        predicate_use = predicate_match
+        if predicate_match_after and HookStage.AFTER in hook_dict['stages']:
+            predicate_use = predicate_match_after
+        elif predicate_match_before and HookStage.BEFORE in hook_dict['stages']:
+            predicate_use = predicate_match_before
+
+        hook = Hook.load(
+            operation_type=operation_type,
+            predicate=predicate_use,
+            resource_type=resource_type,
+            strategies=[HookStrategy.RAISE],
+            uuid=f'hook{idx}_{test_case.faker.unique.name()}',
+            **hook_payload,
+        )
+
+        hook_miss = Hook.load(
+            operation_type=hook.operation_type,
+            resource_type=hook.resource_type,
+            **hook.to_dict(),
+        )
+        hook_miss.uuid = test_case.faker.unique.name()
+        hook_miss.predicate = predicate_miss
+
+        if snapshot:
+            hook.snapshot()
+            hook_miss.snapshot()
+
+        hooks.append(hook)
+        hooks.append(hook_miss)
+        hooks_match.append(hook)
+        hooks_miss.append(hook_miss)
+
+        global_hooks.add_hook(hook)
+        global_hooks.add_hook(hook_miss)
+
+    global_hooks.save()
+
+    return global_hooks, hooks, hooks_match, hooks_miss
+
+
+class GlobalHooksMixin(BaseApiTestCase):
+    def setUp(self):
+        self.pipelines_created_for_testing = []
+
+    async def setUpAsync(
+        self,
+        block_settings: Dict = None,
+        hook_settings: Callable = None,
+        matching_predicate_resource: Dict = None,
+        operation_type: HookOperation = None,
+        pipeline_type: PipelineType = None,
+        predicate_match: HookPredicate = None,
+        predicate_match_after: HookPredicate = None,
+        predicate_match_before: HookPredicate = None,
+        predicate_miss: HookPredicate = None,
+        resource_type: EntityName = None,
+        snapshot: bool = True,
+    ):
+        block_settings_init = {
+            0: dict(content=build_content({
+                'type[]': [PipelineType.STREAMING.value],
+            })),
+            1: dict(content=build_content({
+                'type[]': [PipelineType.INTEGRATION.value],
+            })),
+            2: dict(content=build_content(dict(powers=dict(fire=1)))),
+            3: dict(content=build_content(dict(level=2))),
+        }
+
+        if block_settings:
+            block_settings = merge_dict(block_settings_init, block_settings)
+        else:
+            block_settings = block_settings_init
+
+        pipeline1, blocks1 = await build_pipeline_with_blocks_and_content(
+            self,
+            block_settings=block_settings,
+            pipeline_type=pipeline_type or PipelineType.PYTHON,
+        )
+        pipeline2, blocks2 = await build_pipeline_with_blocks_and_content(
+            self,
+            block_settings=block_settings,
+            pipeline_type=pipeline_type or PipelineType.PYTHON,
+        )
+        block11, block12, block13, block14 = blocks1
+        block21, block22, block23, block24 = blocks2
+
+        hook_settings_use = {
+            0: dict(
+                outputs=[
+                    HookOutputSettings.load(
+                        block=HookOutputBlock.load(uuid=block11.uuid),
+                        key=HookOutputKey.QUERY,
+                    ),
+                    HookOutputSettings.load(
+                        block=HookOutputBlock.load(uuid=block12.uuid),
+                        key=HookOutputKey.QUERY,
+                    ),
+                ],
+                pipeline=dict(
+                    uuid=pipeline1.uuid,
+                ),
+            ),
+            1: dict(
+                pipeline=dict(
+                    uuid=pipeline1.uuid,
+                ),
+            ),
+            2: dict(
+                outputs=[
+                    HookOutputSettings.load(
+                        block=HookOutputBlock.load(uuid=block13.uuid),
+                        key=HookOutputKey.METADATA,
+                    ),
+                ],
+                pipeline=dict(
+                    uuid=pipeline1.uuid,
+                ),
+            ),
+            3: dict(
+                outputs=[
+                    HookOutputSettings.load(
+                        block=HookOutputBlock.load(uuid=block24.uuid),
+                        key=HookOutputKey.METADATA,
+                        keys=['water'],
+                    ),
+                ],
+                pipeline=dict(
+                    uuid=pipeline2.uuid,
+                ),
+            ),
+        }
+
+        if hook_settings:
+            hook_settings_use = merge_dict(hook_settings_use, hook_settings(dict(
+                blocks1=blocks1,
+                blocks2=blocks2,
+                pipeline1=pipeline1,
+                pipeline2=pipeline2,
+            )))
+
+        global_hooks, hooks, hooks_match, hooks_miss = build_hooks(
+            self,
+            hook_settings=hook_settings_use,
+            matching_predicate_resource=matching_predicate_resource,
+            operation_type=operation_type,
+            pipeline=pipeline1,
+            predicate_match=predicate_match,
+            predicate_match_after=predicate_match_after,
+            predicate_match_before=predicate_match_before,
+            predicate_miss=predicate_miss,
+            resource_type=resource_type,
+            snapshot=snapshot,
+        )
+
+        self.blocks1 = blocks1
+        self.blocks2 = blocks2
+        self.global_hooks = global_hooks
+        self.hooks = hooks
+        self.hooks_match = hooks_match
+        self.hooks_miss = hooks_miss
+        self.pipeline1 = pipeline1
+        self.pipeline2 = pipeline2
+
+    def tearDown(self):
+        file_path = GlobalHooks.file_path()
+        if os.path.exists(file_path):
+            os.remove(file_path)
+
+        self.pipeline1.delete()
+        self.pipeline2.delete()
+
+        super().tearDown()
+
+    def bootstrap(self):
+        pass
+
+    def cleanup(self):
+        pass

mage_ai/tests/shared/test_hash.py

@@ -1,4 +1,4 @@
-from mage_ai.shared.hash import get_json_value
+from mage_ai.shared.hash import get_json_value, set_value
 from mage_ai.tests.base_test import TestCase
 
 
@@ -20,3 +20,19 @@ class HashTests(TestCase):
             get_json_value('{"k1": "v1", "k2": {"k21": "v21", "k22": "v22"} }', 'k2.k22'),
             'v22'
         )
+
+    def test_get_value(self):
+        obj = dict(l1=dict(power=1))
+        value = 'mage'
+
+        self.assertEqual(
+            set_value(obj, ['l1', 'l2', 'l3'], value),
+            dict(
+                l1=dict(
+                    l2=dict(
+                        l3=value,
+                    ),
+                    power=1,
+                ),
+            ),
+        )

mage_ai/tests/streaming/sinks/test_rabbitmq.py

@@ -0,0 +1,36 @@
+from unittest.mock import patch
+
+from mage_ai.streaming.sinks.rabbitmq import RabbitMQSink
+from mage_ai.tests.base_test import TestCase
+
+
+class KinesisTests(TestCase):
+    def test_init(self):
+        with patch.object(RabbitMQSink, 'init_client') as mock_init_client:
+            RabbitMQSink(dict(
+                connector_type='rabbitmq',
+                connection_host='localhost',
+                connection_port='5672',
+                queue_name='queue_name',
+                username='guest',
+                password='guest',
+                amqp_url_virtual_host='%2f',
+            ))
+            mock_init_client.assert_called_once()
+
+    def test_init_invalid_config(self):
+        with patch.object(RabbitMQSink, 'init_client') as mock_init_client:
+            with self.assertRaises(Exception) as context:
+                RabbitMQSink(dict(
+                    connector_type='rabbitmq',
+                    connection_port='5672',
+                    queue_name='queue_name',
+                    username='guest',
+                    password='guest',
+                    amqp_url_virtual_host='%2f',
+                ))
+            self.assertTrue(
+                '__init__() missing 1 required positional argument: \'connection_host\''
+                in str(context.exception),
+            )
+            self.assertEqual(mock_init_client.call_count, 0)

mage_ai/tests/streaming/sources/test_activemq.py

@@ -0,0 +1,32 @@
+from unittest.mock import patch
+
+from mage_ai.streaming.sources.activemq import ActiveMQSource
+from mage_ai.tests.base_test import TestCase
+
+
+class ActiveMQTests(TestCase):
+    def test_init(self):
+        with patch.object(ActiveMQSource, 'init_client') as mock_init_client:
+            ActiveMQSource(dict(
+                connection_host='localhost',
+                connection_port='61613',
+                queue_name='test_queue',
+                username='guest',
+                password='guest',
+            ))
+            mock_init_client.assert_called_once()
+
+    def test_init_invalid_config(self):
+        with patch.object(ActiveMQSource, 'init_client') as mock_init_client:
+            with self.assertRaises(Exception) as context:
+                ActiveMQSource(dict(
+                    connection_port='61613',
+                    queue_name='test_queue',
+                    username='guest',
+                    password='guest',
+                ))
+            self.assertTrue(
+                '__init__() missing 1 required positional argument: \'connection_host\''
+                in str(context.exception),
+            )
+            self.assertEqual(mock_init_client.call_count, 0)

mage_ai/tests/streaming/sources/test_nats_js.py

@@ -0,0 +1,32 @@
+from unittest.mock import patch
+
+from mage_ai.streaming.sources.nats_js import NATSSource
+from mage_ai.tests.base_test import TestCase
+
+
+class NATSTests(TestCase):
+    def test_init(self):
+        with patch.object(NATSSource, 'init_client') as mock_init_client:
+            NATSSource(
+                dict(
+                    connector_type='nats',
+                    server_url='nats://localhost:4222',
+                    subject='test_subject',
+                    consumer_name='test_consumer',
+                    stream_name='test_stream',
+                )
+            )
+            mock_init_client.assert_called_once()
+
+    def test_init_invalid_config(self):
+        invalid_config = {
+            'subject': 'test_subject',
+            'consumer_name': 'test_consumer',
+        }
+
+        with patch.object(NATSSource, 'init_client') as mock_init_client:
+            with self.assertRaises(Exception) as context:
+                NATSSource(invalid_config)
+            error_msg = "missing 2 required positional arguments: 'server_url' and 'stream_name'"
+            self.assertIn(error_msg, str(context.exception))
+            self.assertEqual(mock_init_client.call_count, 0)

mage_ai/tests/streaming/sources/test_source_factory.py

@@ -1,13 +1,27 @@
+from unittest.mock import patch
+
+from mage_ai.streaming.sources.activemq import ActiveMQSource
 from mage_ai.streaming.sources.azure_event_hub import AzureEventHubSource
 from mage_ai.streaming.sources.google_cloud_pubsub import GoogleCloudPubSubSource
 from mage_ai.streaming.sources.kafka import KafkaSource
+from mage_ai.streaming.sources.nats_js import NATSSource
 from mage_ai.streaming.sources.rabbitmq import RabbitMQSource
 from mage_ai.streaming.sources.source_factory import SourceFactory
 from mage_ai.tests.base_test import TestCase
-from unittest.mock import patch
 
 
 class SourceFactoryTests(TestCase):
+    def test_get_source_nats(self):
+        with patch.object(NATSSource,
+                          '__init__',
+                          return_value=None) as mock_init:
+            config = dict(
+                connector_type='nats',
+            )
+            source = SourceFactory.get_source(config)
+            self.assertIsInstance(source, NATSSource)
+            mock_init.assert_called_once_with(config)
+
     def test_get_source_kafka(self):
         with patch.object(KafkaSource,
                           '__init__',
@@ -52,6 +66,17 @@ class SourceFactoryTests(TestCase):
             self.assertIsInstance(source, RabbitMQSource)
             mock_init.assert_called_once_with(config)
 
+    def test_get_source_activemq(self):
+        with patch.object(ActiveMQSource,
+                          '__init__',
+                          return_value=None) as mock_init:
+            config = dict(
+                connector_type='activemq',
+            )
+            source = SourceFactory.get_source(config)
+            self.assertIsInstance(source, ActiveMQSource)
+            mock_init.assert_called_once_with(config)
+
     def test_get_source_other(self):
         s = "Consuming data from random is not supported " \
             "in streaming pipelines yet."

{mage_ai-0.9.45.dist-info → mage_ai-0.9.47.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: mage-ai
-Version: 0.9.45
+Version: 0.9.47
 Summary: Mage is a tool for building and deploying data pipelines.
 Home-page: https://github.com/mage-ai/mage-ai
 Author: Mage
@@ -41,6 +41,7 @@ Requires-Dist: pandas >=1.3.0
 Requires-Dist: polars <=0.19.2
 Requires-Dist: protobuf
 Requires-Dist: pyarrow ==10.0.1
+Requires-Dist: pyarrow-hotfix ==0.5
 Requires-Dist: python-dateutil ==2.8.2
 Requires-Dist: pytz ==2022.2.1
 Requires-Dist: pyyaml ~=6.0
@@ -55,7 +56,7 @@ Requires-Dist: sqlalchemy <2.0.0,>=1.4.20
 Requires-Dist: terminado ==0.17.1
 Requires-Dist: thefuzz[speedup] ==0.19.0
 Requires-Dist: tornado ==6.1
-Requires-Dist: typer[all] ==0.7.0
+Requires-Dist: typer[all] ==0.9.0
 Requires-Dist: typing-extensions ==4.5.0
 Provides-Extra: ai
 Requires-Dist: astor >=0.8.1 ; extra == 'ai'
@@ -73,6 +74,7 @@ Requires-Dist: azure-mgmt-containerinstance ==10.1.0 ; extra == 'all'
 Requires-Dist: azure-storage-blob ==12.14.1 ; extra == 'all'
 Requires-Dist: boto3 ==1.26.60 ; extra == 'all'
 Requires-Dist: botocore ==1.29.60 ; extra == 'all'
+Requires-Dist: chroma >=0.4.17 ; extra == 'all'
 Requires-Dist: clickhouse-connect ==0.5.20 ; extra == 'all'
 Requires-Dist: confluent-avro ; extra == 'all'
 Requires-Dist: db-dtypes ==1.0.5 ; extra == 'all'
@@ -99,12 +101,14 @@ Requires-Dist: great-expectations ==0.15.50 ; extra == 'all'
 Requires-Dist: gspread ==5.7.2 ; extra == 'all'
 Requires-Dist: influxdb-client ==1.36.1 ; extra == 'all'
 Requires-Dist: kafka-python ==2.0.2 ; extra == 'all'
-Requires-Dist: kubernetes ==25.3.0 ; extra == 'all'
+Requires-Dist: kubernetes >=28.1.0 ; extra == 'all'
 Requires-Dist: langchain >=0.0.222 ; extra == 'all'
 Requires-Dist: ldap3 ==2.9.1 ; extra == 'all'
+Requires-Dist: nats-py ==2.6.0 ; extra == 'all'
+Requires-Dist: nkeys ; extra == 'all'
 Requires-Dist: openai <1.0.0,>=0.27.8 ; extra == 'all'
 Requires-Dist: opensearch-py ==2.0.0 ; extra == 'all'
-Requires-Dist: opentelemetry-exporter-prometheus >=1.12.0rc1 ; extra == 'all'
+Requires-Dist: opentelemetry-exporter-prometheus >=0.41b0 ; extra == 'all'
 Requires-Dist: opentelemetry-instrumentation-tornado >=0.41b0 ; extra == 'all'
 Requires-Dist: oracledb ==1.3.1 ; extra == 'all'
 Requires-Dist: pika ==1.3.1 ; extra == 'all'
@@ -119,6 +123,7 @@ Requires-Dist: redshift-connector ==2.0.909 ; extra == 'all'
 Requires-Dist: requests-aws4auth ==1.1.2 ; extra == 'all'
 Requires-Dist: snowflake-connector-python ==3.2.1 ; 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'
 Provides-Extra: azure
 Requires-Dist: azure-eventhub ==5.11.2 ; extra == 'azure'
@@ -130,6 +135,8 @@ Requires-Dist: azure-storage-blob ==12.14.1 ; extra == 'azure'
 Provides-Extra: bigquery
 Requires-Dist: google-cloud-bigquery ~=3.0 ; extra == 'bigquery'
 Requires-Dist: db-dtypes ==1.0.5 ; extra == 'bigquery'
+Provides-Extra: chroma
+Requires-Dist: chroma >=0.4.17 ; extra == 'chroma'
 Provides-Extra: clickhouse
 Requires-Dist: clickhouse-connect ==0.5.20 ; extra == 'clickhouse'
 Provides-Extra: dbt
@@ -173,9 +180,12 @@ Requires-Dist: confluent-avro ; extra == 'streaming'
 Requires-Dist: influxdb-client ==1.36.1 ; extra == 'streaming'
 Requires-Dist: kafka-python ==2.0.2 ; extra == 'streaming'
 Requires-Dist: opensearch-py ==2.0.0 ; extra == 'streaming'
+Requires-Dist: nats-py ==2.6.0 ; extra == 'streaming'
+Requires-Dist: nkeys ; extra == 'streaming'
 Requires-Dist: pika ==1.3.1 ; extra == 'streaming'
 Requires-Dist: pymongo ==4.3.3 ; extra == 'streaming'
 Requires-Dist: requests-aws4auth ==1.1.2 ; extra == 'streaming'
+Requires-Dist: stomp.py ==8.1.0 ; extra == 'streaming'
 Requires-Dist: elasticsearch ==8.9.0 ; extra == 'streaming'
 
 <h1 align="center">