mage-ai 0.9.70__py3-none-any.whl → 0.9.71__py3-none-any.whl

mage_ai/tests/orchestration/db/models/test_schedules_dynamic_blocks.py

@@ -0,0 +1,279 @@
+from unittest.mock import patch
+
+from mage_ai.data_preparation.models.block.settings.dynamic.constants import ModeType
+from mage_ai.orchestration.db.models.schedules import BlockRun
+from mage_ai.tests.data_preparation.models.test_blocks_helper import (
+    BlockHelperTest,
+    load_dataframe,
+    load_dynamic_block_child_data,
+    passthrough,
+)
+from mage_ai.tests.factory import create_pipeline_run
+
+
+class PipelineRunWithDynamicBlocksTests(BlockHelperTest):
+    def setUp(self):
+        super().setUp()
+
+    def tearDown(self):
+        BlockRun.query.delete()
+        self.pipeline.delete()
+        super().tearDown()
+
+    @patch.multiple('mage_ai.settings.server', DYNAMIC_BLOCKS_V2=True)
+    def test_executable_block_runs(self):
+        """
+        Combinations of upstream blocks:
+            - Basic
+            - Dynamic
+            - Dynamic child
+            - Dynamic duo
+            - Reduce output
+
+        Base structure: 14 blocks
+        block1
+        |   dynamic1
+        |   |---dynamic1_child1a
+        |   |   |---dynamic1_child1a_reduce1a
+        |   |---|---dynamic1_child1a_dynamic
+        |   |   |   |---end1
+        |---|---dynamic1_child2a
+                |---dynamic1_child1a_child2a_reduce2a (also depends on dynamic1_child1a)
+                    |---end2 (also depends on dynamic1_child1a_reduce1a)
+
+        block2
+        |---end3
+        |---end4 (also depends on block1)
+        |---dynamic2
+            | |-dynamic2_child1b
+            | | |---dynamic2_child1b_reduce1b
+            |-------dynamic2_child1b_dynamic
+            |---dynamic2_child2b
+            | |-|---dynamic2_child1b_child2b_reduce2b
+        """
+
+        # Basic blocks
+        block1 = self.create_block('block1', func=load_dataframe)
+        self.pipeline.add_block(block1)
+        block2 = self.create_block('block2', func=load_dataframe)
+        self.pipeline.add_block(block2)
+
+        # Dynamic blocks
+        dynamic1 = self.create_block(
+            'dynamic1',
+            func=load_dynamic_block_child_data,
+            configuration=dict(
+                dynamic=dict(parent=True, modes=[dict(type=ModeType.STREAM.value)]),
+            ),
+        )
+        self.pipeline.add_block(dynamic1)
+
+        dynamic2 = self.create_block(
+            'dynamic2',
+            func=load_dynamic_block_child_data,
+            configuration=dict(
+                dynamic=dict(parent=True, modes=[dict(type=ModeType.STREAM.value)]),
+            ),
+        )
+        self.pipeline.add_block(dynamic2, upstream_block_uuids=[block2.uuid])
+
+        # Dynamic children
+        dynamic1_child1a = self.create_block('dynamic1_child1a', func=passthrough)
+        self.pipeline.add_block(dynamic1_child1a, upstream_block_uuids=[dynamic1.uuid])
+
+        dynamic1_child2a = self.create_block('dynamic1_child2a', func=passthrough)
+        self.pipeline.add_block(
+            dynamic1_child2a, upstream_block_uuids=[dynamic1.uuid, block1.uuid]
+        )
+
+        dynamic2_child1b = self.create_block('dynamic2_child1b', func=passthrough)
+        self.pipeline.add_block(dynamic2_child1b, upstream_block_uuids=[dynamic2.uuid])
+
+        dynamic2_child2b = self.create_block('dynamic2_child2b', func=passthrough)
+        self.pipeline.add_block(dynamic2_child2b, upstream_block_uuids=[dynamic2.uuid])
+
+        # Dynamic duo
+        dynamic1_child1a_dynamic = self.create_block('dynamic1_child1a_dynamic', func=passthrough)
+        self.pipeline.add_block(
+            dynamic1_child1a_dynamic, upstream_block_uuids=[dynamic1_child1a.uuid, dynamic1.uuid]
+        )
+
+        dynamic2_child1b_dynamic = self.create_block('dynamic2_child1b_dynamic', func=passthrough)
+        self.pipeline.add_block(
+            dynamic2_child1b_dynamic, upstream_block_uuids=[dynamic2_child1b.uuid, dynamic2.uuid]
+        )
+
+        # Reduce outputs
+        # # Reduce output from dynamic1 children
+        dynamic1_child1a_reduce1a = self.create_block(
+            'dynamic1_child1a_reduce1a',
+            func=passthrough,
+            configuration=dict(dynamic=dict(reduce_output=True)),
+        )
+        self.pipeline.add_block(
+            dynamic1_child1a_reduce1a, upstream_block_uuids=[dynamic1_child1a.uuid]
+        )
+
+        dynamic1_child1a_child2a_reduce2a = self.create_block(
+            'dynamic1_child1a_child2a_reduce2a',
+            func=passthrough,
+            configuration=dict(dynamic=dict(reduce_output=True)),
+        )
+        self.pipeline.add_block(
+            dynamic1_child1a_child2a_reduce2a,
+            upstream_block_uuids=[dynamic1_child1a.uuid, dynamic1_child2a.uuid],
+        )
+        # Reduce output from dynamic2 children
+        dynamic2_child1b_reduce1b = self.create_block(
+            'dynamic2_child1b_reduce1b',
+            func=passthrough,
+            configuration=dict(dynamic=dict(reduce_output=True)),
+        )
+        self.pipeline.add_block(
+            dynamic2_child1b_reduce1b, upstream_block_uuids=[dynamic2_child1b.uuid]
+        )
+
+        dynamic2_child1b_child2b_reduce2b = self.create_block(
+            'dynamic2_child1b_child2b_reduce2b',
+            func=passthrough,
+            configuration=dict(dynamic=dict(reduce_output=True)),
+        )
+        self.pipeline.add_block(
+            dynamic2_child1b_child2b_reduce2b,
+            upstream_block_uuids=[dynamic2_child1b.uuid, dynamic2_child2b.uuid],
+        )
+
+        # Final blocks
+        end1 = self.create_block('end1', func=passthrough)
+        self.pipeline.add_block(end1, upstream_block_uuids=[dynamic1_child1a_reduce1a.uuid])
+
+        end2 = self.create_block('end2', func=passthrough)
+        self.pipeline.add_block(
+            end2,
+            upstream_block_uuids=[
+                dynamic1_child1a_reduce1a.uuid,
+                dynamic1_child1a_child2a_reduce2a.uuid,
+            ],
+        )
+
+        end3 = self.create_block('end3', func=passthrough)
+        self.pipeline.add_block(end3, upstream_block_uuids=[block2.uuid])
+
+        end4 = self.create_block('end4', func=passthrough)
+        self.pipeline.add_block(end4, upstream_block_uuids=[block1.uuid, block2.uuid])
+
+        pipeline_run = create_pipeline_run(pipeline_uuid=self.pipeline.uuid)
+
+        def get_run(block, block_runs=pipeline_run.block_runs):
+            return next(br for br in block_runs if br.block_uuid == block.uuid)
+
+        def complete(blocks):
+            updates(blocks, BlockRun.BlockRunStatus.COMPLETED)
+
+        def running(blocks):
+            updates(blocks, BlockRun.BlockRunStatus.RUNNING)
+
+        def updates(blocks, status):
+            for block in blocks:
+                get_run(block).update(status=status)
+
+        def test_status(blocks, status=BlockRun.BlockRunStatus.COMPLETED):
+            for block in blocks:
+                self.assertEqual(get_run(block).status, status)
+
+        def test_runs(blocks, cls=self, pipeline_run=pipeline_run):
+            cls.assertEqual(
+                [br.block_uuid for br in pipeline_run.executable_block_runs()],
+                [b.uuid for b in blocks],
+            )
+
+        test_runs([block1, block2, dynamic1])
+
+        complete([block2])
+        running([dynamic1])
+
+        # If a block run is running,
+        # it’s no longer considered an executable block run to be scheduled.
+        test_runs([block1, dynamic2, dynamic1_child1a, end3])
+
+        complete([block1])
+        test_runs([dynamic2, dynamic1_child1a, dynamic1_child2a, end3, end4])
+
+        running([dynamic2, end3, end4])
+        test_runs([dynamic1_child1a, dynamic1_child2a, dynamic2_child1b, dynamic2_child2b])
+
+        running([dynamic1_child1a])
+        test_runs([
+            dynamic1_child2a,
+            dynamic2_child1b,
+            dynamic2_child2b,
+            dynamic1_child1a_dynamic,
+            dynamic1_child1a_reduce1a,
+        ])
+
+        running([dynamic1_child2a])
+        test_runs([
+            dynamic2_child1b,
+            dynamic2_child2b,
+            dynamic1_child1a_dynamic,
+            dynamic1_child1a_reduce1a,
+            dynamic1_child1a_child2a_reduce2a,
+        ])
+
+        running([dynamic1_child1a_reduce1a])
+        test_runs([
+            dynamic2_child1b,
+            dynamic2_child2b,
+            dynamic1_child1a_dynamic,
+            dynamic1_child1a_child2a_reduce2a,
+        ])
+
+        complete([dynamic1_child1a_reduce1a])
+        test_runs([
+            dynamic2_child1b,
+            dynamic2_child2b,
+            dynamic1_child1a_dynamic,
+            dynamic1_child1a_child2a_reduce2a,
+            end1,
+        ])
+
+        running([dynamic1_child1a_child2a_reduce2a])
+        test_runs([
+            dynamic2_child1b,
+            dynamic2_child2b,
+            dynamic1_child1a_dynamic,
+            end1,
+        ])
+
+        complete([dynamic1_child1a_child2a_reduce2a])
+        test_runs([
+            dynamic2_child1b,
+            dynamic2_child2b,
+            dynamic1_child1a_dynamic,
+            end1,
+            end2,
+        ])
+
+        complete([
+            dynamic2_child1b,
+            dynamic2_child2b,
+            dynamic1_child1a_dynamic,
+            end1,
+            end2,
+        ])
+        test_runs([
+            dynamic2_child1b_dynamic,
+            dynamic2_child1b_reduce1b,
+            dynamic2_child1b_child2b_reduce2b,
+        ])
+
+        running([
+            dynamic2_child1b_dynamic,
+            dynamic2_child1b_reduce1b,
+            dynamic2_child1b_child2b_reduce2b,
+        ])
+
+        self.assertEqual(len(pipeline_run.executable_block_runs()), 0)
+
+        self.assertEqual(len(self.pipeline.blocks_by_uuid), len(pipeline_run.block_runs))
+        self.assertEqual(len(self.pipeline.blocks_by_uuid), 18)

mage_ai/tests/orchestration/test_pipeline_scheduler.py

@@ -453,6 +453,7 @@ class PipelineSchedulerTests(DBTestCase):
                 execution_date=None,
                 execution_partition=f'{pipeline_run.pipeline_schedule_id}',
                 event={},
+                pipeline_run_id=pipeline_run.id,
             ),
             dict(
                 pipeline_run_id=pipeline_run.id,

mage_ai/tests/orchestration/triggers/test_global_data_product.py

@@ -33,7 +33,7 @@ class TriggerGlobalDataProductTest(DBTestCase):
             self.pipeline.add_block(
                 Block('data_exporter', 'data_exporter', BlockType.DATA_EXPORTER))
         except Exception:
-            self.pipeline = Pipeline.get('test_pipeline')
+            self.pipeline = Pipeline.get('test_pipeline', repo_path=self.repo_path)
 
         self.global_data_product = GlobalDataProduct(
             object_type='pipeline',
@@ -55,6 +55,7 @@ class TriggerGlobalDataProductTest(DBTestCase):
                 week_of_month=8,
                 week_of_year=9,
             ),
+            repo_path=self.repo_path,
             settings=dict(
                 data_exporter={},
                 data_loader=dict(partitions=1),
@@ -67,7 +68,7 @@ class TriggerGlobalDataProductTest(DBTestCase):
             self.repo_path,
             'global_data_products.yaml',
         )
-        self.global_data_product.save(file_path=self.file_path)
+        self.global_data_product.save()
 
         PipelineRun.query.filter(
             PipelineRun.id.in_([pr.id for pr in self.global_data_product.pipeline_runs()]),

mage_ai/tests/orchestration/triggers/test_utils.py

@@ -29,7 +29,7 @@ class TriggerUtilsTest(DBTestCase):
             self.pipeline.add_block(
                 Block('data_exporter', 'data_exporter', BlockType.DATA_EXPORTER))
         except Exception:
-            self.pipeline = Pipeline.get('test_pipeline')
+            self.pipeline = Pipeline.get('test_pipeline', repo_path=self.repo_path)
 
         self.global_data_product = GlobalDataProduct(
             object_type='pipeline',
@@ -51,6 +51,7 @@ class TriggerUtilsTest(DBTestCase):
                 week_of_month=8,
                 week_of_year=9,
             ),
+            repo_path=self.repo_path,
             settings=dict(
                 data_exporter={},
                 data_loader=dict(partitions=1),
@@ -64,7 +65,7 @@ class TriggerUtilsTest(DBTestCase):
             self.repo_path,
             'global_data_products.yaml',
         )
-        self.global_data_product.save(file_path=self.file_path)
+        self.global_data_product.save()
 
     def tearDown(self):
         super().tearDown()

mage_ai/tests/services/k8s/test_job_manager.py

@@ -7,6 +7,7 @@ from kubernetes.client import (
     V1Container,
     V1EnvFromSource,
     V1EnvVar,
+    V1JobSpec,
     V1NodeAffinity,
     V1NodeSelector,
     V1NodeSelectorRequirement,
@@ -58,6 +59,10 @@ MOCK_POD_CONFIG = V1Pod(
         volumes=[],
     )
 )
+MOCK_JOB_CONFIG = V1JobSpec(template=MOCK_POD_CONFIG,
+                            active_deadline_seconds=120,
+                            backoff_limit=5,
+                            ttl_seconds_after_finished=86400)
 
 
 class TestK8sUtilities(TestCase):
@@ -227,7 +232,9 @@ class JobManagerTests(TestCase):
         )
         mock_client.V1JobSpec.assert_called_once_with(
             template=mock_v1_pod_template_spec,
+            active_deadline_seconds=None,
             backoff_limit=0,
+            ttl_seconds_after_finished=None
         )
         mock_client.V1Job.assert_called_once_with(
             api_version='batch/v1',
@@ -255,6 +262,7 @@ class JobManagerTests(TestCase):
             'env': [dict(name='spark_host', value='127.0.0.1')],
             'resources': dict(limits={'cpu': '1000m'}, requests={'cpu': '500m'}),
         }
+        k8s_config.job_config = MOCK_JOB_CONFIG
 
         # Call the method to create the job object
         command = "echo 'hello world'"
@@ -355,6 +363,11 @@ class JobManagerTests(TestCase):
                     },
                 ],
             },
+            'job': {
+                'active_deadline_seconds': 0,
+                'backoff_limit': 3,
+                'ttl_seconds_after_finished': 100
+            }
         }
         mock_getenv.return_value = 'pod_name'
 
@@ -428,6 +441,11 @@ class JobManagerTests(TestCase):
             )
         )
 
+        # Job
+        self.assertEqual(job.spec.active_deadline_seconds, 0)
+        self.assertEqual(job.spec.backoff_limit, 3)
+        self.assertEqual(job.spec.ttl_seconds_after_finished, 100)
+
     @patch('mage_ai.services.k8s.job_manager.client')
     @patch('mage_ai.services.k8s.job_manager.config')
     def test_create_job(self, mock_config, mock_client):

mage_ai/tests/streaming/sinks/test_oracledb.py

@@ -0,0 +1,38 @@
+from unittest.mock import patch
+
+from mage_ai.streaming.sinks.oracledb import OracleDbSink
+from mage_ai.tests.base_test import TestCase
+
+
+class OracleDbTests(TestCase):
+    def test_init(self):
+        with patch.object(OracleDbSink, 'init_client') as mock_init_client:
+            OracleDbSink(dict(
+                connector_type='oracledb',
+                user='test',
+                password='123456',
+                host='oracle-test',
+                port=1521,
+                table='test_table',
+                service_name='test_service',
+                mode='thin',
+            ))
+            mock_init_client.assert_called_once()
+
+    def test_init_invalid_config(self):
+        with patch.object(OracleDbSink, 'init_client') as mock_init_client:
+            with self.assertRaises(Exception) as context:
+                OracleDbSink(dict(
+                    connector_type='oracledb',
+                    user='test',
+                    password='123456',
+                    host='oracle-test',
+                    port=1521,
+                    table='test_table',
+                    mode='thin',
+                ))
+            self.assertTrue(
+                '__init__() missing 1 required positional argument: \'service_name\''
+                in str(context.exception),
+            )
+            self.assertEqual(mock_init_client.call_count, 0)

mage_ai/tests/test_shared.py

@@ -1,3 +1,7 @@
+from scipy.sparse import csr_matrix
+
+from mage_ai.data.models.generator import DataGenerator
+from mage_ai.data.tabular.mocks import create_dataframe
 from mage_ai.data_cleaner.shared.utils import wrap_column_name
 from mage_ai.tests.base_test import TestCase
 
@@ -22,3 +26,60 @@ class SharedTests(TestCase):
         ]
         for name, expected_name in zip(column_names, expected_names):
             self.assertEqual(wrap_column_name(name), expected_name)
+
+
+def build_pandas(*args, **kwargs):
+    return create_dataframe(n_rows=1_000, use_pandas=True)
+
+
+def build_pandas_series(*args, **kwargs):
+    df = create_dataframe(n_rows=2_000, use_pandas=True)
+    return df[df.columns[0]]
+
+
+def build_polars(*args, **kwargs):
+    return create_dataframe(n_rows=3_000)
+
+
+def build_polars_series(*args, **kwargs):
+    df = create_dataframe(n_rows=4_000)
+    return df[df.columns[0]]
+
+
+def build_iterable(*args, **kwargs):
+    return [i for i in range(5_000)]
+
+
+def build_matrix_sparse(*args, **kwargs):
+    row_count = 6000
+
+    data = []
+    rows = []
+    cols = []
+
+    for i in range(int(row_count / 2)):
+        # Replicate the 3x3 pattern for each block of 2 rows
+        data.extend([1, 2, 1, 2, 3, 2])  # Replicated entries for simplicity
+        rows.extend([
+            i * 2,
+            i * 2,
+            i * 2 + 1,
+            i * 2 + 1,
+            i * 2 + 1,
+            i * 2 + 1,
+        ])  # Repeat every two rows
+        cols.extend([0, 2, 2, 0, 1, 2])  # Columns indices are consistent with your pattern
+
+    # Create the CSR matrix with specified shape
+    return csr_matrix((data, (rows, cols)), shape=(row_count, 3))
+
+
+def build_list_complex(*args, **kwargs):
+    return [DataGenerator([i for i in range(100)]) for i in range(100)] + [
+        build_pandas(),
+        build_pandas_series(),
+        build_polars(),
+        build_polars_series(),
+        build_iterable(),
+        build_matrix_sparse(),
+    ]

mage_ai/usage_statistics/logger.py

@@ -45,8 +45,13 @@ from mage_ai.usage_statistics.utils import build_event_data_for_chart
 
 
 class UsageStatisticLogger():
-    def __init__(self, project=None):
-        self.project = project or Project()
+    def __init__(
+        self,
+        context_data: Dict = None,
+        project=None,
+        repo_path: str = None,
+    ):
+        self.project = project or Project(context_data=context_data, repo_path=repo_path)
 
     async def block_create(
         self,

mage_ai/utils/code.py

@@ -2,12 +2,14 @@ import importlib
 import re
 from typing import Tuple
 
+from jupyter_client import KernelClient
+
 from mage_ai.autocomplete.utils import extract_all_imports
 from mage_ai.settings.repo import get_repo_path
 
 
-def reload_all_repo_modules(content: str) -> None:
-    parts = get_repo_path().split('/')
+def reload_all_repo_modules(content: str, client: KernelClient) -> None:
+    parts = get_repo_path(root_project=True).split('/')
     project_name = parts[-1]
 
     for line in extract_all_imports(content):
@@ -21,6 +23,12 @@ def reload_all_repo_modules(content: str) -> None:
             m1, m2 = matches[0]
             try:
                 importlib.reload(importlib.import_module(m1 or m2))
+                client.execute(
+                    f"""
+import importlib
+importlib.reload(importlib.import_module("{m1}" or "{m2}"))
+                """
+                )
             except Exception:
                 pass
 
@@ -35,34 +43,40 @@ def extract_decorated_function(code: str, decorated_function_name: str) -> Tuple
     number_of_lines = len(lines)
 
     for idx, line in enumerate(lines):
-        if line and \
-                span_current_start is not None and \
-                span_current_def is not None:
+        if line and span_current_start is not None and span_current_def is not None:
 
             if line.startswith('@') or line.startswith('def '):
-                spans.append((
-                    span_current_start,
-                    idx - 1,
-                ))
+                spans.append(
+                    (
+                        span_current_start,
+                        idx - 1,
+                    )
+                )
                 span_current_start = None
                 span_current_def = None
 
-        if line and \
-                span_current_start is not None and \
-                span_current_def is None and \
-                line.startswith('def '):
+        if (
+            line
+            and span_current_start is not None
+            and span_current_def is None
+            and line.startswith('def ')
+        ):
 
             span_current_def = idx
 
-        if line and \
-                span_current_start is None and \
-                line.startswith(f'@{decorated_function_name}'):
+        if (
+            line
+            and span_current_start is None
+            and line.startswith(f'@{decorated_function_name}')
+        ):
 
             span_current_start = idx
 
-        if span_current_start is not None and \
-                span_current_def is not None and \
-                idx == number_of_lines - 1:
+        if (
+            span_current_start is not None
+            and span_current_def is not None
+            and idx == number_of_lines - 1
+        ):
 
             spans.append((span_current_start, idx))
 

{mage_ai-0.9.70.dist-info → mage_ai-0.9.71.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: mage-ai
-Version: 0.9.70
+Version: 0.9.71
 Summary: Mage is a tool for building and deploying data pipelines.
 Home-page: https://github.com/mage-ai/mage-ai
 Author: Mage
@@ -14,6 +14,7 @@ License-File: LICENSE
 Requires-Dist: Faker ==4.14.0
 Requires-Dist: GitPython ==3.1.41
 Requires-Dist: Jinja2 ==3.1.3
+Requires-Dist: Pillow ==10.3.0
 Requires-Dist: PyGithub ==1.59.0
 Requires-Dist: PyJWT ==2.6.0
 Requires-Dist: aiofiles ==22.1.0
@@ -35,14 +36,16 @@ Requires-Dist: jupyter-server-proxy ==3.2.3
 Requires-Dist: jupyter-server ==1.23.5
 Requires-Dist: jupyter-client ==7.4.4
 Requires-Dist: ldap3 ==2.9.1
+Requires-Dist: memory-profiler
 Requires-Dist: newrelic ==8.8.0
 Requires-Dist: numpy >=1.22.0
 Requires-Dist: pandas >=1.3.0
 Requires-Dist: polars <0.19.2,>=0.18.0
 Requires-Dist: protobuf ~=4.21.12
+Requires-Dist: psutil ==5.9.8
 Requires-Dist: pyarrow ==14.0.1
 Requires-Dist: python-dateutil ==2.8.2
-Requires-Dist: pytz ==2022.2.1
+Requires-Dist: pytz >=2022.2.1
 Requires-Dist: pyyaml ~=6.0
 Requires-Dist: redis ~=5.0.1
 Requires-Dist: requests ~=2.31.0