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

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

@@ -28,6 +28,11 @@ class GlobalHookUtilsTest(AsyncDBTestCase):
             extracted_data = extract_valid_data(entity_name, data)
 
             if entity_name in types:
-                self.assertEqual(extracted_data, dict(hook=data['hook']))
+                self.assertEqual(extracted_data, dict(
+                    hook=data['hook'],
+                    resource_id=data['resource_id'],
+                    resource_parent_id=data['resource_parent_id'],
+                    user=data['user'],
+                ))
             else:
                 self.assertEqual(extracted_data, ignore_keys(data, ['mage']))

mage_ai/tests/data_preparation/models/test_block.py

@@ -13,6 +13,7 @@ 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.tests.base_test import DBTestCase
+from mage_ai.tests.factory import create_integration_pipeline_with_blocks
 
 
 class BlockTest(DBTestCase):
@@ -707,3 +708,28 @@ def test_output(output, *args) -> None:
             block3.upstream_block_uuids[1],
             'test_data_loader_1',
         )
+
+    def test_output_variables_for_integration_pipeline_blocks(self):
+        pipeline = create_integration_pipeline_with_blocks(
+            'test integration pipeline',
+            repo_path=self.repo_path,
+        )
+        block = pipeline.get_block('test_integration_source')
+
+        df = pd.DataFrame(
+            [
+                [1, 'abc@xyz.com', '32132'],
+                [2, 'abc2@xyz.com', '12345'],
+                [3, 'test', '1234'],
+                [4, 'abc@test.net', 'abcde'],
+                [5, 'abc12345@', '54321'],
+                [6, 'abcdef@123.com', '56789'],
+            ],
+            columns=['id', 'email', 'zip_code'],
+        )
+
+        block.store_variables({
+            'output_sample_data_stream1': df
+        })
+        output_variables = block.output_variables()
+        self.assertTrue('output_sample_data_stream1' in output_variables)

mage_ai/tests/data_preparation/models/test_pipeline.py

@@ -35,6 +35,7 @@ class PipelineTest(DBTestCase):
         pipeline = Pipeline('test_pipeline_2', self.repo_path)
 
         self.assertEqual(pipeline.to_dict(), dict(
+            cache_block_output_in_memory=False,
             concurrency_config=dict(),
             data_integration=None,
             description=None,
@@ -49,6 +50,8 @@ class PipelineTest(DBTestCase):
             run_pipeline_in_one_process=False,
             spark_config=dict(),
             type='python',
+            remote_variables_dir=None,
+            variables_dir=self.repo_path,
             blocks=[
                 dict(
                     language='python',
@@ -198,6 +201,7 @@ class PipelineTest(DBTestCase):
         pipeline.delete_block(block)
         pipeline = Pipeline('test_pipeline_3', self.repo_path)
         self.assertEqual(pipeline.to_dict(), dict(
+            cache_block_output_in_memory=False,
             concurrency_config=dict(),
             data_integration=None,
             description=None,
@@ -212,6 +216,8 @@ class PipelineTest(DBTestCase):
             run_pipeline_in_one_process=False,
             spark_config=dict(),
             type='python',
+            remote_variables_dir=None,
+            variables_dir=self.repo_path,
             blocks=[
                 dict(
                     language='python',
@@ -289,6 +295,7 @@ class PipelineTest(DBTestCase):
         pipeline.add_block(block4, upstream_block_uuids=['block2', 'block3'])
         pipeline.execute_sync()
         self.assertEqual(pipeline.to_dict(), dict(
+            cache_block_output_in_memory=False,
             concurrency_config=dict(),
             data_integration=None,
             description=None,
@@ -303,6 +310,8 @@ class PipelineTest(DBTestCase):
             run_pipeline_in_one_process=False,
             spark_config=dict(),
             type='python',
+            remote_variables_dir=None,
+            variables_dir=self.repo_path,
             blocks=[
                 dict(
                     language='python',
@@ -403,6 +412,7 @@ class PipelineTest(DBTestCase):
         pipeline.add_block(block7, upstream_block_uuids=['block2', 'block3', 'block6'])
         pipeline.execute_sync()
         self.assertEqual(pipeline.to_dict(), dict(
+            cache_block_output_in_memory=False,
             concurrency_config=dict(),
             data_integration=None,
             description=None,
@@ -417,6 +427,8 @@ class PipelineTest(DBTestCase):
             run_pipeline_in_one_process=False,
             spark_config=dict(),
             type='python',
+            remote_variables_dir=None,
+            variables_dir=self.repo_path,
             blocks=[
                 dict(
                     language='python',
@@ -653,6 +665,7 @@ class PipelineTest(DBTestCase):
             self.assertEqual(
                 config_json,
                 dict(
+                    cache_block_output_in_memory=False,
                     concurrency_config=dict(),
                     created_at='2023-08-01 08:08:24+00:00',
                     data_integration=None,
@@ -670,6 +683,8 @@ class PipelineTest(DBTestCase):
                     type='integration',
                     updated_at=None,
                     uuid='test_pipeline_9',
+                    remote_variables_dir=None,
+                    variables_dir=self.repo_path,
                     blocks=[
                         dict(
                             all_upstream_blocks_executed=True,

mage_ai/tests/factory.py

@@ -1,5 +1,5 @@
 import os
-from datetime import datetime
+from datetime import datetime, timedelta
 from typing import Dict, Union
 
 from faker import Faker
@@ -10,7 +10,11 @@ from mage_ai.data_preparation.models.constants import BlockType, PipelineType
 from mage_ai.data_preparation.models.pipeline import Pipeline
 from mage_ai.data_preparation.models.triggers import ScheduleType
 from mage_ai.orchestration.db.models.oauth import User
-from mage_ai.orchestration.db.models.schedules import PipelineRun, PipelineSchedule
+from mage_ai.orchestration.db.models.schedules import (
+    Backfill,
+    PipelineRun,
+    PipelineSchedule,
+)
 from mage_ai.shared.hash import extract, ignore_keys, merge_dict
 
 faker = Faker()
@@ -159,6 +163,40 @@ def create_pipeline_run_with_schedule(
     return pipeline_run
 
 
+def create_backfill(
+    pipeline_uuid: str,
+    end_datetime: datetime = None,
+    interval_type: Backfill.IntervalType = Backfill.IntervalType.DAY,
+    interval_units: int = 1,
+    pipeline_schedule_id: int = None,
+    start_datetime: datetime = None,
+    **kwargs,
+):
+    if start_datetime is None:
+        start_datetime = datetime.today() - timedelta(days=1)
+    if end_datetime is None:
+        end_datetime = datetime.today()
+    if pipeline_schedule_id is None:
+        pipeline_schedule = PipelineSchedule.create(
+            name=f'{pipeline_uuid}_trigger',
+            pipeline_uuid=pipeline_uuid,
+            schedule_type=ScheduleType.TIME,
+        )
+        pipeline_schedule_id = pipeline_schedule.id
+    backfill = Backfill.create(
+        end_datetime=end_datetime,
+        interval_type=interval_type,
+        interval_units=interval_units,
+        pipeline_uuid=pipeline_uuid,
+        pipeline_schedule_id=pipeline_schedule_id,
+        start_datetime=start_datetime,
+        status=Backfill.Status.INITIAL,
+        **kwargs,
+    )
+
+    return backfill
+
+
 def create_user(
     as_dict: bool = False,
     save: bool = True,

mage_ai/tests/orchestration/test_pipeline_scheduler.py

@@ -14,12 +14,14 @@ from mage_ai.data_preparation.models.triggers import (
 )
 from mage_ai.data_preparation.preferences import get_preferences
 from mage_ai.data_preparation.variable_manager import VariableManager
+from mage_ai.orchestration.backfills.service import start_backfill
 from mage_ai.orchestration.db.models.schedules import (
+    Backfill,
     BlockRun,
     PipelineRun,
     PipelineSchedule,
 )
-from mage_ai.orchestration.job_manager import JobType
+from mage_ai.orchestration.job_manager import JobManager, JobType
 from mage_ai.orchestration.notification.sender import NotificationSender
 from mage_ai.orchestration.pipeline_scheduler import (
     PipelineScheduler,
@@ -30,6 +32,7 @@ from mage_ai.shared.array import find
 from mage_ai.shared.hash import ignore_keys, merge_dict
 from mage_ai.tests.base_test import DBTestCase
 from mage_ai.tests.factory import (
+    create_backfill,
     create_integration_pipeline_with_blocks,
     create_pipeline_run_with_schedule,
     create_pipeline_with_blocks,
@@ -141,6 +144,84 @@ class PipelineSchedulerTests(DBTestCase):
                 pipeline_run=pipeline_run,
             )
 
+    @freeze_time('2023-11-11 12:30:00')
+    def test_backfill_status_with_completed_pipeline_run(self):
+        pipeline_schedule = PipelineSchedule.create(
+            name='test_pipeline_trigger',
+            pipeline_uuid='test_pipeline',
+        )
+        backfill = create_backfill(
+            'test_pipeline',
+            end_datetime=datetime.today(),
+            pipeline_schedule_id=pipeline_schedule.id,
+            start_datetime=datetime.today(),
+        )
+        pipeline_runs = start_backfill(backfill)
+        pipeline_run = pipeline_runs[0]
+        pipeline_run.update(status=PipelineRun.PipelineRunStatus.RUNNING)
+        scheduler = PipelineScheduler(pipeline_run=pipeline_run)
+        with patch.object(JobManager, 'add_job') as mock_add_br_job:
+            scheduler.schedule()
+            self.assertEqual(
+                backfill.status,
+                Backfill.Status.RUNNING,
+            )
+            mock_add_br_job.assert_called()
+
+            for b in pipeline_run.block_runs:
+                b.update(status=BlockRun.BlockRunStatus.COMPLETED)
+            scheduler.schedule()
+            self.assertEqual(
+                backfill.status,
+                Backfill.Status.COMPLETED,
+            )
+
+    @freeze_time('2023-11-11 12:30:00')
+    def test_backfill_status_with_failed_pipeline_run(self):
+        pipeline_schedule = PipelineSchedule.create(
+            name='test_pipeline_trigger',
+            pipeline_uuid='test_pipeline',
+        )
+        backfill = create_backfill(
+            'test_pipeline',
+            end_datetime=datetime.today(),
+            pipeline_schedule_id=pipeline_schedule.id,
+            start_datetime=datetime.today(),
+        )
+        pipeline_runs = start_backfill(backfill)
+        pipeline_run = pipeline_runs[0]
+        scheduler = PipelineScheduler(pipeline_run=pipeline_run)
+
+        for b in pipeline_run.block_runs:
+            b.update(status=BlockRun.BlockRunStatus.FAILED)
+        scheduler.schedule()
+        # Backfill with 1 pipeline run that fails
+        self.assertEqual(
+            backfill.status,
+            Backfill.Status.FAILED,
+        )
+
+        # Retry failed pipeline run that was associated with the existing backfill
+        retried_pipeline_run = create_pipeline_run_with_schedule(
+            'test_pipeline',
+            execution_date=pipeline_run.execution_date,
+            pipeline_schedule_id=pipeline_schedule.id,
+        )
+        scheduler2 = PipelineScheduler(pipeline_run=retried_pipeline_run)
+        backfill.update(status=Backfill.Status.INITIAL)
+        self.assertEqual(
+            backfill.status,
+            Backfill.Status.INITIAL,
+        )
+        # Retried pipeline run successfully completes, so backfill status updates to "completed"
+        for b in retried_pipeline_run.block_runs:
+            b.update(status=BlockRun.BlockRunStatus.COMPLETED)
+        scheduler2.schedule()
+        self.assertEqual(
+            backfill.status,
+            Backfill.Status.COMPLETED,
+        )
+
     @freeze_time('2023-10-11 12:13:14')
     def test_schedule_all_with_integration_pipeline(self):
         integration_pipeline_schedule = PipelineSchedule.create(

mage_ai/tests/services/k8s/test_job_manager.py

@@ -41,6 +41,14 @@ MOCK_POD_CONFIG = V1Pod(
         )],
         restart_policy='Never',
         service_account_name='old_service_account_name',
+        tolerations=[
+            {
+                'effect': 'NoExecute',
+                'key': 'node.kubernetes.io/not-ready',
+                'operator': 'Exists',
+                'toleration_seconds': 300,
+            },
+        ],
         volumes=[],
     )
 )
@@ -188,6 +196,14 @@ class JobManagerTests(TestCase):
                     secret_ref=V1SecretEnvSource(name='mysecret', optional=None),
                 ),
         ])
+        self.assertEqual(pod_config.tolerations, [
+            {
+                'effect': 'NoExecute',
+                'key': 'node.kubernetes.io/not-ready',
+                'operator': 'Exists',
+                'toleration_seconds': 300,
+            },
+        ])
 
         mock_client.V1ResourceRequirements.assert_called_once_with(
             limits=dict(

mage_ai/tests/shared/mixins.py

@@ -1,11 +1,15 @@
 import json
 import os
-import uuid
 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
+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,
@@ -13,10 +17,13 @@ from mage_ai.data_preparation.models.global_hooks.models import (
     HookOperation,
     HookOutputBlock,
     HookOutputSettings,
-    HookPredicate,
     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
@@ -62,8 +69,10 @@ def build_hooks(
     resource_type: EntityName = None,
     test_predicate_after: bool = False,
     test_predicate_before: bool = False,
-    predicates_match: List[HookPredicate] = None,
-    predicates_miss: List[HookPredicate] = None,
+    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:
@@ -75,19 +84,37 @@ def build_hooks(
             type=pipeline.type,
         )
 
-    if predicates_match is None:
-        predicates_match = [[
-            HookPredicate.load(resource=matching_predicate_resource),
-        ]]
+    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 predicates_miss is None:
-        predicates_miss = [[
-            HookPredicate.load(resource=dict(
-                name=uuid.uuid4().hex,
-                type=pipeline.type,
-            )),
-            predicates_match[0][0],
-        ]]
+    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
@@ -120,9 +147,15 @@ def build_hooks(
             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,
-            predicates=predicates_match + predicates_miss,
+            predicate=predicate_use,
             resource_type=resource_type,
             strategies=[HookStrategy.RAISE],
             uuid=f'hook{idx}_{test_case.faker.unique.name()}',
@@ -135,7 +168,7 @@ def build_hooks(
             **hook.to_dict(),
         )
         hook_miss.uuid = test_case.faker.unique.name()
-        hook_miss.predicates = predicates_miss
+        hook_miss.predicate = predicate_miss
 
         if snapshot:
             hook.snapshot()
@@ -165,8 +198,10 @@ class GlobalHooksMixin(BaseApiTestCase):
         matching_predicate_resource: Dict = None,
         operation_type: HookOperation = None,
         pipeline_type: PipelineType = None,
-        predicates_match: List[HookPredicate] = None,
-        predicates_miss: List[HookPredicate] = 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,
     ):
@@ -259,8 +294,10 @@ class GlobalHooksMixin(BaseApiTestCase):
             matching_predicate_resource=matching_predicate_resource,
             operation_type=operation_type,
             pipeline=pipeline1,
-            predicates_match=predicates_match,
-            predicates_miss=predicates_miss,
+            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,
         )

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-0.9.46.dist-info → mage_ai-0.9.48.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: mage-ai
-Version: 0.9.46
+Version: 0.9.48
 Summary: Mage is a tool for building and deploying data pipelines.
 Home-page: https://github.com/mage-ai/mage-ai
 Author: Mage
@@ -56,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'
@@ -74,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: chromadb >=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'
@@ -100,14 +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'
@@ -134,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: chromadb >=0.4.17 ; extra == 'chroma'
 Provides-Extra: clickhouse
 Requires-Dist: clickhouse-connect ==0.5.20 ; extra == 'clickhouse'
 Provides-Extra: dbt