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

mage_ai/streaming/sinks/rabbitmq.py

@@ -0,0 +1,76 @@
+import json
+from dataclasses import dataclass
+from typing import Dict, List
+
+import pika
+
+from mage_ai.shared.config import BaseConfig
+from mage_ai.streaming.sinks.base import BaseSink
+
+
+@dataclass
+class RabbitMQConfig(BaseConfig):
+    connection_host: str
+    connection_port: int
+    queue_name: str
+    username: str = 'guest'
+    password: str = 'guest'
+    amqp_url_virtual_host: str = r'%2f'
+
+
+class RabbitMQSink(BaseSink):
+    config_class = RabbitMQConfig
+
+    def init_client(self):
+        self._print('Start initializing producer.')
+        # Initialize RabbitMQ producer
+        queue_name = self.config.queue_name
+        username = self.config.username
+        password = self.config.password
+        connection_host = self.config.connection_host
+        connection_port = self.config.connection_port
+        vt_host = self.config.amqp_url_virtual_host
+
+        self._print(f'Starting to initialize producer for queue {queue_name}')
+
+        try:
+            generated_url = f"amqp://{username}:{password}@" \
+                            f"{connection_host}:{connection_port}/{vt_host}"
+
+            self._print(f'Trying to connect on {generated_url}')
+            self.connection = pika.BlockingConnection(
+                pika.URLParameters(
+                    generated_url
+                )
+            )
+
+            self.main_channel = self.connection.channel()
+            self._print('Finish initializing producer.')
+        except Exception as e:
+            self._print('Connection Error! Please check RabbitMQ connection')
+            raise e
+
+    def write(self, message: Dict):
+        pass
+
+    def batch_write(self, messages: List[Dict]):
+        if not messages:
+            return
+        self._print(
+            f'Batch ingest {len(messages)} messages. Sample: {messages[0]}'
+        )
+        for message in messages:
+            if isinstance(message, dict):
+                data = message.get('data', message)
+                metadata = message.get('metadata', None)
+            else:
+                data = message
+                metadata = None
+            message_properties = pika.BasicProperties(headers=metadata)
+            self.main_channel.basic_publish(
+                exchange='',
+                routing_key=self.config.queue_name,
+                body=json.dumps(data).encode('utf-8'),
+                properties=message_properties,
+                mandatory=True,
+            )

mage_ai/streaming/sinks/sink_factory.py

@@ -53,6 +53,10 @@ class SinkFactory:
             from mage_ai.streaming.sinks.postgres import PostgresSink
 
             return PostgresSink(config, **kwargs)
+        elif connector_type == SinkType.RABBITMQ:
+            from mage_ai.streaming.sinks.rabbitmq import RabbitMQSink
+
+            return RabbitMQSink(config, **kwargs)
         elif connector_type in GENERIC_IO_SINK_TYPES:
             from mage_ai.streaming.sinks.generic_io import GenericIOSink
 

mage_ai/tests/api/operations/test_operations_with_hooks.py

@@ -3,16 +3,25 @@
 
 # from mage_ai.authentication.permissions.constants import EntityName
 # from mage_ai.data_preparation.models.constants import BlockType, PipelineType
-# from mage_ai.data_preparation.models.global_hooks.constants import HookOutputKey
+# from mage_ai.data_preparation.models.global_hooks.constants import (
+#     HookOutputKey,
+#     PredicateAndOrOperator,
+#     PredicateObjectType,
+#     PredicateOperator,
+#     PredicateValueDataType,
+# )
 # from mage_ai.data_preparation.models.global_hooks.models import (
 #     GlobalHooks,
 #     HookCondition,
 #     HookOperation,
 #     HookOutputBlock,
 #     HookOutputSettings,
-#     HookPredicate,
 #     HookStage,
 # )
+# from mage_ai.data_preparation.models.global_hooks.predicates import (
+#     HookPredicate,
+#     PredicateValueType,
+# )
 # from mage_ai.data_preparation.models.pipeline import Pipeline
 # from mage_ai.data_preparation.models.project.constants import FeatureUUID
 # from mage_ai.data_preparation.repo_manager import get_repo_config
@@ -45,6 +54,62 @@
 #                 except FileNotFoundError as err:
 #                     print(err)
 
+#     def build_predicates(self,
+#         before_values,
+#         after_values,
+#         and_or_operator: PredicateAndOrOperator = None,
+#     ):
+#         predicate_match_before = HookPredicate.load(
+#             and_or_operator=and_or_operator or PredicateAndOrOperator.AND,
+#             predicates=[HookPredicate.load(
+#                 left_object_keys=left_object_keys,
+#                 left_object_type=left_object_type,
+#                 left_value_type=PredicateValueType.load(
+#                     value_data_type=PredicateValueDataType.STRING,
+#                 ),
+#                 operator=PredicateOperator.EQUALS,
+#                 right_value=value,
+#                 right_value_type=PredicateValueType.load(
+#                     value_data_type=PredicateValueDataType.STRING,
+#                 ),
+#             ) for left_object_keys, left_object_type, value in before_values],
+#         )
+
+#         predicate_match_after = HookPredicate.load(
+#             and_or_operator=and_or_operator or PredicateAndOrOperator.AND,
+#             predicates=[HookPredicate.load(
+#                 left_object_keys=left_object_keys,
+#                 left_object_type=left_object_type,
+#                 left_value_type=PredicateValueType.load(
+#                     value_data_type=PredicateValueDataType.STRING,
+#                 ),
+#                 operator=PredicateOperator.EQUALS,
+#                 right_value=value,
+#                 right_value_type=PredicateValueType.load(
+#                     value_data_type=PredicateValueDataType.STRING,
+#                 ),
+#             ) for left_object_keys, left_object_type, value in after_values],
+#         )
+
+#         predicate_miss = HookPredicate.load(
+#             and_or_operator=and_or_operator or PredicateAndOrOperator.AND,
+#             predicates=[
+#                 HookPredicate.load(
+#                     left_value=uuid.uuid4().hex,
+#                     left_value_type=PredicateValueType.load(
+#                         value_data_type=PredicateValueDataType.STRING,
+#                     ),
+#                     operator=PredicateOperator.EQUALS,
+#                     right_value=uuid.uuid4().hex,
+#                     right_value_type=PredicateValueType.load(
+#                         value_data_type=PredicateValueDataType.STRING,
+#                     ),
+#                 ),
+#             ],
+#         )
+
+#         return predicate_match_before, predicate_match_after, predicate_miss
+
 #     async def test_list(self):
 #         await self.setUpAsync(
 #             block_settings={
@@ -54,8 +119,8 @@
 #             },
 #             operation_type=HookOperation.LIST,
 #             pipeline_type=PipelineType.PYTHON,
-#             predicates_match=[],
-#             predicates_miss=[],
+#             predicate_match=None,
+#             predicate_miss=None,
 #         )
 
 #         pipeline3, _blocks = await build_pipeline_with_blocks_and_content(
@@ -97,8 +162,8 @@
 #             },
 #             operation_type=HookOperation.LIST,
 #             pipeline_type=PipelineType.PYTHON,
-#             predicates_match=[],
-#             predicates_miss=[],
+#             predicate_match=None,
+#             predicate_miss=None,
 #             snapshot=False,
 #         )
 
@@ -136,6 +201,17 @@
 #         payload = dict(name=self.faker.unique.name(), type=PipelineType.STREAMING)
 #         name_final = self.faker.unique.name()
 
+#         predicate_match_before, predicate_match_after, predicate_miss = self.build_predicates(
+#             before_values=[
+#                 (['type'], PredicateObjectType.PAYLOAD, payload['type'].value),
+#                 (['name'], PredicateObjectType.PAYLOAD, payload['name']),
+#             ],
+#             after_values=[
+#                 (['type'], PredicateObjectType.RESOURCE, PipelineType.INTEGRATION.value),
+#                 (['name'], PredicateObjectType.RESOURCE, name_final),
+#             ],
+#         )
+
 #         await self.setUpAsync(
 #             block_settings={
 #                 0: dict(content=build_content(dict(type=PipelineType.INTEGRATION))),
@@ -159,25 +235,9 @@
 #                 ),
 #             },
 #             operation_type=HookOperation.CREATE,
-#             predicates_match=[
-#                 [
-#                     HookPredicate.load(resource=dict(
-#                         type=PipelineType.STREAMING.value,
-#                     )),
-#                 ],
-#                 [
-#                     HookPredicate.load(resource=dict(
-#                         name=name_final,
-#                     )),
-#                 ],
-#             ],
-#             predicates_miss=[
-#                 [
-#                     HookPredicate.load(resource=dict(
-#                         name=uuid.uuid4().hex,
-#                     )),
-#                 ],
-#             ],
+#             predicate_match_before=predicate_match_before,
+#             predicate_match_after=predicate_match_after,
+#             predicate_miss=predicate_miss,
 #         )
 
 #         response = await self.build_create_operation(payload=payload).execute()
@@ -198,6 +258,15 @@
 #     async def test_detail(self):
 #         name_final = self.faker.unique.name()
 
+#         predicate_match_before, predicate_match_after, predicate_miss = self.build_predicates(
+#             before_values=[
+#                 (['type'], PredicateObjectType.OPERATION_RESOURCE, PipelineType.PYTHON.value),
+#             ],
+#             after_values=[
+#                 (['type'], PredicateObjectType.OPERATION_RESOURCE, PipelineType.PYTHON.value),
+#             ],
+#         )
+
 #         await self.setUpAsync(
 #             block_settings={
 #                 0: dict(content=build_content(dict())),
@@ -233,25 +302,9 @@
 #             },
 #             operation_type=HookOperation.DETAIL,
 #             pipeline_type=PipelineType.PYTHON.value,
-#             predicates_match=[
-#                 [
-#                     HookPredicate.load(resource=dict(
-#                         type=PipelineType.PYTHON.value,
-#                     )),
-#                 ],
-#                 [
-#                     HookPredicate.load(resource=dict(
-#                         name=PipelineType.INTEGRATION.value,
-#                     )),
-#                 ],
-#             ],
-#             predicates_miss=[
-#                 [
-#                     HookPredicate.load(resource=dict(
-#                         name=uuid.uuid4().hex,
-#                     )),
-#                 ],
-#             ],
+#             predicate_match_before=predicate_match_before,
+#             predicate_match_after=predicate_match_after,
+#             predicate_miss=predicate_miss,
 #         )
 
 #         response = await self.build_detail_operation(self.pipeline2.uuid).execute()
@@ -273,6 +326,15 @@
 
 #         payload = dict(description=description_init, tags=tags_init)
 
+#         predicate_match_before, predicate_match_after, predicate_miss = self.build_predicates(
+#             before_values=[
+#                 (['type'], PredicateObjectType.OPERATION_RESOURCE, PipelineType.PYTHON.value),
+#             ],
+#             after_values=[
+#                 (['type'], PredicateObjectType.OPERATION_RESOURCE, PipelineType.PYTHON.value),
+#             ],
+#         )
+
 #         await self.setUpAsync(
 #             block_settings={
 #                 0: dict(content=build_content(dict(
@@ -333,20 +395,9 @@
 #             },
 #             operation_type=HookOperation.UPDATE,
 #             pipeline_type=PipelineType.PYTHON.value,
-#             predicates_match=[
-#                 [
-#                     HookPredicate.load(resource=dict(
-#                         type=PipelineType.PYTHON.value,
-#                     )),
-#                 ],
-#             ],
-#             predicates_miss=[
-#                 [
-#                     HookPredicate.load(resource=dict(
-#                         name=uuid.uuid4().hex,
-#                     )),
-#                 ],
-#             ],
+#             predicate_match_before=predicate_match_before,
+#             predicate_match_after=predicate_match_after,
+#             predicate_miss=predicate_miss,
 #         )
 
 #         response = await self.build_update_operation(
@@ -367,6 +418,15 @@
 #         name_final = uuid.uuid4().hex
 #         uuid_final = uuid.uuid4().hex
 
+#         predicate_match_before, predicate_match_after, predicate_miss = self.build_predicates(
+#             before_values=[
+#                 (['type'], PredicateObjectType.OPERATION_RESOURCE, PipelineType.PYTHON.value),
+#             ],
+#             after_values=[
+#                 (['type'], PredicateObjectType.OPERATION_RESOURCE, PipelineType.PYTHON.value),
+#             ],
+#         )
+
 #         await self.setUpAsync(
 #             block_settings={
 #                 0: dict(content=build_content(dict(
@@ -431,20 +491,9 @@
 #             },
 #             operation_type=HookOperation.DELETE,
 #             pipeline_type=PipelineType.PYTHON.value,
-#             predicates_match=[
-#                 [
-#                     HookPredicate.load(resource=dict(
-#                         type=PipelineType.PYTHON.value,
-#                     )),
-#                 ],
-#             ],
-#             predicates_miss=[
-#                 [
-#                     HookPredicate.load(resource=dict(
-#                         name=uuid.uuid4().hex,
-#                     )),
-#                 ],
-#             ],
+#             predicate_match_before=predicate_match_before,
+#             predicate_match_after=predicate_match_after,
+#             predicate_miss=predicate_miss,
 #         )
 
 #         response = await self.build_delete_operation(self.pipeline2.uuid).execute()
@@ -468,6 +517,18 @@
 #         color = uuid.uuid4().hex
 #         configuration = dict(power=uuid.uuid4().hex)
 
+#         predicate_match_before, predicate_match_after, predicate_miss = self.build_predicates(
+#             before_values=[
+#                 (['type'], PredicateObjectType.OPERATION_RESOURCE, BlockType.DATA_LOADER.value),
+#                 (['type'], PredicateObjectType.OPERATION_RESOURCE, BlockType.TRANSFORMER.value),
+#             ],
+#             after_values=[
+#                 (['type'], PredicateObjectType.OPERATION_RESOURCE, BlockType.DATA_LOADER.value),
+#                 (['type'], PredicateObjectType.OPERATION_RESOURCE, BlockType.TRANSFORMER.value),
+#             ],
+#             and_or_operator=PredicateAndOrOperator.OR,
+#         )
+
 #         await self.setUpAsync(
 #             block_settings={
 #                 0: dict(content=build_content(dict(
@@ -519,25 +580,9 @@
 #             },
 #             operation_type=HookOperation.UPDATE_ANYWHERE,
 #             pipeline_type=PipelineType.PYTHON.value,
-#             predicates_match=[
-#                 [
-#                     HookPredicate.load(resource=dict(
-#                         type=BlockType.DATA_LOADER.value,
-#                     )),
-#                 ],
-#                 [
-#                     HookPredicate.load(resource=dict(
-#                         type=BlockType.TRANSFORMER.value,
-#                     )),
-#                 ],
-#             ],
-#             predicates_miss=[
-#                 [
-#                     HookPredicate.load(resource=dict(
-#                         type=BlockType.MARKDOWN.value,
-#                     )),
-#                 ],
-#             ],
+#             predicate_match_before=predicate_match_before,
+#             predicate_match_after=predicate_match_after,
+#             predicate_miss=predicate_miss,
 #             resource_type=EntityName.Block,
 #         )
 

mage_ai/tests/api/policies/test_oauth_policy.py

@@ -0,0 +1,38 @@
+from mage_ai.api.constants import AttributeOperationType
+from mage_ai.api.policies.OauthPolicy import OauthPolicy
+from mage_ai.api.presenters.OauthPresenter import OauthPresenter
+from mage_ai.tests.base_test import AsyncDBTestCase
+
+
+class OauthPolicyTestCase(AsyncDBTestCase):
+    def test_attribute_rule_with_permissions(self):
+        error = False
+
+        try:
+            write_attributes = [
+                'action_type',
+                'provider',
+                'token',
+            ]
+            for key in write_attributes:
+                OauthPolicy.attribute_rule_with_permissions(
+                    AttributeOperationType.WRITE,
+                    key,
+                )
+
+            read_attributes = OauthPresenter.default_attributes
+            for key in read_attributes:
+                OauthPolicy.attribute_rule_with_permissions(
+                    AttributeOperationType.READ,
+                    key,
+                )
+
+            for key in ['doesn’t exist', 'mage']:
+                OauthPolicy.attribute_rule_with_permissions(
+                    AttributeOperationType.QUERY,
+                    key,
+                )
+        except Exception:
+            error = True
+
+        self.assertFalse(error)

mage_ai/tests/data_preparation/models/global_hooks/test_global_hooks.py

@@ -15,11 +15,11 @@ from mage_ai.data_preparation.models.global_hooks.models import (
     Hook,
     HookCondition,
     HookOperation,
-    HookPredicate,
     HookRunSettings,
     HookStage,
     HookStrategy,
 )
+from mage_ai.data_preparation.models.global_hooks.predicates import HookPredicate
 from mage_ai.shared.array import find
 from mage_ai.shared.io import safe_write
 from mage_ai.tests.api.operations.test_base import BaseApiTestCase
@@ -53,7 +53,7 @@ def build_hook(
     operation_type: HookOperation = None,
     output: Dict = None,
     pipeline: Dict = None,
-    predicates: List[List[HookPredicate]] = None,
+    predicate: HookPredicate = None,
     resource_type: EntityName = None,
     run_settings: HookRunSettings = None,
     stages: List[HookStage] = None,
@@ -65,7 +65,7 @@ def build_hook(
         operation_type=operation_type or HookOperation.DETAIL,
         output=output,
         pipeline=pipeline,
-        predicates=predicates,
+        predicate=predicate,
         resource_type=resource_type or EntityName.Pipeline,
         run_settings=run_settings,
         stages=stages or [HookStage.BEFORE],
@@ -491,19 +491,44 @@ class GlobalHooksTest(BaseApiTestCase):
                     stage=HookStage.BEFORE,
                     conditions=[HookCondition.SUCCESS, HookCondition.FAILURE],
                     operation_resource=dict(mage=1),
+                    resource_id=1,
+                    resource_parent_id=3,
+                    user=dict(id=7),
                 )
 
                 self.global_hooks.get_and_run_hooks(
-                    fire=2,
-                    water=3,
                     **options,
                 )
 
-                mock_get_hooks.assert_called_once_with(**options)
+                mock_get_hooks.assert_called_once_with(
+                    [HookOperation.LIST, HookOperation.CREATE],
+                    EntityName.Tag,
+                    HookStage.BEFORE,
+                    conditions=[HookCondition.SUCCESS, HookCondition.FAILURE],
+                    error=None,
+                    meta=None,
+                    metadata=None,
+                    operation_resource=dict(mage=1),
+                    payload=None,
+                    query=None,
+                    resource=None,
+                    resource_id=1,
+                    resource_parent_id=3,
+                    resources=None,
+                    user=dict(id=7),
+                )
                 mock_run_hooks.assert_called_once_with(
                     hooks,
-                    fire=2,
-                    water=3,
+                    error=None,
+                    meta=None,
+                    metadata=None,
+                    payload=None,
+                    query=None,
+                    resource=None,
+                    resource_id=1,
+                    resource_parent_id=3,
+                    resources=None,
+                    user=dict(id=7),
                 )
 
     def test_save(self):