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

mage_ai/server/server.py

@@ -95,7 +95,13 @@ from mage_ai.settings import (
     SHELL_COMMAND,
     USE_UNIQUE_TERMINAL,
 )
-from mage_ai.settings.repo import DEFAULT_MAGE_DATA_DIR, get_repo_name, set_repo_path
+from mage_ai.settings.repo import (
+    DEFAULT_MAGE_DATA_DIR,
+    MAGE_CLUSTER_TYPE_ENV_VAR,
+    MAGE_PROJECT_TYPE_ENV_VAR,
+    get_repo_name,
+    set_repo_path,
+)
 from mage_ai.shared.constants import ENV_VAR_INSTANCE_TYPE, InstanceType
 from mage_ai.shared.io import chmod
 from mage_ai.shared.logger import LoggingLevel
@@ -328,6 +334,8 @@ def make_app(template_dir: str = None, update_routes: bool = False):
         (r'/files', MainPageHandler),
         (r'/global-data-products', MainPageHandler),
         (r'/global-data-products/(?P<uuid>\w+)', MainPageHandler),
+        (r'/global-hooks', MainPageHandler),
+        (r'/global-hooks/(?P<uuid>\w+)', MainPageHandler),
         (r'/templates', MainPageHandler),
         (r'/templates/(?P<uuid>\w+)', MainPageHandler),
         (r'/version-control', MainPageHandler),
@@ -648,8 +656,8 @@ if __name__ == '__main__':
     manage = args.manage_instance == '1'
     dbt_docs = args.dbt_docs_instance == '1'
     instance_type = os.getenv(ENV_VAR_INSTANCE_TYPE, args.instance_type)
-    project_type = os.getenv('PROJECT_TYPE', ProjectType.STANDALONE)
-    cluster_type = os.getenv('CLUSTER_TYPE')
+    project_type = os.getenv(MAGE_PROJECT_TYPE_ENV_VAR, ProjectType.STANDALONE)
+    cluster_type = os.getenv(MAGE_CLUSTER_TYPE_ENV_VAR)
 
     start_server(
         host=host,

mage_ai/services/k8s/job_manager.py

@@ -154,6 +154,7 @@ class JobManager():
         pod_spec.containers = [container]
         mage_server_pod_spec = self.pod_config.spec
         pod_spec.volumes.extend(mage_server_pod_spec.volumes)
+        pod_spec.tolerations = mage_server_pod_spec.tolerations
         pod_spec.node_selector = mage_server_pod_spec.node_selector
         pod_spec.image_pull_secrets = pod_spec.image_pull_secrets if pod_spec.image_pull_secrets \
             else mage_server_pod_spec.image_pull_secrets

mage_ai/settings/repo.py

@@ -3,6 +3,9 @@ import sys
 
 from mage_ai.shared.environments import is_test
 
+MAGE_PROJECT_TYPE_ENV_VAR = 'PROJECT_TYPE'
+MAGE_CLUSTER_TYPE_ENV_VAR = 'CLUSTER_TYPE'
+
 """
 Moved from repo_manager because repo_manager has too many dependencies
 which can cause circular import errors.

mage_ai/shared/files.py

@@ -0,0 +1,47 @@
+import os
+
+
+def reverse_readline(filename, buf_size=8192):
+    """A generator that returns the lines of a file in reverse order"""
+    with open(filename, 'rb') as fh:
+        segment = None
+        offset = 0
+        fh.seek(0, os.SEEK_END)
+        file_size = remaining_size = fh.tell()
+        while remaining_size > 0:
+            offset = min(file_size, offset + buf_size)
+            fh.seek(file_size - offset)
+            buffer = fh.read(min(remaining_size, buf_size)).decode(encoding='utf-8')
+            remaining_size -= buf_size
+            lines = buffer.split('\n')
+            # The first line of the buffer is probably not a complete line so
+            # we'll save it and append it to the last line of the next buffer
+            # we read
+            if segment is not None:
+                # If the previous chunk starts right from the beginning of line
+                # do not concat the segment to the last line of new chunk.
+                # Instead, yield the segment first
+                if buffer[-1] != '\n':
+                    lines[-1] += segment
+                else:
+                    yield segment
+            segment = lines[0]
+            for index in range(len(lines) - 1, 0, -1):
+                if lines[index]:
+                    yield lines[index]
+        # Don't yield None if the file was empty
+        if segment is not None:
+            yield segment
+
+
+def read_last_line(filename: str) -> str:
+    with open(filename, 'rb') as f:
+        try:  # catch OSError in case of a one line file
+            f.seek(-2, os.SEEK_END)
+            while f.read(1) != b'\n':
+                f.seek(-2, os.SEEK_CUR)
+        except OSError:
+            f.seek(0)
+        last_line = f.readline().decode()
+
+        return last_line

mage_ai/shared/models.py

@@ -59,6 +59,7 @@ class BaseClass:
                     model.set_value(key, value)
             except AttributeError as err:
                 print(f'[WARNING] {self.__name__}.load: {err}')
+                raise err
 
         return model
 

mage_ai/streaming/constants.py

@@ -37,6 +37,7 @@ class SinkType(str, Enum):
     OPENSEARCH = 'opensearch'
     ORACLEDB = 'oracledb'
     POSTGRES = 'postgres'
+    RABBITMQ = 'rabbitmq'
     REDSHIFT = 'redshift'
     SNOWFLAKE = 'snowflake'
     TRINO = 'trino'

mage_ai/streaming/sinks/postgres.py

@@ -44,6 +44,8 @@ class PostgresSink(BaseSink):
         self.batch_write([message])
 
     def batch_write(self, messages: List[Dict]):
+        if not messages:
+            return
         self._print(
             f'Batch ingest {len(messages)} records, time={time.time()}. Sample: {messages[0]}')
 

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/streaming/sources/nats_js.py

@@ -91,8 +91,19 @@ class NATSSource(BaseSource):
             # Default consumer_name to stream_name if not provided
             consumer_name = self.config.consumer_name or self.config.stream_name
 
-            await self.js.add_stream(name=self.config.stream_name, subjects=[self.config.subject])
+            # Check if the stream exists, and create it if it doesn't
+            try:
+                await self.js.stream_info(self.config.stream_name)
+            except Exception as e:
+                # Check the exception type or message to ensure it's about a missing stream
+                if 'stream not found' in str(e).lower():
+                    await self.js.add_stream(
+                        name=self.config.stream_name,
+                        subjects=[self.config.subject],
+                    )
+
             self.psub = await self.js.pull_subscribe(self.config.subject, consumer_name)
+
         except NoServersError as e:
             self._print(f'Caught NoServersError while connecting to NATS server: {e}')
         except Exception as e:

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/executors/test_block_executor.py

@@ -99,6 +99,8 @@ class BlockExecutorTest(TestCase):
         self.block_executor._execute.assert_called_once_with(
             analyze_outputs=analyze_outputs,
             block_run_id=None,
+            block_run_outputs_cache=None,
+            cache_block_output_in_memory=False,
             callback_url=callback_url,
             global_vars=global_vars,
             input_from_output=input_from_output,