mage-ai 0.9.73__py3-none-any.whl → 0.9.75__py3-none-any.whl

mage_ai/streaming/sources/kafka.py

@@ -1,19 +1,25 @@
 import importlib
 import json
+import sys
+
+if sys.version_info >= (3, 12, 0):
+    import six
+    sys.modules['kafka.vendor.six.moves'] = six.moves
+
 import time
 from dataclasses import dataclass, field
-from enum import Enum
 from typing import Callable, Dict, List
 
 from kafka import KafkaConsumer, TopicPartition
 
 from mage_ai.shared.config import BaseConfig
+from mage_ai.shared.enum import StrEnum
 from mage_ai.streaming.constants import DEFAULT_BATCH_SIZE, DEFAULT_TIMEOUT_MS
 from mage_ai.streaming.sources.base import BaseSource
 from mage_ai.streaming.sources.shared import SerDeConfig, SerializationMethod
 
 
-class SecurityProtocol(str, Enum):
+class SecurityProtocol(StrEnum):
     SASL_PLAINTEXT = 'SASL_PLAINTEXT'
     SASL_SSL = 'SASL_SSL'
     SSL = 'SSL'
@@ -102,8 +108,18 @@ class KafkaSource(BaseSource):
             consumer_kwargs['sasl_plain_username'] = self.config.sasl_config.username
             consumer_kwargs['sasl_plain_password'] = self.config.sasl_config.password
 
-            if self.config.ssl_config is not None and self.config.ssl_config.cafile:
-                consumer_kwargs['ssl_cafile'] = self.config.ssl_config.cafile
+            if self.config.ssl_config:
+                consumer_kwargs[
+                    'ssl_check_hostname'] = self.config.ssl_config.check_hostname
+                if self.config.ssl_config.cafile:
+                    consumer_kwargs['ssl_cafile'] = self.config.ssl_config.cafile
+                if self.config.ssl_config.certfile:
+                    consumer_kwargs['ssl_certfile'] = self.config.ssl_config.certfile
+                if self.config.ssl_config.keyfile:
+                    consumer_kwargs['ssl_keyfile'] = self.config.ssl_config.keyfile
+                if self.config.ssl_config.password:
+                    consumer_kwargs['ssl_password'] = self.config.ssl_config.password
+
         elif self.config.security_protocol == SecurityProtocol.SASL_PLAINTEXT:
             consumer_kwargs['security_protocol'] = SecurityProtocol.SASL_PLAINTEXT
             consumer_kwargs['sasl_mechanism'] = self.config.sasl_config.mechanism

mage_ai/streaming/sources/nats_js.py

@@ -1,18 +1,20 @@
 import asyncio
-import enum
 import json
 import ssl
 import threading
 from dataclasses import dataclass
-from typing import Callable, Dict, Optional
+from typing import Callable, Dict, Optional, Tuple, Union
 
 import nats
 from nats.errors import NoServersError, TimeoutError
 
 from mage_ai.shared.config import BaseConfig
+from mage_ai.shared.enum import StrEnum
 from mage_ai.streaming.constants import DEFAULT_BATCH_SIZE, DEFAULT_TIMEOUT_MS
 from mage_ai.streaming.sources.base import BaseSource, SourceConsumeMethod
 
+Credentials = Union[str, Tuple[str, str]]
+
 
 @dataclass
 class SSLConfig:
@@ -22,7 +24,7 @@ class SSLConfig:
     check_hostname: bool = False
 
 
-class ConsumerType(str, enum.Enum):
+class ConsumerType(StrEnum):
     PULL = "PULL"
     PUSH = "PUSH"
 
@@ -33,6 +35,7 @@ class NATSConfig(BaseConfig):
     stream_name: str
     subject: str = None
     nkeys_seed_str: Optional[str] = None
+    user_credentials: Optional[Credentials] = None
     use_tls: bool = False
     ssl_config: Optional[SSLConfig] = None
     consumer_name: Optional[str] = None
@@ -99,6 +102,10 @@ class NATSSource(BaseSource):
             if self.config.nkeys_seed_str:
                 connect_opts["nkeys_seed_str"] = self.config.nkeys_seed_str
 
+            # Use credentials if provided
+            if self.config.user_credentials:
+                connect_opts["user_credentials"] = self.config.user_credentials
+
             # Establish connection with the configured options
             self.nc = await nats.connect(**connect_opts)
             self.js = self.nc.jetstream()

mage_ai/streaming/sources/shared.py

@@ -1,8 +1,9 @@
 from dataclasses import dataclass
-from enum import Enum
 
+from mage_ai.shared.enum import StrEnum
 
-class SerializationMethod(str, Enum):
+
+class SerializationMethod(StrEnum):
     AVRO = 'AVRO'
     JSON = 'JSON'
     PROTOBUF = 'PROTOBUF'

mage_ai/system/constants.py

@@ -1,7 +1,7 @@
-from enum import Enum
+from mage_ai.shared.enum import StrEnum
 
 
-class LogType(str, Enum):
+class LogType(StrEnum):
     END = 'END'
     GENERIC = 'GENERIC'
     MEMORY = 'MEMORY'

mage_ai/tests/api/endpoints/test_blocks.py

@@ -1,6 +1,4 @@
-import os
-
-from mage_ai.data_preparation.models.constants import BlockLanguage, BlockType
+from mage_ai.data_preparation.models.constants import BlockType
 from mage_ai.data_preparation.models.pipeline import Pipeline
 from mage_ai.settings.repo import get_repo_path
 from mage_ai.tests.api.endpoints.mixins import (
@@ -11,7 +9,6 @@ from mage_ai.tests.api.endpoints.mixins import (
     build_list_endpoint_tests,
     build_update_endpoint_tests,
 )
-from mage_ai.tests.shared.mixins import DBTMixin
 
 
 def get_pipeline(self):
@@ -109,100 +106,3 @@ build_delete_endpoint_tests(
     resource='block',
     resource_parent='pipelines',
 )
-
-
-class DBTBlockAPIEndpointTest(BaseAPIEndpointTest, DBTMixin):
-    pass
-
-
-def __assert_after_create(self):
-    pipeline = get_pipeline(self)
-    blocks = pipeline.blocks_by_uuid.values()
-    blocks_dbt = list(filter(
-        lambda x: x.type == BlockType.DBT and x.language == BlockLanguage.SQL,
-        blocks,
-    ))
-
-    self.assertEqual(len(blocks), 5)
-    self.assertEqual(len(blocks_dbt), 1)
-
-    block = blocks_dbt[0]
-    self.assertEqual({
-      'dbt_project_name': 'dbt/test_project',
-      'file_path': 'dbt/test_project/models/example/my_first_dbt_model.sql',
-      'file_source': {
-        'path': 'dbt/test_project/models/example/my_first_dbt_model.sql',
-        'project_path': 'dbt/test_project',
-      },
-    }, block.configuration)
-
-    return True
-
-
-def __assert_after_create_yaml(self):
-    pipeline = get_pipeline(self)
-    blocks = pipeline.blocks_by_uuid.values()
-    blocks_dbt = list(filter(
-        lambda x: x.type == BlockType.DBT and x.language == BlockLanguage.YAML,
-        blocks,
-    ))
-
-    self.assertEqual(len(blocks), 5)
-    self.assertEqual(len(blocks_dbt), 1)
-
-    block = blocks_dbt[0]
-    self.assertEqual({
-      'dbt_project_name': 'dbt/test_project',
-      'file_source': {
-        'path': 'dbts/dbt_yaml.yaml',
-        'project_path': 'dbt/test_project',
-      },
-    }, block.configuration)
-
-    return True
-
-
-build_create_endpoint_tests(
-    DBTBlockAPIEndpointTest,
-    assert_after_create_count=__assert_after_create,
-    assert_before_create_count=lambda self: len(get_pipeline(self).blocks_by_uuid.values()) == 4,
-    build_payload=lambda self: dict(
-        name=self.faker.unique.name(),
-        type=BlockType.DBT,
-        language=BlockLanguage.SQL,
-        configuration=dict(
-            file_path=os.path.join(
-                os.path.basename(self.dbt_directory),
-                'test_project',
-                'models',
-                'example',
-                'my_first_dbt_model.sql',
-            ),
-        ),
-    ),
-    get_resource_parent_id=lambda self: self.pipeline.uuid,
-    resource='block',
-    resource_parent='pipelines',
-)
-
-
-build_create_endpoint_tests(
-    DBTBlockAPIEndpointTest,
-    test_uuid='dbt_yaml',
-    assert_after_create_count=__assert_after_create_yaml,
-    assert_before_create_count=lambda self: len(get_pipeline(self).blocks_by_uuid.values()) == 4,
-    build_payload=lambda self: dict(
-        name='dbt_yaml',
-        type=BlockType.DBT,
-        language=BlockLanguage.YAML,
-        configuration=dict(
-            dbt_project_name=os.path.join(
-                os.path.basename(self.dbt_directory),
-                'test_project',
-            ),
-        ),
-    ),
-    get_resource_parent_id=lambda self: self.pipeline.uuid,
-    resource='block',
-    resource_parent='pipelines',
-)

mage_ai/tests/api/endpoints/test_configuration_options.py

@@ -1,6 +1,5 @@
 import json
 import os
-from unittest.mock import patch
 
 from mage_ai.authentication.permissions.constants import EntityName
 from mage_ai.settings.models.configuration_option import ConfigurationType, OptionType
@@ -8,7 +7,7 @@ from mage_ai.tests.api.endpoints.mixins import (
     BaseAPIEndpointTest,
     build_list_endpoint_tests,
 )
-from mage_ai.tests.shared.mixins import DBTMixin, ProjectPlatformMixin
+from mage_ai.tests.shared.mixins import DBTMixin
 
 CURRENT_FILE_PATH = os.path.dirname(os.path.realpath(__file__))
 
@@ -51,49 +50,3 @@ build_list_endpoint_tests(
     assert_after=__assert_after_list,
     build_query=__build_query,
 )
-
-
-@patch('mage_ai.settings.platform.utils.project_platform_activated', lambda: True)
-@patch('mage_ai.settings.repo.project_platform_activated', lambda: True)
-class ConfigurationOptionProjectPlatformAPIEndpointTest(
-    DBTMixin,
-    ProjectPlatformMixin,
-    BaseAPIEndpointTest,
-):
-    pass
-
-
-def __assert_after_list2(self, result, **kwargs):
-    with open(os.path.join(
-        CURRENT_FILE_PATH,
-        'mocks',
-        'mock_dbt_configuration_options_project_platform.json',
-    )) as f:
-        self.assertEqual(
-            result,
-            json.loads(f.read()),
-        )
-
-
-build_list_endpoint_tests(
-    ConfigurationOptionProjectPlatformAPIEndpointTest,
-    list_count=1,
-    get_resource_parent_id=lambda test_case: test_case.pipeline.uuid,
-    resource='configuration_option',
-    resource_parent='pipeline',
-    result_keys_to_compare=[
-        'configuration_type',
-        'metadata',
-        'option',
-        'option_type',
-        'resource_type',
-        'uuid',
-    ],
-    assert_after=__assert_after_list2,
-    build_query=__build_query,
-    patch_function_settings=[
-        ('mage_ai.settings.models.configuration_option.project_platform_activated', lambda: True),
-        ('mage_ai.settings.platform.utils.project_platform_activated', lambda: True),
-        ('mage_ai.settings.repo.project_platform_activated', lambda: True),
-    ],
-)

mage_ai/tests/api/endpoints/test_configuration_options_project_platform.py

@@ -0,0 +1,68 @@
+import json
+import os
+from unittest.mock import patch
+
+from mage_ai.authentication.permissions.constants import EntityName
+from mage_ai.settings.models.configuration_option import ConfigurationType, OptionType
+from mage_ai.tests.api.endpoints.mixins import (
+    BaseAPIEndpointTest,
+    build_list_endpoint_tests,
+)
+from mage_ai.tests.shared.mixins import DBTMixin, ProjectPlatformMixin
+
+CURRENT_FILE_PATH = os.path.dirname(os.path.realpath(__file__))
+
+
+def __build_query(test_case):
+    return dict(
+        configuration_type=[ConfigurationType.DBT],
+        option_type=[OptionType.PROJECTS],
+        resource_type=[EntityName.Block],
+        resource_uuid=[test_case.blocks[0].uuid],
+    )
+
+
+@patch('mage_ai.settings.platform.utils.project_platform_activated', lambda: True)
+@patch('mage_ai.settings.repo.project_platform_activated', lambda: True)
+class ConfigurationOptionProjectPlatformAPIEndpointTest(
+    DBTMixin,
+    ProjectPlatformMixin,
+    BaseAPIEndpointTest,
+):
+    pass
+
+
+def __assert_after_list2(self, result, **kwargs):
+    with open(os.path.join(
+        CURRENT_FILE_PATH,
+        'mocks',
+        'mock_dbt_configuration_options_project_platform.json',
+    )) as f:
+        self.assertEqual(
+            result,
+            json.loads(f.read()),
+        )
+
+
+build_list_endpoint_tests(
+    ConfigurationOptionProjectPlatformAPIEndpointTest,
+    list_count=1,
+    get_resource_parent_id=lambda test_case: test_case.pipeline.uuid,
+    resource='configuration_option',
+    resource_parent='pipeline',
+    result_keys_to_compare=[
+        'configuration_type',
+        'metadata',
+        'option',
+        'option_type',
+        'resource_type',
+        'uuid',
+    ],
+    assert_after=__assert_after_list2,
+    build_query=__build_query,
+    patch_function_settings=[
+        ('mage_ai.settings.models.configuration_option.project_platform_activated', lambda: True),
+        ('mage_ai.settings.platform.utils.project_platform_activated', lambda: True),
+        ('mage_ai.settings.repo.project_platform_activated', lambda: True),
+    ],
+)

mage_ai/tests/api/endpoints/test_custom_designs.py

@@ -1,15 +1,13 @@
 import json
 import os
-from unittest.mock import patch
 
 from mage_ai.data_preparation.models.project.constants import FeatureUUID
-from mage_ai.settings.utils import base_repo_path
 from mage_ai.tests.api.endpoints.mixins import (
     BaseAPIEndpointTest,
     build_detail_endpoint_tests,
     build_list_endpoint_tests,
 )
-from mage_ai.tests.shared.mixins import CustomDesignMixin, ProjectPlatformMixin
+from mage_ai.tests.shared.mixins import CustomDesignMixin
 
 CURRENT_FILE_PATH = os.path.dirname(os.path.realpath(__file__))
 
@@ -102,106 +100,3 @@ build_detail_endpoint_tests(
         ),
     ],
 )
-
-
-@patch('mage_ai.presenters.design.models.project_platform_activated', lambda: True)
-@patch('mage_ai.settings.platform.utils.project_platform_activated', lambda: True)
-@patch('mage_ai.settings.repo.project_platform_activated', lambda: True)
-class CustomDesignProjectPlatformAPIEndpointTest(
-    CustomDesignMixin,
-    ProjectPlatformMixin,
-    BaseAPIEndpointTest,
-):
-    pass
-
-
-def __assert_after_list2(self, result, **kwargs):
-    with open(os.path.join(CURRENT_FILE_PATH, 'mocks', 'mock_design_project_platform.json')) as f:
-        arr = json.loads(f.read())
-        for item in arr:
-            item['project']['full_path'] = os.path.join(
-                base_repo_path(),
-                item['project']['full_path'],
-            )
-            item['project']['root_project_full_path'] = base_repo_path()
-
-        self.assertEqual(
-            result.sort(key=lambda x: x['uuid']),
-            arr.sort(key=lambda x: x['uuid']),
-        )
-
-
-def __assert_after_detail2(self, result, **kwargs):
-    with open(os.path.join(CURRENT_FILE_PATH, 'mocks', 'mock_design_project_platform.json')) as f:
-        item = json.loads(f.read())[0]
-        item['project']['full_path'] = os.path.join(
-            base_repo_path(),
-            item['project']['full_path'],
-        )
-        item['project']['root_project_full_path'] = base_repo_path()
-
-        self.assertEqual(
-            result,
-            item,
-        )
-
-
-build_list_endpoint_tests(
-    CustomDesignProjectPlatformAPIEndpointTest,
-    list_count=2,
-    resource='custom_design',
-    result_keys_to_compare=[
-        'components',
-        'pages',
-        'project',
-        'uuid',
-    ],
-    assert_after=__assert_after_list2,
-    build_query=__build_query,
-    patch_function_settings=[
-        (
-            '.'.join([
-                'mage_ai',
-                'api',
-                'resources',
-                'CustomDesignResource',
-                'Project',
-                'is_feature_enabled_in_root_or_active_project',
-            ]),
-            lambda feature_name, user=None, **kwargs: FeatureUUID.CUSTOM_DESIGN == feature_name,
-        ),
-        ('mage_ai.presenters.design.models.project_platform_activated', lambda: True),
-        ('mage_ai.settings.platform.utils.project_platform_activated', lambda: True),
-        ('mage_ai.settings.repo.project_platform_activated', lambda: True),
-    ],
-)
-
-
-build_detail_endpoint_tests(
-    CustomDesignProjectPlatformAPIEndpointTest,
-    resource='custom_design',
-    get_resource_id=lambda self: 'pipelines',
-    result_keys_to_compare=[
-        'components',
-        'pages',
-        'project',
-        'uuid',
-    ],
-    assert_after=__assert_after_detail2,
-    patch_function_settings=[
-        (
-            '.'.join([
-                'mage_ai',
-                'api',
-                'resources',
-                'CustomDesignResource',
-                'Project',
-                'is_feature_enabled_in_root_or_active_project',
-            ]),
-            lambda feature_name, user=None, **kwargs: FeatureUUID.CUSTOM_DESIGN == feature_name,
-        ),
-        ('mage_ai.presenters.design.models.project_platform_activated', lambda: True),
-        ('mage_ai.settings.platform.utils.project_platform_activated', lambda: True),
-        ('mage_ai.settings.repo.project_platform_activated', lambda: True),
-    ],
-)

mage_ai/tests/api/endpoints/test_custom_designs_project_platform.py

@@ -0,0 +1,132 @@
+import json
+import os
+from unittest.mock import patch
+
+from mage_ai.data_preparation.models.project.constants import FeatureUUID
+from mage_ai.settings.utils import base_repo_path
+from mage_ai.tests.api.endpoints.mixins import (
+    BaseAPIEndpointTest,
+    build_detail_endpoint_tests,
+    build_list_endpoint_tests,
+)
+from mage_ai.tests.shared.mixins import CustomDesignMixin, ProjectPlatformMixin
+
+CURRENT_FILE_PATH = os.path.dirname(os.path.realpath(__file__))
+
+
+@patch('mage_ai.presenters.design.models.project_platform_activated', lambda: True)
+@patch('mage_ai.settings.platform.utils.project_platform_activated', lambda: True)
+@patch('mage_ai.settings.repo.project_platform_activated', lambda: True)
+class CustomDesignProjectPlatformAPIEndpointTest(
+    CustomDesignMixin,
+    ProjectPlatformMixin,
+    BaseAPIEndpointTest,
+):
+    pass
+
+
+def __build_query(test_case):
+    return dict(
+        operation=test_case.faker.name(),
+        page_path=test_case.faker.name(),
+        page_pathname=test_case.faker.name(),
+        page_query=test_case.faker.name(),
+        page_type=test_case.faker.name(),
+        page_uuid=test_case.faker.name(),
+        resource=test_case.faker.name(),
+        resource_id=test_case.faker.name(),
+        resource_parent=test_case.faker.name(),
+        resource_parent_id=test_case.faker.name(),
+    )
+
+
+def __assert_after_list2(self, result, **kwargs):
+    with open(os.path.join(CURRENT_FILE_PATH, 'mocks', 'mock_design_project_platform.json')) as f:
+        arr = json.loads(f.read())
+        for item in arr:
+            item['project']['full_path'] = os.path.join(
+                base_repo_path(),
+                item['project']['full_path'],
+            )
+            item['project']['root_project_full_path'] = base_repo_path()
+
+        self.assertEqual(
+            result.sort(key=lambda x: x['uuid']),
+            arr.sort(key=lambda x: x['uuid']),
+        )
+
+
+def __assert_after_detail2(self, result, **kwargs):
+    with open(os.path.join(CURRENT_FILE_PATH, 'mocks', 'mock_design_project_platform.json')) as f:
+        item = json.loads(f.read())[0]
+        item['project']['full_path'] = os.path.join(
+            base_repo_path(),
+            item['project']['full_path'],
+        )
+        item['project']['root_project_full_path'] = base_repo_path()
+
+        self.assertEqual(
+            result,
+            item,
+        )
+
+
+build_list_endpoint_tests(
+    CustomDesignProjectPlatformAPIEndpointTest,
+    list_count=2,
+    resource='custom_design',
+    result_keys_to_compare=[
+        'components',
+        'pages',
+        'project',
+        'uuid',
+    ],
+    assert_after=__assert_after_list2,
+    build_query=__build_query,
+    patch_function_settings=[
+        (
+            '.'.join([
+                'mage_ai',
+                'api',
+                'resources',
+                'CustomDesignResource',
+                'Project',
+                'is_feature_enabled_in_root_or_active_project',
+            ]),
+            lambda feature_name, user=None, **kwargs: FeatureUUID.CUSTOM_DESIGN == feature_name,
+        ),
+        ('mage_ai.presenters.design.models.project_platform_activated', lambda: True),
+        ('mage_ai.settings.platform.utils.project_platform_activated', lambda: True),
+        ('mage_ai.settings.repo.project_platform_activated', lambda: True),
+    ],
+)
+
+
+build_detail_endpoint_tests(
+    CustomDesignProjectPlatformAPIEndpointTest,
+    resource='custom_design',
+    get_resource_id=lambda self: 'pipelines',
+    result_keys_to_compare=[
+        'components',
+        'pages',
+        'project',
+        'uuid',
+    ],
+    assert_after=__assert_after_detail2,
+    patch_function_settings=[
+        (
+            '.'.join([
+                'mage_ai',
+                'api',
+                'resources',
+                'CustomDesignResource',
+                'Project',
+                'is_feature_enabled_in_root_or_active_project',
+            ]),
+            lambda feature_name, user=None, **kwargs: FeatureUUID.CUSTOM_DESIGN == feature_name,
+        ),
+        ('mage_ai.presenters.design.models.project_platform_activated', lambda: True),
+        ('mage_ai.settings.platform.utils.project_platform_activated', lambda: True),
+        ('mage_ai.settings.repo.project_platform_activated', lambda: True),
+    ],
+)