mage-ai 0.9.67__py3-none-any.whl → 0.9.69__py3-none-any.whl

mage_ai/shared/files.py

@@ -1,7 +1,7 @@
 import glob
 import os
 from pathlib import Path
-from typing import Callable, List, Tuple
+from typing import Callable, Dict, List, Tuple
 
 import aiofiles
 
@@ -63,6 +63,24 @@ def get_full_file_paths_containing_item(root_full_path: str, comparator: Callabl
     return configfiles
 
 
+def get_full_file_paths_containing_multi_items(
+    root_full_path: str,
+    comparators: Dict[str, Callable],
+    exclude_hidden_dir: bool = False,
+) -> Dict[str, List]:
+    configfiles = dict()
+    for key, _ in comparators.items():
+        configfiles[key] = []
+    for dirpath, dirnames, files in os.walk(root_full_path):
+        if exclude_hidden_dir:
+            dirnames[:] = [d for d in dirnames if not d.startswith('.')]
+        for f in files:
+            for key, comparator in comparators.items():
+                if comparator(f):
+                    configfiles[key].append(os.path.join(dirpath, f))
+    return configfiles
+
+
 def find_directory(top_level_path: str, comparator: Callable) -> str:
     for path, _subdirs, files in os.walk(top_level_path):
         for name in files:

mage_ai/shared/path_fixer.py

@@ -29,6 +29,9 @@ def add_absolute_path(file_path: str, add_base_repo_path: bool = True) -> str:
                 full_path = os.path.join(repo_path, parts[0])
                 if os.path.exists(full_path):
                     full_path = os.path.join(repo_path, file_path)
+                elif os.path.exists(os.path.join(repo_path, 'dbt', parts[0])):
+                    # DBT v1 paths
+                    full_path = os.path.join(repo_path, 'dbt', file_path)
                 else:
                     full_path = find_directory(repo_path, lambda x: x.endswith(file_path))
 

mage_ai/streaming/sources/base.py

@@ -3,6 +3,8 @@ from abc import ABC, abstractmethod
 from enum import Enum
 from typing import Callable, Dict
 
+from mage_ai.shared.environments import is_test
+
 
 class SourceConsumeMethod(str, Enum):
     BATCH_READ = 'BATCH_READ'
@@ -22,6 +24,9 @@ class BaseSource(ABC):
         self.checkpoint_path = kwargs.get('checkpoint_path')
         self.checkpoint = self.read_checkpoint()
         self.init_client()
+        if not is_test():
+            # Not test the connection in unit tests
+            self.test_connection()
 
     @abstractmethod
     def init_client(self):

mage_ai/streaming/sources/influxdb.py

@@ -164,6 +164,7 @@ class InfluxDbSource(BaseSource):
                         'data': {record.get_field(): record.get_value()},
                         'metadata': {
                             'time': int(1e3 * record.get_time().timestamp()),
+                            'measurement': record['_measurement'],
                             'tags': {
                                 k: v
                                 for k, v in record.values.items()
@@ -201,6 +202,7 @@ class InfluxDbSource(BaseSource):
                         'data': {record.get_field(): record.get_value()},
                         'metadata': {
                             'time': int(1e3 * record.get_time().timestamp()),
+                            'measurement': record['_measurement'],
                             'tags': {
                                 k: v
                                 for k, v in record.values.items()

mage_ai/streaming/sources/kafka.py

@@ -256,7 +256,8 @@ class KafkaSource(BaseSource):
                 handler(message_values)
 
     def test_connection(self):
-        return True
+        self.consumer._client.check_version(timeout=5)
+        self._print('Test connection successfully.')
 
     def __deserialize_message(self, message):
         if self.config.serde_config is None:

mage_ai/streaming/sources/mongodb.py

@@ -18,6 +18,7 @@ class MongoDBConfig(BaseConfig):
     pipeline: Optional[_Pipeline] = None
     operation_time: Optional[Timestamp] = None
     start_after: Optional[Mapping[str, Any]] = None
+    full_document: Optional[str] = None
 
 
 class MongoSource(BaseSource):
@@ -46,6 +47,9 @@ class MongoSource(BaseSource):
         if hasattr(self.config, 'start_after') and self.config.start_after is not None:
             watch_args['start_after'] = self.config.start_after
 
+        if hasattr(self.config, 'full_document'):
+            watch_args['full_document'] = self.config.full_document
+
         return watch_args
 
     def batch_read(self, handler: Callable):

mage_ai/tests/api/endpoints/mixins.py

@@ -813,7 +813,7 @@ class BaseAPIEndpointTest(AsyncDBTestCase):
                             validation = assert_after_create_count(self)
                         self.assertTrue(validation)
                     else:
-                        after_count = len(get_resource().model_class.all())
+                        after_count = len(get_resource(resource).model_class.all())
                         self.assertEqual(
                             after_count,
                             before_count + 1 if after_create_count is None else after_create_count,
@@ -1096,7 +1096,7 @@ class BaseAPIEndpointTest(AsyncDBTestCase):
 
                         self.assertTrue(validation)
                     else:
-                        after_count = len(get_resource().model_class.all())
+                        after_count = len(get_resource(resource).model_class.all())
                         self.assertEqual(
                             after_count,
                             before_count - 1 if after_delete_count is None else after_delete_count,
@@ -1132,15 +1132,16 @@ class BaseAPIEndpointTest(AsyncDBTestCase):
                 )
                 RolePermission.create(permission_id=permission.id, role_id=role.id)
             else:
-                permission = Permission.create(
-                    access=access_for_permissions,
-                    entity_name=EntityName(entity_name(resource)),
-                    options=permission_options,
-                )
+                e_name = entity_name(resource)
+                if getattr(EntityName, e_name, None) is not None:
+                    permission = Permission.create(
+                        access=access_for_permissions,
+                        entity_name=EntityName(e_name),
+                        options=permission_options,
+                    )
+                    RolePermission.create(permission_id=permission.id, role_id=role.id)
 
                 if permission_settings:
                     for permission_setting in permission_settings:
                         permission_more = Permission.create(**permission_setting)
                         RolePermission.create(permission_id=permission_more.id, role_id=role.id)
-
-            RolePermission.create(permission_id=permission.id, role_id=role.id)

mage_ai/tests/api/endpoints/test_seeds.py

@@ -0,0 +1,24 @@
+from mage_ai.authentication.permissions.constants import PermissionAccess
+from mage_ai.tests.api.endpoints.mixins import (
+    BaseAPIEndpointTest,
+    build_create_endpoint_tests,
+)
+
+
+class SeedAPIEndpointTest(BaseAPIEndpointTest):
+    pass
+
+
+build_create_endpoint_tests(
+    SeedAPIEndpointTest,
+    resource='seed',
+    authentication_accesses=[PermissionAccess.OWNER],
+    permissions_accesses=[PermissionAccess.OWNER],
+    build_payload=lambda self: dict(
+        roles=True,
+        permissions=True,
+        policy_names=['Block', 'Pipeline', 'Download'],
+    ),
+    assert_before_create_count=lambda self: True,
+    assert_after_create_count=lambda self: True,
+)

mage_ai/tests/api/operations/test_sessions.py

@@ -178,7 +178,9 @@ class SessionOperationTests(BaseApiTestCase):
     @patch('mage_ai.api.resources.SessionResource.AUTHENTICATION_MODE', 'ldap')
     @patch.object(LDAPConnection, 'authorize')
     @patch.object(LDAPConnection, 'authenticate')
-    async def test_ldap_login_with_role_mapping(self, mock_authenticate, mock_authorize):
+    async def test_ldap_login_with_role_mapping(
+        self, mock_authenticate, mock_authorize
+    ):
         mock_authenticate.return_value = (
             True,
             "Yami_Sukehiro",
@@ -188,7 +190,9 @@ class SessionOperationTests(BaseApiTestCase):
 
         Role.create_default_roles()
 
-        with patch.dict(os.environ, dict(LDAP_ROLES_MAPPING=json.dumps(dict(Admin=['Admin'])))):
+        with patch.dict(
+            os.environ, dict(LDAP_ROLES_MAPPING=json.dumps(dict(Admin=['Admin'])))
+        ):
             username = self.faker.email()
             operation = self.build_operation(
                 action=constants.CREATE,
@@ -208,3 +212,50 @@ class SessionOperationTests(BaseApiTestCase):
             user = User.query.filter(User.username == username).first()
             self.assertIsNotNone(user)
             self.assertEqual(user.roles_new[0].name, 'Admin')
+
+    @patch('mage_ai.api.resources.SessionResource.AUTHENTICATION_MODE', 'ldap')
+    @patch.object(LDAPConnection, 'authorize')
+    @patch.object(LDAPConnection, 'authenticate')
+    async def test_ldap_update_roles_on_login(self, mock_authenticate, mock_authorize):
+        mock_authenticate.return_value = (
+            True,
+            "Yami_Sukehiro",
+            dict(memberOf=['Admin']),
+        )
+        mock_authorize.return_value = True
+
+        Role.create_default_roles()
+
+        username = self.faker.email()
+
+        User.create(
+            username=username,
+            email=username,
+            roles_new=[Role.query.filter(Role.name == 'Admin').first()],
+        )
+
+        with patch.dict(
+            os.environ,
+            dict(
+                LDAP_ROLES_MAPPING=json.dumps(dict(Admin=['Editor'])),
+                UPDATE_ROLES_ON_LOGIN='1',
+            ),
+        ):
+            operation = self.build_operation(
+                action=constants.CREATE,
+                payload=dict(
+                    session=dict(
+                        email=username,
+                        password="black bull",
+                    )
+                ),
+                resource='sessions',
+                user=None,
+            )
+            await operation.execute()
+
+            mock_authenticate.assert_called_once_with(username, "black bull")
+            mock_authorize.assert_called_once_with("Yami_Sukehiro")
+            user = User.query.filter(User.username == username).first()
+            self.assertIsNotNone(user)
+            self.assertEqual(user.roles_new[0].name, 'Editor')

mage_ai/tests/authentication/providers/test_okta.py

@@ -0,0 +1,43 @@
+import os
+import unittest
+from unittest.mock import patch
+
+from mage_ai.authentication.providers.okta import OktaProvider
+
+test_parameters = [
+    ('https://samples.auth0.com', 'test-client-id', 'test-client-secret'),
+    ('samples.auth0.com', 'test-client-id', 'test-client-secret'),
+    ('', 'test-client-id', 'test-client-secret'),
+    ('samples.auth0.com', '', ''),
+]
+
+
+class OktaProviderTest(unittest.TestCase):
+    def setUp(self):
+        self.authorization_endpoint = 'https://samples.auth0.com/authorize'
+        self.token_endpoint = 'https://samples.auth0.com/oauth/token'
+        self.userinfo_endpoint = 'https://samples.auth0.com/userinfo'
+
+    def test_okta_provider_initialization(self):
+        for url, id, secret in test_parameters:
+            with self.subTest():
+                with patch.dict(
+                    os.environ,
+                    dict(
+                        OKTA_DOMAIN_URL=url,
+                        OKTA_CLIENT_ID=id,
+                        OKTA_CLIENT_SECRET=secret,
+                    ),
+                ):
+                    if not all([url, id]):
+                        with self.assertRaises(ValueError):
+                            provider = OktaProvider()
+                    else:
+                        provider = OktaProvider()
+                        self.assertEqual(
+                            provider.authorization_endpoint, self.authorization_endpoint
+                        )
+                        self.assertEqual(provider.token_endpoint, self.token_endpoint)
+                        self.assertEqual(
+                            provider.userinfo_endpoint, self.userinfo_endpoint
+                        )

mage_ai/tests/data_preparation/models/block/dbt/test_block.py

@@ -38,7 +38,7 @@ class DBTBlockTest(AsyncDBTestCase):
         self.block_sql.language = BlockLanguage.SQL
         self.block_sql.pipeline = pipeline
 
-        self.dbt_block_sql = DBTBlock(
+        self.dbt_block_sql = DBTBlock.create(
             name='test_dbt_block_sql',
             uuid='test_dbt_block_sql',
             block_type=BlockType.DBT,
@@ -50,7 +50,7 @@ class DBTBlockTest(AsyncDBTestCase):
             pipeline=pipeline,
         )
 
-        self.dbt_block_yaml = DBTBlock(
+        self.dbt_block_yaml = DBTBlock.create(
             name='test_dbt_block_yaml',
             uuid='test_dbt_block_yaml',
             block_type=BlockType.DBT,

mage_ai/tests/data_preparation/models/block/dbt/test_block_sql.py

@@ -23,7 +23,7 @@
 #     pipeline.repo_path = 'test_repo_path'
 #     pipeline.get_block.return_value = None
 
-#     return DBTBlock(
+#     return DBTBlock.create(
 #         name='test_dbt_block_sql',
 #         uuid='test_dbt_block_sql',
 #         block_type=BlockType.DBT,

mage_ai/tests/data_preparation/models/block/dbt/test_block_yaml.py

@@ -18,7 +18,7 @@
 
 
 # def build_block(pipeline, content: str) -> DBTBlock:
-#     return DBTBlock(
+#     return DBTBlock.create(
 #         name='test_dbt_block_yaml',
 #         uuid='test_dbt_block_yaml',
 #         block_type=BlockType.DBT,

mage_ai/tests/data_preparation/models/test_block.py

@@ -8,6 +8,7 @@ from pandas.testing import assert_frame_equal
 
 # from mage_ai.data_cleaner.column_types.constants import ColumnType
 from mage_ai.data_preparation.models.block import Block, BlockType, CallbackBlock
+from mage_ai.data_preparation.models.block.block_factory import BlockFactory
 from mage_ai.data_preparation.models.block.errors import HasDownstreamDependencies
 from mage_ai.data_preparation.models.pipeline import Pipeline
 from mage_ai.data_preparation.repo_manager import get_repo_config
@@ -606,7 +607,7 @@ SELECT 1 AS id;
             upstream_block_uuids=['test_data_loader'],
         )
 
-        test_block = Block.get_block('test_data_loader', 'test_data_loader', 'data_loader')
+        test_block = BlockFactory.get_block('test_data_loader', 'test_data_loader', 'data_loader')
 
         self.assertRaises(HasDownstreamDependencies, block1.delete)
         self.assertRaises(HasDownstreamDependencies, test_block.delete)

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

@@ -17,7 +17,7 @@ class RoleTests(DBTestCase):
         Role.create_default_roles(
             entity=Entity.PROJECT,
             entity_id=test_entity_id,
-            prefix='test',
+            name_func=lambda x: f'test_{x}',
         )
         owner = Role.query.filter(Role.name == 'test_Owner').one_or_none()
 
@@ -29,12 +29,12 @@ class RoleTests(DBTestCase):
         Role.create_default_roles(
             entity=Entity.PROJECT,
             entity_id=test_entity_id,
-            prefix='test-2',
+            name_func=lambda x: f'test-2_{x}',
         )
         Role.create_default_roles(
             entity=Entity.PROJECT,
             entity_id=test_entity_id2,
-            prefix='test-2',
+            name_func=lambda x: f'test-2_{x}',
         )
         owner = Role.query.filter(Role.name == 'test-2_Owner').one_or_none()
 

mage_ai/tests/orchestration/queue/test_process_queue.py

@@ -13,6 +13,7 @@ class ProcessQueueTests(TestCase):
     def setUp(self):
         queue_config = QueueConfig.load(config=dict(concurrency=100))
         self.queue = ProcessQueue(queue_config=queue_config)
+        self.queue.start()
 
     def test_init(self):
         self.assertEqual(self.queue.size, 100)

mage_ai/tests/orchestration/test_pipeline_scheduler.py

@@ -396,6 +396,7 @@ class PipelineSchedulerTests(DBTestCase):
                 error=f'Failed blocks: {block_runs[0].block_uuid}.',
                 pipeline=scheduler.pipeline,
                 pipeline_run=pipeline_run,
+                stacktrace=None,
             )
 
     def test_schedule_with_block_failures(self):
@@ -419,6 +420,7 @@ class PipelineSchedulerTests(DBTestCase):
                 error=f'Failed blocks: {block_runs[0].block_uuid}.',
                 pipeline=scheduler.pipeline,
                 pipeline_run=pipeline_run,
+                stacktrace=None,
             )
 
     @patch('mage_ai.orchestration.pipeline_scheduler_original.run_pipeline')

mage_ai/tests/server/test_server.py

@@ -8,7 +8,7 @@ import tornado.ioloop
 
 import mage_ai.server.server as server_module
 from mage_ai.authentication.passwords import create_bcrypt_hash, generate_salt
-from mage_ai.data_preparation.repo_manager import ProjectType
+from mage_ai.data_preparation.repo_manager import ProjectType, get_project_uuid
 from mage_ai.orchestration.db.models.oauth import (
     Oauth2Application,
     Permission,
@@ -21,7 +21,6 @@ from mage_ai.server.server import (
     make_app,
     replace_base_path,
 )
-from mage_ai.settings.repo import get_repo_name
 from mage_ai.tests.base_test import AsyncDBTestCase
 
 
@@ -150,7 +149,10 @@ class ServerTests(AsyncDBTestCase):
     def test_initialize_user_authentication_subproject(self):
         initialize_user_authentication(ProjectType.SUB)
 
-        owner_role = Role.get_role(f'{get_repo_name()}_{Role.DefaultRole.OWNER}')
+        project_uuid = get_project_uuid()
+        project_uuid_truncated = project_uuid[:8]
+
+        owner_role = Role.get_role(f'{Role.DefaultRole.OWNER}_{project_uuid_truncated}')
         self.assertTrue(len(owner_role.users) > 0)
 
         owner_user = User.query.filter(User.email == 'admin@admin.com').one_or_none()
@@ -169,7 +171,9 @@ class ServerTests(AsyncDBTestCase):
         )
         initialize_user_authentication(ProjectType.SUB)
 
-        owner_role = Role.get_role(f'{get_repo_name()}_{Role.DefaultRole.OWNER}')
+        project_uuid = get_project_uuid()
+        project_uuid_truncated = project_uuid[:8]
+        owner_role = Role.get_role(f'{Role.DefaultRole.OWNER}_{project_uuid_truncated}')
         self.assertTrue(len(owner_role.users) == 0)
 
         owner_user = User.query.filter(User.email == 'admin@admin.com').one_or_none()

mage_ai/tests/settings/models/test_configuration_option.py

@@ -100,7 +100,7 @@ class ConfigurationOptionTest(AsyncDBTestCase):
             self.faker.unique.name(),
             os.path.join(self.repo_path, 'mage_platform'),
         )
-        block = DBTBlock(
+        block = DBTBlock.create(
             block_type='dbt',
             configuration=dict(file_path='dir2/mage.sql'),
             language='sql',
@@ -139,7 +139,7 @@ class ConfigurationOptionTest(AsyncDBTestCase):
                         self.faker.unique.name(),
                         os.path.join(self.repo_path, 'mage_platform'),
                     )
-                    block = DBTBlock(
+                    block = DBTBlock.create(
                         block_type='dbt',
                         configuration=dict(file_path='dir2/mage.sql'),
                         language='sql',

{mage_ai-0.9.67.dist-info → mage_ai-0.9.69.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: mage-ai
-Version: 0.9.67
+Version: 0.9.69
 Summary: Mage is a tool for building and deploying data pipelines.
 Home-page: https://github.com/mage-ai/mage-ai
 Author: Mage
@@ -24,7 +24,7 @@ Requires-Dist: cryptography ==41.0.6
 Requires-Dist: dask >=2022.2.0
 Requires-Dist: datadog ==0.44.0
 Requires-Dist: freezegun ==1.2.2
-Requires-Dist: great-expectations ==0.15.50
+Requires-Dist: great-expectations ==0.18.12
 Requires-Dist: httpx ==0.23.1
 Requires-Dist: inflection ==0.5.1
 Requires-Dist: ipykernel ==6.15.0
@@ -57,7 +57,7 @@ Requires-Dist: terminado ==0.17.1
 Requires-Dist: thefuzz[speedup] ==0.19.0
 Requires-Dist: tornado ==6.3.3
 Requires-Dist: typer[all] ==0.9.0
-Requires-Dist: typing-extensions ==4.5.0
+Requires-Dist: typing-extensions ==4.10.0
 Requires-Dist: watchdog ==4.0.0
 Provides-Extra: ai
 Requires-Dist: astor >=0.8.1 ; extra == 'ai'
@@ -98,7 +98,7 @@ Requires-Dist: google-cloud-iam ~=2.13.0 ; extra == 'all'
 Requires-Dist: google-cloud-pubsub ~=2.19.0 ; extra == 'all'
 Requires-Dist: google-cloud-run ~=0.10.1 ; extra == 'all'
 Requires-Dist: google-cloud-storage ~=2.5.0 ; extra == 'all'
-Requires-Dist: great-expectations ==0.15.50 ; extra == 'all'
+Requires-Dist: great-expectations ==0.18.12 ; 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'
@@ -128,7 +128,7 @@ Requires-Dist: pyodbc ==4.0.35 ; extra == 'all'
 Requires-Dist: redshift-connector ==2.0.915 ; extra == 'all'
 Requires-Dist: lxml ==4.9.4 ; extra == 'all'
 Requires-Dist: requests-aws4auth ==1.1.2 ; extra == 'all'
-Requires-Dist: snowflake-connector-python ==3.4.0 ; extra == 'all'
+Requires-Dist: snowflake-connector-python ==3.5.0 ; extra == 'all'
 Requires-Dist: sshtunnel ==0.4.0 ; extra == 'all'
 Requires-Dist: stomp.py ==8.1.0 ; extra == 'all'
 Requires-Dist: thefuzz[speedup] ==0.19.0 ; extra == 'all'
@@ -183,7 +183,7 @@ Provides-Extra: s3
 Requires-Dist: boto3 ==1.26.60 ; extra == 's3'
 Requires-Dist: botocore ==1.29.60 ; extra == 's3'
 Provides-Extra: snowflake
-Requires-Dist: snowflake-connector-python ==3.4.0 ; extra == 'snowflake'
+Requires-Dist: snowflake-connector-python ==3.5.0 ; extra == 'snowflake'
 Provides-Extra: spark
 Requires-Dist: boto3 ==1.26.60 ; extra == 'spark'
 Requires-Dist: botocore ==1.29.60 ; extra == 'spark'