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

mage_ai/tests/data_preparation/models/block/platform/test_mixins.py

@@ -106,7 +106,7 @@ class BlockWithProjectPlatformInactiveTest(BaseAPIEndpointTest, BlockWithProject
 
         with patch(
             'mage_ai.data_preparation.models.block.File.from_path',
-            lambda _x: file,
+            lambda _x, repo_path=None: file,
         ):
             with patch.object(file, 'create_parent_directories') as mock_create_parent_directories:
                 with patch.object(file, 'update_content') as mock_update_content:

mage_ai/tests/data_preparation/models/block/sql/utils/test_shared.py

@@ -55,7 +55,7 @@ class SQLBlockSharedUtilsTest(DBTestCase):
         self.assertEqual(extract_create_statement_table_name(text), expected_result)
 
     def test_extract_insert_statement_table_names(self):
-        # Test case 1: Basic test with one INSERT statement
+        # Test case 1: Basic test with one INSERT statement and remove comment
         text = """
         -- This is a comment
         INSERT INTO table_name
@@ -79,6 +79,31 @@ class SQLBlockSharedUtilsTest(DBTestCase):
         expected_result = []
         self.assertEqual(extract_insert_statement_table_names(text), expected_result)
 
+        # Test case 4: Multiple INSERT statements
+        text = """
+            INSERT INTO table1 VALUES (1, 'John');
+            INSERT INTO table2 VALUES (2, 'Jane');
+            """
+        expected_result = ["table1", "table2"]
+        self.assertEqual(extract_insert_statement_table_names(text), expected_result)
+
+        # Test case 5: IGNORE INSERT statement
+        text = "INSERT IGNORE INTO table3 VALUES (3, 'Alice');"
+        expected_result = ["table3"]
+        self.assertEqual(extract_insert_statement_table_names(text), expected_result)
+
+        # Test case 6: case insensitivity
+        text = "insert into TABLE5 values (5, 'Eve');"
+        expected_result = ["TABLE5"]
+        self.assertEqual(extract_insert_statement_table_names(text), expected_result)
+
+        # Test case 7: Multiple spaces
+        text = """
+        INSERT   INTO   table_name VALUES (1, 'John');
+        """
+        expected_result = ['table_name']
+        self.assertEqual(extract_insert_statement_table_names(text), expected_result)
+
     def test_extract_drop_statement_table_names(self):
         # Test case 1: Basic DROP TABLE statement
         text = """

mage_ai/tests/data_preparation/models/block/test_global_data_product.py

@@ -20,7 +20,7 @@ class GlobalDataProductBlockTest(DBTestCase):
                 repo_path=self.repo_path,
             )
         except Exception:
-            self.pipeline = Pipeline.get('test_pipeline')
+            self.pipeline = Pipeline.get('test_pipeline', repo_path=self.repo_path)
 
         self.block = GlobalDataProductBlock(
             'GDP',
@@ -42,6 +42,7 @@ class GlobalDataProductBlockTest(DBTestCase):
             object_uuid=self.pipeline.uuid,
             outdated_after=dict(months=1),
             outdated_starting_at=dict(day_of_month=1),
+            repo_path=self.repo_path,
             settings=dict(data_loader=dict(partitions=1)),
             uuid='mage',
         )
@@ -50,7 +51,7 @@ class GlobalDataProductBlockTest(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/data_preparation/models/custom_templates/test_utils.py

@@ -1,10 +1,11 @@
 from unittest.mock import patch
 
-
 from mage_ai.data_preparation.models.custom_templates.custom_block_template import (
     CustomBlockTemplate,
 )
-from mage_ai.data_preparation.models.custom_templates.utils import group_and_hydrate_files
+from mage_ai.data_preparation.models.custom_templates.utils import (
+    group_and_hydrate_files,
+)
 from mage_ai.shared.array import find
 from mage_ai.tests.base_test import DBTestCase
 
@@ -20,8 +21,8 @@ class CustomTemplatesUtilsTest(DBTestCase):
             dict(parent_names=['mage', 'fire']),
         ]
 
-        def load(template_uuid: str):
-            return CustomBlockTemplate(template_uuid=template_uuid)
+        def load(repo_path: str, template_uuid: str):
+            return CustomBlockTemplate(repo_path=repo_path, template_uuid=template_uuid)
 
         with patch.object(
             CustomBlockTemplate,

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

@@ -334,6 +334,7 @@ class HookTest(GlobalHooksMixin):
 
             mock_trigger_pipeline.assert_called_once_with(
                 hook.pipeline.uuid,
+                repo_path=self.repo_path,
                 variables=merge_dict(variables, kwargs),
                 check_status=True,
                 error_on_failure=True,
@@ -381,6 +382,7 @@ class HookTest(GlobalHooksMixin):
 
             mock_trigger_pipeline.assert_called_once_with(
                 hook.pipeline.uuid,
+                repo_path=self.repo_path,
                 variables=merge_dict(variables, kwargs),
                 check_status=False,
                 error_on_failure=True,
@@ -426,6 +428,7 @@ class HookTest(GlobalHooksMixin):
 
             mock_trigger_pipeline.assert_called_once_with(
                 hook.pipeline.uuid,
+                repo_path=self.repo_path,
                 variables=merge_dict(variables, dict(hook=kwargs['hook'], project=None)),
                 check_status=True,
                 error_on_failure=True,

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

@@ -18,6 +18,7 @@ from mage_ai.data_preparation.models.global_hooks.predicates import (
     PredicateValueType,
 )
 from mage_ai.data_preparation.models.pipeline import Pipeline
+from mage_ai.settings.repo import get_repo_path
 from mage_ai.shared.array import find
 from mage_ai.shared.hash import dig, ignore_keys, merge_dict, set_value
 from mage_ai.tests.base_test import AsyncDBTestCase
@@ -39,6 +40,7 @@ def build_operation_resource_settings(
     first_object_key=first_object_key,
     repo_path: str = None,
 ) -> Dict:
+    repo_path = repo_path or get_repo_path()
     settings = {
         BaseResource.__name__: dict(
             right_object_keys=[first_object_key],
@@ -62,9 +64,13 @@ def build_operation_resource_settings(
         ),
         Pipeline.__name__: dict(
             right_object_keys=['variables', first_object_key],
-            operation_resource=Pipeline(uuid.uuid4().hex, config=dict(variables={
-                first_object_key: right_value,
-            })),
+            operation_resource=Pipeline(
+                uuid.uuid4().hex,
+                config=dict(variables={
+                    first_object_key: right_value,
+                }),
+                repo_path=repo_path,
+            ),
         ),
     }
 

mage_ai/tests/data_preparation/models/test_block.py

@@ -26,14 +26,18 @@ class BlockTest(DBTestCase):
     def test_create(self):
         block1 = Block.create('test_transformer', 'transformer', self.repo_path)
         block2 = Block.create('test data loader', BlockType.DATA_LOADER, self.repo_path)
-        self.assertTrue(os.path.exists(os.path.join(
-            self.repo_path, 'transformers', 'test_transformer.py')))
-        self.assertTrue(os.path.exists(os.path.join(
-            self.repo_path, 'transformers', '__init__.py')))
-        self.assertTrue(os.path.exists(os.path.join(
-            self.repo_path, 'data_loaders', 'test_data_loader.py')))
-        self.assertTrue(os.path.exists(os.path.join(
-            self.repo_path, 'data_loaders', '__init__.py')))
+        self.assertTrue(
+            os.path.exists(os.path.join(self.repo_path, 'transformers', 'test_transformer.py'))
+        )
+        self.assertTrue(
+            os.path.exists(os.path.join(self.repo_path, 'transformers', '__init__.py'))
+        )
+        self.assertTrue(
+            os.path.exists(os.path.join(self.repo_path, 'data_loaders', 'test_data_loader.py'))
+        )
+        self.assertTrue(
+            os.path.exists(os.path.join(self.repo_path, 'data_loaders', '__init__.py'))
+        )
         self.assertEqual(block1.name, 'test_transformer')
         self.assertEqual(block1.uuid, 'test_transformer')
         self.assertEqual(block1.type, 'transformer')
@@ -55,20 +59,20 @@ class BlockTest(DBTestCase):
             upstream_block_uuids=['test_data_loader'],
         )
         with open(block1.file_path, 'w') as file:
-            file.write('''import pandas as pd
+            file.write("""import pandas as pd
 @data_loader
 def load_data():
     data = {'col1': [1, 1, 3], 'col2': [2, 2, 4]}
     df = pd.DataFrame(data)
     return [df]
-            ''')
+            """)
         with open(block2.file_path, 'w') as file:
-            file.write('''import pandas as pd
+            file.write("""import pandas as pd
 @transformer
 def remove_duplicate_rows(df):
     df_transformed = df.drop_duplicates()
     return [df_transformed]
-            ''')
+            """)
         asyncio.run(block1.execute(analyze_outputs=True))
         asyncio.run(block2.execute(analyze_outputs=True))
 
@@ -76,10 +80,7 @@ def remove_duplicate_rows(df):
             variables_dir=get_repo_config(self.repo_path).variables_dir,
         )
         data = variable_manager.get_variable(
-            pipeline.uuid,
-            block2.uuid,
-            'output_0',
-            variable_type='dataframe'
+            pipeline.uuid, block2.uuid, 'output_0', variable_type='dataframe'
         )
         # analysis = variable_manager.get_variable(
         #     pipeline.uuid,
@@ -111,7 +112,7 @@ def remove_duplicate_rows(df):
             upstream_block_uuids=['test_data_loader'],
         )
         with open(block1.file_path, 'w') as file:
-            file.write('''import pandas as pd
+            file.write("""import pandas as pd
 @preprocesser
 def preprocesser0(*args, **kwargs):
     kwargs['context']['count'] = 1
@@ -129,14 +130,14 @@ def load_data(**kwargs):
     data = {'col1': [i * count for i in [1, 1, 3]], 'col2': [i * count for i in [2, 2, 4]]}
     df = pd.DataFrame(data)
     return [df]
-            ''')
+            """)
         with open(block2.file_path, 'w') as file:
-            file.write('''import pandas as pd
+            file.write("""import pandas as pd
 @transformer
 def remove_duplicate_rows(df):
     df_transformed = df.drop_duplicates()
     return [df_transformed]
-            ''')
+            """)
         asyncio.run(block1.execute())
         asyncio.run(block2.execute())
 
@@ -144,10 +145,7 @@ def remove_duplicate_rows(df):
             variables_dir=get_repo_config(self.repo_path).variables_dir,
         )
         data = variable_manager.get_variable(
-            pipeline.uuid,
-            block2.uuid,
-            'output_0',
-            variable_type='dataframe'
+            pipeline.uuid, block2.uuid, 'output_0', variable_type='dataframe'
         )
 
         df_final = pd.DataFrame({'col1': [2, 2, 6], 'col2': [4, 4, 8]}).drop_duplicates()
@@ -172,7 +170,7 @@ def remove_duplicate_rows(df):
             upstream_block_uuids=['test_data_loader_2'],
         )
         with open(block1.file_path, 'w') as file:
-            file.write('''import pandas as pd
+            file.write("""import pandas as pd
 @data_loader
 def load_data():
     data = {
@@ -183,14 +181,14 @@ def load_data():
     }
     df = pd.DataFrame(data)
     return [df]
-            ''')
+            """)
         with open(block2.file_path, 'w') as file:
-            file.write('''import pandas as pd
+            file.write("""import pandas as pd
 @transformer
 def remove_duplicate_rows(df):
     df_transformed = df
     return [df_transformed]
-            ''')
+            """)
         asyncio.run(block1.execute())
         asyncio.run(block2.execute())
 
@@ -198,10 +196,7 @@ def remove_duplicate_rows(df):
             variables_dir=get_repo_config(self.repo_path).variables_dir,
         )
         data = variable_manager.get_variable(
-            pipeline.uuid,
-            block2.uuid,
-            'output_0',
-            variable_type='dataframe'
+            pipeline.uuid, block2.uuid, 'output_0', variable_type='dataframe'
         )
         df_final = pd.DataFrame({
             'col1': [1, 1, 3],
@@ -236,28 +231,28 @@ def remove_duplicate_rows(df):
             upstream_block_uuids=['test_data_loader_1', 'test_data_loader_2'],
         )
         with open(block1.file_path, 'w') as file:
-            file.write('''import pandas as pd
+            file.write("""import pandas as pd
 @data_loader
 def load_data():
     data = {'col1': [1, 3], 'col2': [2, 4]}
     df = pd.DataFrame(data)
     return [df]
-            ''')
+            """)
         with open(block2.file_path, 'w') as file:
-            file.write('''import pandas as pd
+            file.write("""import pandas as pd
 @data_loader
 def load_data():
     data = {'col1': [5], 'col2': [6]}
     df = pd.DataFrame(data)
     return [df]
-            ''')
+            """)
         with open(block3.file_path, 'w') as file:
-            file.write('''import pandas as pd
+            file.write("""import pandas as pd
 @transformer
 def union_datasets(df1, df2):
     df_union = pd.concat([df1, df2]).reset_index(drop=True)
     return [df_union]
-            ''')
+            """)
         asyncio.run(block1.execute(analyze_outputs=True))
         asyncio.run(block2.execute(analyze_outputs=True))
         asyncio.run(block3.execute(analyze_outputs=True))
@@ -266,10 +261,7 @@ def union_datasets(df1, df2):
             variables_dir=get_repo_config(self.repo_path).variables_dir,
         )
         data = variable_manager.get_variable(
-            pipeline.uuid,
-            block3.uuid,
-            'output_0',
-            variable_type='dataframe'
+            pipeline.uuid, block3.uuid, 'output_0', variable_type='dataframe'
         )
         analysis = variable_manager.get_variable(
             pipeline.uuid,
@@ -319,48 +311,48 @@ def union_datasets(df1, df2):
             upstream_block_uuids=['test_data_loader_1', 'test_data_loader_2'],
         )
         with open(block1.file_path, 'w') as file:
-            file.write('''import pandas as pd
+            file.write("""import pandas as pd
 @data_loader
 def load_data():
     data = {'col1': [1, 3], 'col2': [2, 4]}
     df = pd.DataFrame(data)
     return [df]
-            ''')
+            """)
         with open(block2.file_path, 'w') as file:
-            file.write('''import pandas as pd
+            file.write("""import pandas as pd
 @data_loader
 def load_data():
     data = {'col1': [5], 'col2': [6]}
     df = pd.DataFrame(data)
     return [df]
-            ''')
+            """)
         with open(block3.file_path, 'w') as file:
-            file.write('''import pandas as pd
+            file.write("""import pandas as pd
 @transformer
 def incorrect_function(df1):
     return df1
-            ''')
+            """)
         asyncio.run(block1.execute())
         asyncio.run(block2.execute())
         with self.assertRaisesRegex(
             Exception,
-            'Block test_transformer may have too many upstream dependencies. It expected to have' +
-            ' 1 arguments, but received 2. Confirm that the @transformer method declaration has ' +
-            'the correct number of arguments.'
+            'Block test_transformer may have too many upstream dependencies. It expected to have'
+            + ' 1 arguments, but received 2. Confirm that the @transformer method declaration has '
+            + 'the correct number of arguments.',
         ):
             asyncio.run(block3.execute())
 
         with open(block3.file_path, 'w') as file:
-            file.write('''import pandas as pd
+            file.write("""import pandas as pd
 @transformer
 def incorrect_function(df1, df2, df3):
     return df1
-            ''')
+            """)
         with self.assertRaisesRegex(
             Exception,
-            'Block test_transformer may have too many upstream dependencies. It expected to have' +
-            ' 1 arguments, but received 2. Confirm that the @transformer method declaration has ' +
-            'the correct number of arguments.'
+            'Block test_transformer may have too many upstream dependencies. It expected to have'
+            + ' 1 arguments, but received 2. Confirm that the @transformer method declaration has '
+            + 'the correct number of arguments.',
         ):
             asyncio.run(block3.execute())
 
@@ -384,20 +376,20 @@ def incorrect_function(df1, df2, df3):
         block2.upstream_blocks = [block1]
         block1.downstream_blocks = [block2]
         with open(block1.file_path, 'w') as file:
-            file.write('''import pandas as pd
+            file.write("""import pandas as pd
 @data_loader
 def load_data():
     data = {'col1': [1, 3], 'col2': [2, 4]}
     df = pd.DataFrame(data)
     return [df]
-            ''')
+            """)
         with open(block2.file_path, 'w') as file:
-            file.write('''import pandas as pd
+            file.write("""import pandas as pd
 @sensor
 def sensor(*args):
     df = args[0]
     return df is not None
-            ''')
+            """)
         asyncio.run(block1.execute())
         output = block2.execute_sync(from_notebook=True)
         self.assertEqual(output['output'][0], True)
@@ -415,20 +407,20 @@ def sensor(*args):
             pipeline=pipeline,
         )
         with open(block1.file_path, 'w') as file:
-            file.write('''import pandas as pd
+            file.write("""import pandas as pd
 @data_loader
 def load_data():
     data = {'col1': [1, 3], 'col2': [2, 4]}
     df = pd.DataFrame(data)
     return [df]
-            ''')
+            """)
         block1.update(dict(has_callback=True))
         with open(block1.callback_block.file_path, 'w') as file:
-            file.write('''
+            file.write("""
 @on_success
 def on_success_callback(**kwargs):
     print('SUCCESS')
-            ''')
+            """)
         block1.execute_with_callback()
         mock_print.assert_called_with('SUCCESS')
 
@@ -445,18 +437,18 @@ def on_success_callback(**kwargs):
             pipeline=pipeline,
         )
         with open(block1.file_path, 'w') as file:
-            file.write('''import pandas as pd
+            file.write("""import pandas as pd
 @data_loader
 def load_data():
     raise Exception('failed')
-            ''')
+            """)
         block1.update(dict(has_callback=True))
         with open(block1.callback_block.file_path, 'w') as file:
-            file.write('''
+            file.write("""
 @on_failure
 def on_failure_callback(**kwargs):
     print('FAILED')
-            ''')
+            """)
 
         with self.assertRaisesRegex(Exception, 'failed'):
             block1.execute_with_callback()
@@ -477,40 +469,46 @@ def on_failure_callback(**kwargs):
         )
         block2.upstream_blocks = [block1]
         block1.downstream_blocks = [block2]
-        self.assertEqual(block1.to_dict(), dict(
-            all_upstream_blocks_executed=True,
-            color=None,
-            configuration={},
-            downstream_blocks=['test_data_exporter'],
-            executor_config=None,
-            executor_type='local_python',
-            has_callback=False,
-            language='sql',
-            name='test_transformer_2',
-            retry_config=None,
-            status='not_executed',
-            timeout=None,
-            type='transformer',
-            upstream_blocks=[],
-            uuid='test_transformer_2',
-        ))
-        self.assertEqual(block2.to_dict(), dict(
-            all_upstream_blocks_executed=False,
-            color=None,
-            configuration={},
-            downstream_blocks=[],
-            executor_config=None,
-            executor_type='local_python',
-            has_callback=False,
-            language='python',
-            name='test_data_exporter',
-            retry_config=None,
-            status='not_executed',
-            timeout=None,
-            type='data_exporter',
-            upstream_blocks=['test_transformer_2'],
-            uuid='test_data_exporter',
-        ))
+        self.assertEqual(
+            block1.to_dict(),
+            dict(
+                all_upstream_blocks_executed=True,
+                color=None,
+                configuration={},
+                downstream_blocks=['test_data_exporter'],
+                executor_config=None,
+                executor_type='local_python',
+                has_callback=False,
+                language='sql',
+                name='test_transformer_2',
+                retry_config=None,
+                status='not_executed',
+                timeout=None,
+                type='transformer',
+                upstream_blocks=[],
+                uuid='test_transformer_2',
+            ),
+        )
+        self.assertEqual(
+            block2.to_dict(),
+            dict(
+                all_upstream_blocks_executed=False,
+                color=None,
+                configuration={},
+                downstream_blocks=[],
+                executor_config=None,
+                executor_type='local_python',
+                has_callback=False,
+                language='python',
+                name='test_data_exporter',
+                retry_config=None,
+                status='not_executed',
+                timeout=None,
+                type='data_exporter',
+                upstream_blocks=['test_transformer_2'],
+                uuid='test_data_exporter',
+            ),
+        )
 
     def test_full_table_name(self):
         faker = Faker()
@@ -518,7 +516,7 @@ def on_failure_callback(**kwargs):
         block = Block.create(faker.name(), 'data_loader', self.repo_path, pipeline=pipeline)
 
         def get_block():
-            return Pipeline.get(pipeline.uuid).get_block(block.uuid)
+            return Pipeline.get(pipeline.uuid, repo_path=self.repo_path).get_block(block.uuid)
 
         with open(block.file_path, 'w') as file:
             file.write('')
@@ -614,7 +612,7 @@ SELECT 1 AS id;
 
         test_block.delete(force=True)
 
-        pipeline = Pipeline.get('test_pipeline_delete_force')
+        pipeline = Pipeline.get('test_pipeline_delete_force', repo_path=self.repo_path)
         self.assertIsNone(pipeline.get_block('test_data_loader'))
 
         block2 = pipeline.get_block('test_transformer')
@@ -732,19 +730,20 @@ def test_output(output, *args) -> None:
             columns=['id', 'email', 'zip_code'],
         )
 
-        block.store_variables({
-            'output_sample_data_stream1': df
-        })
+        block.store_variables({'output_sample_data_stream1': df})
         output_variables = block.output_variables()
         self.assertTrue('output_sample_data_stream1' in output_variables)
 
     def test_file_path(self):
         block = Block.create('test_transformer', 'transformer', self.repo_path)
-        self.assertEqual(block.file_path, os.path.join(
-            self.repo_path,
-            'transformers',
-            f'{block.uuid}.py',
-        ))
+        self.assertEqual(
+            block.file_path,
+            os.path.join(
+                self.repo_path,
+                'transformers',
+                f'{block.uuid}.py',
+            ),
+        )
 
 
 class BlockProjectPlatformTests(ProjectPlatformMixin):
@@ -763,7 +762,7 @@ class BlockProjectPlatformTests(ProjectPlatformMixin):
                     'test_transformer',
                     'transformer',
                     self.repo_path,
-                    configuration=dict(file_source=dict(path=path))
+                    configuration=dict(file_source=dict(path=path)),
                 )
 
                 self.assertEqual(block.file_path, add_root_repo_path_to_relative_path(path))
@@ -783,14 +782,10 @@ class CallbackBlockTest(DBTestCase):
             self.repo_path,
             pipeline=self.pipeline,
         )
-        callback_block = CallbackBlock.create(parent_block.name)
+        callback_block = CallbackBlock.create(parent_block.name, self.repo_path)
         self.pipeline.add_block(callback_block)
-        parent_block = parent_block.update(
-            dict(callback_blocks=[callback_block.uuid])
-        )
-        parent_block.global_vars = dict(
-            configuration=dict(table_name='load_data_table')
-        )
+        parent_block = parent_block.update(dict(callback_blocks=[callback_block.uuid]))
+        parent_block.global_vars = dict(configuration=dict(table_name='load_data_table'))
 
         global_vars = dict(
             random_var=1,