mage-ai 0.9.55__py3-none-any.whl → 0.9.57__py3-none-any.whl

mage_ai/settings/platform/__init__.py

@@ -163,10 +163,10 @@ def project_platform_activated() -> bool:
     return ENABLE_PROJECT_PLATFORM and os.path.exists(platform_settings_full_path())
 
 
-def platform_settings() -> Dict:
+def platform_settings(mage_projects_only: bool = False) -> Dict:
     config = __load_platform_settings(platform_settings_full_path()) or {}
     config['projects'] = merge_dict(
-        __get_projects_of_any_type() or {},
+        {} if mage_projects_only else (__get_projects_of_any_type() or {}),
         (config.get('projects') if config else {}) or {},
     )
     return config
@@ -181,26 +181,31 @@ def active_project_settings(
         settings = project_platform_settings(repo_path=repo_path, mage_projects_only=True)
 
     items = list(settings.items())
-    if items:
-        project_settings_tup = find(
-            lambda tup: tup and len(tup) >= 2 and (tup[1] or {}).get('active'),
-            items,
-        )
+    if not items:
+        return
 
-        if not project_settings_tup and get_default:
-            project_settings_tup = items[0]
+    project_settings_tup = find(
+        lambda tup: tup and len(tup) >= 2 and (tup[1] or {}).get('active'),
+        items,
+    )
 
-        if project_settings_tup:
-            project_name, project_settings = project_settings_tup
+    if not project_settings_tup and get_default:
+        project_settings_tup = items[0]
 
-            return merge_dict(
-                project_settings or {},
-                dict(uuid=project_name),
-            )
+    if project_settings_tup:
+        project_name, project_settings = project_settings_tup
+
+        return merge_dict(
+            project_settings or {},
+            dict(uuid=project_name),
+        )
 
 
 def project_platform_settings(repo_path: str = None, mage_projects_only: bool = False) -> Dict:
-    mapping = (__combined_platform_settings(repo_path=repo_path) or {}).get('projects')
+    mapping = (__combined_platform_settings(
+        repo_path=repo_path,
+        mage_projects_only=mage_projects_only,
+    ) or {}).get('projects')
 
     if mage_projects_only:
         select_keys = []
@@ -214,9 +219,28 @@ def project_platform_settings(repo_path: str = None, mage_projects_only: bool =
     return mapping
 
 
-def __combined_platform_settings(repo_path: str = None) -> Dict:
-    child = (platform_settings() or {}).copy()
-    parent = (__local_platform_settings(repo_path=repo_path) or {}).copy()
+def update_settings(settings: Dict) -> Dict:
+    projects = {}
+    for project_name, project_settings in (settings.get('projects') or {}).items():
+        uuid = project_settings.get('uuid') or project_name
+        projects[uuid] = extract(project_settings or {}, [
+            'path',
+        ])
+
+    settings['projects'] = projects
+    content = yaml.dump(settings)
+
+    safe_write(platform_settings_full_path(), content)
+
+
+def __combined_platform_settings(repo_path: str = None, mage_projects_only: bool = False) -> Dict:
+    parent = (platform_settings() or {}).copy()
+    child = (__local_platform_settings(repo_path=repo_path) or {}).copy()
+
+    if mage_projects_only:
+        keys = (parent.get('projects') or {}).keys()
+        child['projects'] = extract(child.get('projects') or {}, keys)
+
     combine_into(child, parent)
     return parent
 
@@ -249,7 +273,7 @@ def __get_projects_of_any_type() -> Dict:
     repo_path = base_repo_path()
     for path in os.listdir(repo_path):
         project_path = os.path.join(repo_path, path)
-        if not os.path.isdir(project_path):
+        if not os.path.isdir(project_path) or path.startswith('.'):
             continue
 
         is_project = False

mage_ai/shared/io.py

@@ -87,6 +87,9 @@ async def safe_write_async(filepath: str, content: str, write_func: Callable = N
 
 
 async def read_last_line_async(file_path: str) -> str:
+    if not file_path:
+        return
+
     # https://stackoverflow.com/questions/46258499/how-to-read-the-last-line-of-a-file-in-python
     async with aiofiles.open(file_path, mode='rb') as f:
         # catch OSError in case of a one line file

mage_ai/streaming/sinks/rabbitmq.py

@@ -15,6 +15,7 @@ class RabbitMQConfig(BaseConfig):
     queue_name: str
     username: str = 'guest'
     password: str = 'guest'
+    url_protocol: str = 'amqp'
     amqp_url_virtual_host: str = r'%2f'
 
 
@@ -29,12 +30,13 @@ class RabbitMQSink(BaseSink):
         password = self.config.password
         connection_host = self.config.connection_host
         connection_port = self.config.connection_port
+        url_protocol = self.config.url_protocol
         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}@" \
+            generated_url = f"{url_protocol}://{username}:{password}@" \
                             f"{connection_host}:{connection_port}/{vt_host}"
 
             self._print(f'Trying to connect on {generated_url}')

mage_ai/streaming/sources/rabbitmq.py

@@ -3,9 +3,10 @@ from dataclasses import dataclass
 from typing import Callable, Dict
 
 import pika
+from pika.exceptions import AMQPConnectionError
+
 from mage_ai.shared.config import BaseConfig
 from mage_ai.streaming.sources.base import BaseSource
-from pika.exceptions import AMQPConnectionError
 
 
 @dataclass
@@ -24,6 +25,7 @@ class RabbitMQConfig(BaseConfig):
     username: str = 'guest'
     password: str = 'guest'
     amqp_url_virtual_host: str = r'%2f'
+    url_protocol: str = 'amqp'
     consume_config: ConsumeConfig = None
 
     @classmethod
@@ -43,13 +45,14 @@ class RabbitMQSource(BaseSource):
         password = self.config.password
         connection_host = self.config.connection_host
         connection_port = self.config.connection_port
+        url_protocol = self.config.url_protocol
         vt_host = self.config.amqp_url_virtual_host
 
         self._print(f'Starting to initialize consumer for queue {queue_name}')
 
         try:
 
-            generated_url = f"amqp://{username}:{password}@" \
+            generated_url = f"{url_protocol}://{username}:{password}@" \
                             f"{connection_host}:{connection_port}/{vt_host}"
 
             self._print(f'Trying to connect on {generated_url}')

mage_ai/tests/api/endpoints/test_pipeline_runs.py

@@ -84,6 +84,10 @@ class PipelineRunAPIEndpointTest(BaseAPIEndpointTest):
 
 
 @patch('mage_ai.settings.platform.utils.project_platform_activated', lambda: True)
+@patch(
+    'mage_ai.orchestration.db.models.schedules.project_platform_activated',
+    lambda: True,
+)
 class PipelineRunProjectPlatformTests(ProjectPlatformMixin, BaseAPIEndpointTest):
     async def test_collection_project_platform_activated(self):
         pipelines = []

mage_ai/tests/api/endpoints/test_projects.py

@@ -126,5 +126,6 @@ build_update_endpoint_tests(
     assert_after_update=_assert_after_update,
     patch_function_settings=[
         ('mage_ai.api.resources.ProjectResource.activate_project',),
+        ('mage_ai.data_preparation.models.project.project_platform_activated', lambda: True),
     ],
 )

mage_ai/tests/api/resources/test_pipeline_resource.py

@@ -339,8 +339,8 @@ class PipelineResourceTest(BaseApiTestCase):
             lambda: True,
         ):
             with patch(
-                'mage_ai.api.resources.PipelineResource.Pipeline.get_async',
+                'mage_ai.api.resources.PipelineResource.get_pipeline_from_platform_async',
                 wraps=Pipeline.get_async,
             ) as mock:
                 await PipelineResource.member(self.pipeline1.uuid, None)
-                mock.assert_called_once_with(self.pipeline1.uuid, all_projects=True)
+                mock.assert_called_once_with(self.pipeline1.uuid)

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

@@ -7,6 +7,7 @@ from mage_ai.data_preparation.models.block.dbt.block_sql import DBTBlock, DBTBlo
 from mage_ai.data_preparation.models.constants import BlockLanguage, BlockType
 from mage_ai.data_preparation.models.file import File
 from mage_ai.settings.utils import base_repo_path
+from mage_ai.shared.path_fixer import add_root_repo_path_to_relative_path
 from mage_ai.tests.api.endpoints.mixins import BaseAPIEndpointTest
 from mage_ai.tests.shared.mixins import ProjectPlatformMixin
 
@@ -45,7 +46,7 @@ class BlockWithProjectPlatformShared:
             'get_file_path_from_source',
             lambda value=value: value,
         ):
-            self.assertEqual(self.block.file_path, value)
+            self.assertEqual(self.block.file_path, add_root_repo_path_to_relative_path(value))
 
         with patch.object(
             self.block,
@@ -197,7 +198,7 @@ class BlockWithProjectPlatformActivatedTest(ProjectPlatformMixin, BlockWithProje
         with patch('mage_ai.data_preparation.models.block.load_template') as mock_load_template:
             with patch(
                 'mage_ai.data_preparation.models.block.from_another_project',
-                lambda x, file_path=file_path: file_path == x,
+                lambda file_path, other_file_path, x=file_path: file_path == x,
             ):
                 self.build_block(configuration=dict(
                     file_source=dict(
@@ -223,7 +224,7 @@ class BlockWithProjectPlatformActivatedTest(ProjectPlatformMixin, BlockWithProje
                     ):
                         with patch(
                             'mage_ai.data_preparation.models.block.from_another_project',
-                            lambda x, file_path=file_path: file_path == x,
+                            lambda file_path, other_file_path, x=file_path: file_path == x,
                         ):
                             self.build_block(
                                 block_type=BlockType.DBT,

mage_ai/tests/data_preparation/models/test_block.py

@@ -12,6 +12,7 @@ from mage_ai.data_preparation.models.block.errors import HasDownstreamDependenci
 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.shared.path_fixer import add_root_repo_path_to_relative_path
 from mage_ai.tests.base_test import DBTestCase
 from mage_ai.tests.factory import create_integration_pipeline_with_blocks
 from mage_ai.tests.shared.mixins import ProjectPlatformMixin
@@ -763,4 +764,4 @@ class BlockProjectPlatformTests(ProjectPlatformMixin):
                     configuration=dict(file_source=dict(path=path))
                 )
 
-                self.assertEqual(block.file_path, path)
+                self.assertEqual(block.file_path, add_root_repo_path_to_relative_path(path))

mage_ai/tests/data_preparation/models/test_project.py

@@ -1,5 +1,5 @@
 import os
-from unittest.mock import patch
+from unittest.mock import call, patch
 
 import yaml
 
@@ -80,9 +80,19 @@ class ProjectTest(ProjectPlatformMixin, AsyncDBTestCase):
                         self.assertEqual(project.version, VERSION)
                         self.assertFalse(project.root_project)
 
-                        mock_get_repo_config.assert_called_once_with(
-                            repo_path=os.path.join(base_repo_path(), 'mage_platform'),
-                            root_project=False,
+                        self.assertEqual(
+                            mock_get_repo_config.mock_calls[0],
+                            call(
+                                repo_path=os.path.join(base_repo_path(), 'mage_platform'),
+                                root_project=False,
+                            ),
+                        )
+                        self.assertEqual(
+                            mock_get_repo_config.mock_calls[1],
+                            call(
+                                repo_path=base_repo_path(),
+                                root_project=True,
+                            ),
                         )
 
                     with patch(
@@ -96,9 +106,19 @@ class ProjectTest(ProjectPlatformMixin, AsyncDBTestCase):
                         self.assertEqual(project.version, VERSION)
                         self.assertTrue(project.root_project)
 
-                        mock_get_repo_config.assert_called_once_with(
-                            repo_path=base_repo_path(),
-                            root_project=True,
+                        self.assertEqual(
+                            mock_get_repo_config.mock_calls[0],
+                            call(
+                                repo_path=base_repo_path(),
+                                root_project=True,
+                            ),
+                        )
+                        self.assertEqual(
+                            mock_get_repo_config.mock_calls[1],
+                            call(
+                                repo_path=base_repo_path(),
+                                root_project=True,
+                            ),
                         )
 
     def test_repo_path_for_database_query(self):

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

@@ -1791,6 +1791,10 @@ class BlockRunTests(DBTestCase):
             self.assertEqual(b.logs.get('path'), expected_file_path)
 
 
+@patch(
+    'mage_ai.orchestration.db.models.schedules.project_platform_activated',
+    lambda: True,
+)
 class PipelineScheduleProjectPlatformTests(ProjectPlatformMixin):
     def test_repo_query(self):
         with patch(
@@ -1831,7 +1835,9 @@ class PipelineScheduleProjectPlatformTests(ProjectPlatformMixin):
                 self.assertIn(ps.id, ids)
 
     def test_pipeline(self):
-        with patch('mage_ai.orchestration.db.models.schedules.get_pipeline_from_platform') as mock:
+        with patch(
+            'mage_ai.orchestration.db.models.schedules_project_platform.get_pipeline_from_platform',
+        ) as mock:
             pipeline_uuid = self.faker.unique.name()
             repo_path = self.faker.unique.name()
 
@@ -1847,6 +1853,10 @@ class PipelineScheduleProjectPlatformTests(ProjectPlatformMixin):
             )
 
 
+@patch(
+    'mage_ai.orchestration.db.models.schedules.project_platform_activated',
+    lambda: True,
+)
 class PipelineRunProjectPlatformTests(ProjectPlatformMixin, AsyncDBTestCase):
     def test_pipeline(self):
         pipeline_uuid = self.faker.unique.name()
@@ -1864,7 +1874,9 @@ class PipelineRunProjectPlatformTests(ProjectPlatformMixin, AsyncDBTestCase):
             create_block_runs=False,
         )
 
-        with patch('mage_ai.orchestration.db.models.schedules.get_pipeline_from_platform') as mock:
+        with patch(
+            'mage_ai.orchestration.db.models.schedules_project_platform.get_pipeline_from_platform',
+        ) as mock:
             pipeline_run.pipeline
 
             mock.assert_called_once_with(
@@ -1926,11 +1938,12 @@ class PipelineRunProjectPlatformTests(ProjectPlatformMixin, AsyncDBTestCase):
                 return fake_logger_manager
 
         with patch(
-            'mage_ai.orchestration.db.models.schedules.get_pipeline_from_platform_async',
+            'mage_ai.orchestration.db.models.schedules_project_platform.'
+            'get_pipeline_from_platform_async',
             __get_pipeline_from_platform_async,
         ):
             with patch(
-                'mage_ai.orchestration.db.models.schedules.LoggerManagerFactory',
+                'mage_ai.orchestration.db.models.schedules_project_platform.LoggerManagerFactory',
                 FakeLoggerManagerFactory,
             ):
                 await pipeline_run.logs_async()
@@ -1951,7 +1964,7 @@ class PipelineRunProjectPlatformTests(ProjectPlatformMixin, AsyncDBTestCase):
         )
 
         with patch(
-            'mage_ai.orchestration.db.models.schedules.get_global_variables',
+            'mage_ai.orchestration.db.models.schedules_project_platform.get_global_variables',
             wraps=lambda pipeline_uuid, pipeline: {},
         ) as mock:
             with patch(
@@ -1963,6 +1976,10 @@ class PipelineRunProjectPlatformTests(ProjectPlatformMixin, AsyncDBTestCase):
                 self.assertEqual(mock.mock_calls[0][2]['pipeline'].uuid, pipeline_run.pipeline_uuid)
 
 
+@patch(
+    'mage_ai.orchestration.db.models.schedules.project_platform_activated',
+    lambda: True,
+)
 class BlockRunProjectPlatformTests(ProjectPlatformMixin, AsyncDBTestCase):
     async def test_logs_async(self):
         value = self.faker.unique.name()
@@ -2028,11 +2045,12 @@ class BlockRunProjectPlatformTests(ProjectPlatformMixin, AsyncDBTestCase):
                 return fake_logger_manager
 
         with patch(
-            'mage_ai.orchestration.db.models.schedules.get_pipeline_from_platform_async',
+            'mage_ai.orchestration.db.models.schedules_project_platform.' +
+            'get_pipeline_from_platform_async',
             __get_pipeline_from_platform_async,
         ):
             with patch(
-                'mage_ai.orchestration.db.models.schedules.LoggerManagerFactory',
+                'mage_ai.orchestration.db.models.schedules_project_platform.LoggerManagerFactory',
                 FakeLoggerManagerFactory,
             ):
                 await block_run.logs_async()