mage-ai 0.9.12__py3-none-any.whl → 0.9.14__py3-none-any.whl

mage_ai/server/server.py

@@ -22,8 +22,9 @@ from mage_ai.data_preparation.preferences import get_preferences
 from mage_ai.data_preparation.repo_manager import (
     ProjectType,
     get_project_type,
+    get_variables_dir,
+    init_project_uuid,
     init_repo,
-    update_project_uuid,
 )
 from mage_ai.data_preparation.shared.constants import MANAGE_ENV_VAR
 from mage_ai.data_preparation.sync import GitConfig
@@ -52,6 +53,7 @@ from mage_ai.server.constants import DATA_PREP_SERVER_PORT
 from mage_ai.server.docs_server import run_docs_server
 from mage_ai.server.kernel_output_parser import parse_output_message
 from mage_ai.server.kernels import DEFAULT_KERNEL_NAME
+from mage_ai.server.logger import Logger
 from mage_ai.server.scheduler_manager import (
     SCHEDULER_AUTO_RESTART_INTERVAL,
     check_scheduler_status,
@@ -74,7 +76,7 @@ from mage_ai.settings import (
     SHELL_COMMAND,
     USE_UNIQUE_TERMINAL,
 )
-from mage_ai.settings.repo import set_repo_path
+from mage_ai.settings.repo import DEFAULT_MAGE_DATA_DIR, get_repo_name, set_repo_path
 from mage_ai.shared.constants import InstanceType
 from mage_ai.shared.logger import LoggingLevel
 from mage_ai.shared.utils import is_port_in_use
@@ -85,9 +87,11 @@ BASE_PATH_EXPORTS_FOLDER = 'frontend_dist_base_path'
 BASE_PATH_TEMPLATE_EXPORTS_FOLDER = 'frontend_dist_base_path_template'
 BASE_PATH_PLACEHOLDER = 'CLOUD_NOTEBOOK_BASE_PATH_PLACEHOLDER_'
 
+logger = Logger().new_server_logger(__name__)
+
 
 class MainPageHandler(tornado.web.RequestHandler):
-    def get(self, *args):
+    def get(self, *args, **kwargs):
         self.render('index.html')
 
 
@@ -113,16 +117,26 @@ class ApiSchedulerHandler(BaseHandler):
         self.write(dict(scheduler=dict(status=scheduler_manager.get_status())))
 
 
-def replace_base_path(base_path: str) -> None:
+def replace_base_path(base_path: str) -> str:
     """
-    This function will create and go through the BASE_PATH_EXPORTS_FOLDER and replace all the
+    This function will create the BASE_PATH_EXPORTS_FOLDER and replace all the
     occurrences of CLOUD_NOTEBOOK_BASE_PATH_PLACEHOLDER_ with the base_path parameter.
 
     Args:
         base_path (str): The base path to replace the placeholder with.
+
+    Returns:
+        str: The path of the frontend static export folder with the replaced base paths.
     """
     src = os.path.join(os.path.dirname(__file__), BASE_PATH_TEMPLATE_EXPORTS_FOLDER)
-    dst = os.path.join(os.path.dirname(__file__), BASE_PATH_EXPORTS_FOLDER)
+
+    directory = get_variables_dir()
+    # get_variables_dir can return an s3 path. In that case, use the DEFAULT_MAGE_DATA_DIR
+    # directory.
+    if directory.startswith('s3'):
+        directory = os.path.join(os.path.expanduser(DEFAULT_MAGE_DATA_DIR), get_repo_name())
+        os.makedirs(directory, exist_ok=True)
+    dst = os.path.join(directory, BASE_PATH_EXPORTS_FOLDER)
     if os.path.exists(dst):
         shutil.rmtree(dst)
     shutil.copytree(src, dst)
@@ -138,9 +152,10 @@ def replace_base_path(base_path: str) -> None:
                 # replace favicon
                 with open(filepath, 'w', encoding='utf-8') as f:
                     f.write(s)
+    return dst
 
 
-def make_app(update_routes: bool = False):
+def make_app(template_dir: str = None, update_routes: bool = False):
     shell_command = SHELL_COMMAND
     if shell_command is None:
         shell_command = 'bash'
@@ -151,7 +166,8 @@ def make_app(update_routes: bool = False):
         term_klass = MageUniqueTermManager
     term_manager = term_klass(shell_command=[shell_command])
 
-    template_dir = BASE_PATH_EXPORTS_FOLDER if update_routes else EXPORTS_FOLDER
+    if template_dir is None:
+        template_dir = os.path.join(os.path.dirname(__file__), EXPORTS_FOLDER)
     routes = [
         (r'/', MainPageHandler),
         (r'/files', MainPageHandler),
@@ -171,22 +187,22 @@ def make_app(update_routes: bool = False):
         (
             r'/_next/static/(.*)',
             tornado.web.StaticFileHandler,
-            {'path': os.path.join(os.path.dirname(__file__), f'{template_dir}/_next/static')},
+            {'path': os.path.join(template_dir, '_next/static')},
         ),
         (
             r'/fonts/(.*)',
             tornado.web.StaticFileHandler,
-            {'path': os.path.join(os.path.dirname(__file__), f'{template_dir}/fonts')},
+            {'path': os.path.join(template_dir, 'fonts')},
         ),
         (
             r'/images/(.*)',
             tornado.web.StaticFileHandler,
-            {'path': os.path.join(os.path.dirname(__file__), f'{template_dir}/images')},
+            {'path': os.path.join(template_dir, 'images')},
         ),
         (
             r'/(favicon.ico)',
             tornado.web.StaticFileHandler,
-            {'path': os.path.join(os.path.dirname(__file__), template_dir)},
+            {'path': template_dir},
         ),
         (r'/websocket/', WebSocketServer),
         (r'/websocket/terminal', TerminalWebsocketServer, {'term_manager': term_manager}),
@@ -231,7 +247,9 @@ def make_app(update_routes: bool = False):
         (r'/api/(?P<resource>\w+)/(?P<pk>.+)', ApiResourceDetailHandler),
         (r'/files', MainPageHandler),
         (r'/global-data-products', MainPageHandler),
+        (r'/global-data-products/(?P<uuid>\w+)', MainPageHandler),
         (r'/templates', MainPageHandler),
+        (r'/templates/(?P<uuid>\w+)', MainPageHandler),
         (r'/version-control', MainPageHandler),
     ]
 
@@ -249,7 +267,7 @@ def make_app(update_routes: bool = False):
     return tornado.web.Application(
         updated_routes,
         autoreload=True,
-        template_path=os.path.join(os.path.dirname(__file__), template_dir),
+        template_path=template_dir,
     )
 
 
@@ -262,15 +280,22 @@ async def main(
 
     # Update base path if environment variable is set
     update_routes = False
-
+    template_dir = None
     if BASE_PATH:
         try:
-            replace_base_path(BASE_PATH)
+            template_dir = replace_base_path(BASE_PATH)
             update_routes = True
         except Exception:
-            print('Failed to replace base path, using default routes.')
+            logger.warning(
+                'Server failed to replace base path with error:\n%s',
+                traceback.format_exc(),
+            )
+            logger.warning('Continuing with default routes...')
 
-    app = make_app(update_routes=update_routes)
+    app = make_app(
+        update_routes=update_routes,
+        template_dir=template_dir,
+    )
 
     port = int(port)
     max_port = port + 100
@@ -290,7 +315,7 @@ async def main(
     if update_routes:
         url = f'{url}/{BASE_PATH}'
     webbrowser.open_new_tab(url)
-    print(f'Mage is running at {url} and serving project {project}')
+    logger.info(f'Mage is running at {url} and serving project {project}')
 
     db_connection.start_session(force=True)
 
@@ -302,18 +327,18 @@ async def main(
             try:
                 sync = GitSync(sync_config)
                 sync.sync_data()
-                print(
+                logger.info(
                     f'Successfully synced data from git repo: {sync_config.remote_repo_link}'
                     f', branch: {sync_config.branch}'
                 )
             except Exception as err:
-                print(
+                logger.info(
                     f'Failed to sync data from git repo: {sync_config.remote_repo_link}'
                     f', branch: {sync_config.branch} with error: {str(err)}'
                 )
 
     if REQUIRE_USER_AUTHENTICATION:
-        print('User authentication is enabled.')
+        logger.info('User authentication is enabled.')
         # We need to sleep for a few seconds after creating all the tables or else there
         # may be an error trying to create users.
         sleep(5)
@@ -327,7 +352,7 @@ async def main(
         default_owner_role = Role.get_role(Role.DefaultRole.OWNER)
         owner_users = default_owner_role.users if default_owner_role else []
         if not legacy_owner_user and len(owner_users) == 0:
-            print('User with owner permission doesn’t exist, creating owner user.')
+            logger.info('User with owner permission doesn’t exist, creating owner user.')
             if AUTHENTICATION_MODE.lower() == 'ldap':
                 user = User.create(
                     roles_new=[default_owner_role],
@@ -352,7 +377,8 @@ async def main(
             Oauth2Application.client_id == OAUTH2_APPLICATION_CLIENT_ID,
         ).first()
         if not oauth_client:
-            print('OAuth2 application doesn’t exist for frontend, creating OAuth2 application.')
+            logger.info(
+                'OAuth2 application doesn’t exist for frontend, creating OAuth2 application.')
             oauth_client = Oauth2Application.create(
                 client_id=OAUTH2_APPLICATION_CLIENT_ID,
                 client_type=Oauth2Application.ClientType.PUBLIC,
@@ -360,13 +386,13 @@ async def main(
                 user_id=owner_user.id,
             )
 
-    print('Initializing block cache.')
+    logger.info('Initializing block cache.')
     await BlockCache.initialize_cache(replace=True)
 
-    print('Initializing tag cache.')
+    logger.info('Initializing tag cache.')
     await TagCache.initialize_cache(replace=True)
 
-    print('Initializing block action object cache.')
+    logger.info('Initializing block action object cache.')
     await BlockActionObjectCache.initialize_cache(replace=True)
 
     # Check scheduler status periodically
@@ -414,7 +440,7 @@ def start_server(
             project_uuid=project_uuid,
         )
     set_repo_path(project)
-    update_project_uuid()
+    init_project_uuid()
 
     asyncio.run(UsageStatisticLogger().project_impression())
 

mage_ai/services/spark/spark.py

@@ -1,8 +1,63 @@
 from mage_ai.services.spark.config import SparkConfig
+from typing import List
 import os
 
 
+def get_file_names(jars: List) -> List:
+    """
+    Extracts file names from a list of jar files.
+
+    Args:
+        jars (List): A list of jar files.
+
+    Returns:
+        List: A list of file names.
+    """
+    if jars is None:
+        return None
+
+    return [os.path.basename(jar) for jar in jars]
+
+
+def contains_same_jars(jars_1: List, jars_2: List) -> bool:
+    """
+    Checks if two lists of jar files contain the same jars.
+    The order of the jars in the list does not matter.
+
+    Args:
+        jars_1 (List): A list of jar files.
+        jars_2 (List): A list of jar files.
+
+    Returns:
+        bool: True if the lists contain the same jars, False otherwise.
+    """
+    if jars_1 is None and jars_2 is None:
+        return True
+
+    if jars_1 is None or jars_2 is None:
+        return False
+
+    if len(jars_1) != len(jars_2):
+        return False
+
+    file_names_1 = get_file_names(jars_1)
+    file_names_2 = get_file_names(jars_2)
+
+    return set(file_names_1) == set(file_names_2)
+
+
 def has_same_spark_config(spark_session, spark_config: SparkConfig) -> bool:
+    """
+    Checks if the spark session has the same configuration as the spark config.
+
+    Args:
+        spark_session (SparkSession): The spark session.
+        spark_config (SparkConfig): The spark config.
+
+    Returns:
+        bool: True if the spark session has the same configuration as the
+        spark config, False otherwise.
+    """
     if spark_session is None:
         return False
 
@@ -24,8 +79,8 @@ def has_same_spark_config(spark_session, spark_config: SparkConfig) -> bool:
                 if spark_session.conf.get(f'spark.executorEnv.{key}') != value:
                     print(f'Different spark.executorEnv.{key} found!')
                     return False
-        if (spark_config.spark_jars and spark_config.spark_jars
-                != spark_session.conf.get('spark.jars')):
+        if (spark_config.spark_jars and not contains_same_jars(
+                spark_config.spark_jars, spark_session.conf.get('spark.jars'))):
             print('Different spark.jars found!')
             return False
         if spark_config.others:
@@ -37,6 +92,24 @@ def has_same_spark_config(spark_session, spark_config: SparkConfig) -> bool:
 
 
 def get_spark_session(spark_config: SparkConfig):
+    """
+    Gets a Spark session.
+    If the given spark_config is None, then create a Spark session with the
+    default configuration.
+    If the given spark_config is not None, then check if the active Spark
+    session has the same configuration as the given spark_config.
+    If the active Spark session has the same configuration as the given
+    spark_config, then reuse the active Spark session.
+    If the active Spark session does not have the same configuration as the
+    given spark_config, then create a new Spark session with the given
+    spark_config.
+
+    Args:
+        spark_config (SparkConfig): The Spark configuration.
+
+    Returns:
+        SparkSession: The Spark session.
+    """
     from pyspark.conf import SparkConf
     from pyspark.sql import SparkSession
 

mage_ai/streaming/constants.py

@@ -1,6 +1,7 @@
 from enum import Enum
 
 DEFAULT_BATCH_SIZE = 100
+DEFAULT_TIMEOUT_MS = 500
 
 
 class SourceType(str, Enum):

mage_ai/streaming/sources/kafka.py

@@ -1,14 +1,16 @@
+import importlib
+import json
+import time
 from dataclasses import dataclass
+from enum import Enum
+from typing import Callable, Dict
+
 from kafka import KafkaConsumer
+
 from mage_ai.shared.config import BaseConfig
-from mage_ai.streaming.constants import DEFAULT_BATCH_SIZE
+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 SerializationMethod, SerDeConfig
-from enum import Enum
-from typing import Callable, Dict
-import importlib
-import json
-import time
+from mage_ai.streaming.sources.shared import SerDeConfig, SerializationMethod
 
 
 class SecurityProtocol(str, Enum):
@@ -39,6 +41,7 @@ class KafkaConfig(BaseConfig):
     topic: str
     api_version: str = '0.10.2'
     batch_size: int = DEFAULT_BATCH_SIZE
+    timeout_ms: int = DEFAULT_TIMEOUT_MS
     security_protocol: SecurityProtocol = None
     ssl_config: SSLConfig = None
     sasl_config: SASLConfig = None
@@ -138,16 +141,21 @@ class KafkaSource(BaseSource):
             batch_size = self.config.batch_size
         else:
             batch_size = DEFAULT_BATCH_SIZE
+        if self.config.timeout_ms > 0:
+            timeout_ms = self.config.timeout_ms
+        else:
+            timeout_ms = DEFAULT_TIMEOUT_MS
+
         while True:
             # Response format is {TopicPartiton('topic1', 1): [msg1, msg2]}
             msg_pack = self.consumer.poll(
                 max_records=batch_size,
-                timeout_ms=500,
+                timeout_ms=timeout_ms,
             )
 
             message_values = []
             msg_printed = False
-            for tp, messages in msg_pack.items():
+            for _tp, messages in msg_pack.items():
                 for message in messages:
                     if not msg_printed:
                         self.__print_message(message)

mage_ai/tests/data_preparation/models/custom_templates/test_utils.py

@@ -0,0 +1,36 @@
+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.shared.array import find
+from mage_ai.tests.base_test import DBTestCase
+
+
+class CustomTemplatesUtilsTest(DBTestCase):
+    def test_group_and_hydrate_files(self):
+        file_dicts = [
+            None,
+            dict(parent_names=None),
+            dict(parent_names=[]),
+            dict(parent_names=[None]),
+            dict(parent_names=['mage']),
+            dict(parent_names=['mage', 'fire']),
+        ]
+
+        def load(template_uuid: str):
+            return CustomBlockTemplate(template_uuid=template_uuid)
+
+        with patch.object(
+            CustomBlockTemplate,
+            'load',
+            load,
+        ):
+            arr = group_and_hydrate_files(file_dicts, CustomBlockTemplate)
+
+            self.assertTrue(find(lambda x: x.template_uuid == '', arr))
+            self.assertTrue(find(lambda x: x.template_uuid == 'None', arr))
+            self.assertTrue(find(lambda x: x.template_uuid == 'mage', arr))
+            self.assertTrue(find(lambda x: x.template_uuid == 'mage/fire', arr))

mage_ai/tests/data_preparation/models/test_block.py

@@ -542,3 +542,51 @@ SELECT 1 AS id;
 
         block2 = pipeline.get_block('test_transformer')
         self.assertEqual([], block2.upstream_block_uuids)
+
+    def test_execute_block_from_notebook(self):
+        pipeline = Pipeline.create(
+            'test_pipeline_execute_block_function_from_notebook',
+            repo_path=self.repo_path,
+        )
+        block = Block.create('test_data_loader', 'data_loader', self.repo_path, pipeline=pipeline)
+
+        with open(block.file_path, 'w') as file:
+            file.write("""import pandas as pd
+if 'data_loader' not in globals():
+    from mage_ai.data_preparation.decorators import data_loader
+
+@data_loader
+def load_data():
+    data = {'col1': [1, 3], 'col2': [2, 4]}
+    df = pd.DataFrame(data)
+    return df
+""")
+
+        output = block._execute_block(dict(), from_notebook=True)
+        self.assertIsNotNone(block.module)
+        self.assertEqual(len(output), 1)
+
+    def test_execute_block_from_notebook_no_decorator_definition(self):
+        pipeline = Pipeline.create(
+            'test_pipeline_execute_block_function_from_notebook_no_decorator_definition',
+            repo_path=self.repo_path,
+        )
+        block = Block.create('test_data_loader', 'data_loader', self.repo_path, pipeline=pipeline)
+
+        with open(block.file_path, 'w') as file:
+            file.write("""import pandas as pd
+
+@data_loader
+def load_data():
+    data = {'col1': [1, 3], 'col2': [2, 4]}
+    df = pd.DataFrame(data)
+    return df
+
+@test
+def test_output(output, *args) -> None:
+    assert output is not None, 'The output is undefined'
+""")
+
+        output = block._execute_block(dict(), from_notebook=True)
+        self.assertIsNotNone(block.module)
+        self.assertEqual(len(output), 1)

mage_ai/tests/data_preparation/test_repo_manager.py

@@ -1,10 +1,15 @@
 import os
 import shutil
 import uuid
+from unittest.mock import patch
 
 import yaml
 
-from mage_ai.data_preparation.repo_manager import RepoConfig
+from mage_ai.data_preparation.repo_manager import (
+    RepoConfig,
+    get_project_uuid,
+    set_project_uuid_from_metadata,
+)
 from mage_ai.settings.repo import MAGE_DATA_DIR_ENV_VAR
 from mage_ai.tests.base_test import DBTestCase
 
@@ -71,3 +76,16 @@ class RepoManagerTest(DBTestCase):
         )
         test_dir = os.path.expanduser(f'~/{dir_name}')
         shutil.rmtree(test_dir)
+
+    def test_set_project_uuid_from_metadata(self):
+        test_metadata_file = os.path.join(self.repo_path, 'test_repo_manager.yaml')
+        with open(test_metadata_file, 'w', encoding='utf-8') as f:
+            yaml.dump(dict(project_uuid='123456789'), f)
+
+        with patch(
+            'mage_ai.data_preparation.repo_manager.get_metadata_path',
+            return_value=test_metadata_file
+        ):
+            set_project_uuid_from_metadata()
+            self.assertEqual(get_project_uuid(), '123456789')
+        os.remove(test_metadata_file)

mage_ai/tests/server/test_server.py

@@ -52,7 +52,9 @@ class ServerTests(TestCase):
         self.assertIsNone(handler)
 
     @patch('mage_ai.server.server.BASE_PATH', 'random-test-string-1hasdfh')
-    def test_replace_base_path(self):
+    @patch('mage_ai.server.server.get_variables_dir')
+    def test_replace_base_path(self, mock_variables_dir):
+        mock_variables_dir.return_value = os.path.dirname(server_module.__file__)
         with patch(
             'mage_ai.server.server.BASE_PATH_EXPORTS_FOLDER',
             'base_path_test',
@@ -66,7 +68,9 @@ class ServerTests(TestCase):
         shutil.rmtree(test_dir)
 
     @patch('mage_ai.server.server.BASE_PATH', 'test_prefix')
-    def test_replace_base_path_directory_exists(self):
+    @patch('mage_ai.server.server.get_variables_dir')
+    def test_replace_base_path_directory_exists(self, mock_variables_dir):
+        mock_variables_dir.return_value = os.path.dirname(server_module.__file__)
         test_dir = os.path.join(os.path.dirname(server_module.__file__), 'base_path_test')
         os.makedirs(test_dir)
         with patch(
@@ -77,3 +81,17 @@ class ServerTests(TestCase):
 
         self.assertTrue(len(os.listdir(test_dir)) > 0)
         shutil.rmtree(test_dir)
+
+    @patch('mage_ai.server.server.BASE_PATH', 'test_prefix')
+    @patch('mage_ai.server.server.get_variables_dir')
+    def test_replace_base_path_s3_directory(self, mock_variables_dir):
+        mock_variables_dir.return_value = 's3://test-bucket/test-prefix'
+        test_dir = os.path.join(self.repo_path, 'base_path_test')
+        with patch(
+            'mage_ai.server.server.DEFAULT_MAGE_DATA_DIR',
+            test_dir,
+        ):
+            replace_base_path(base_path='test_prefix')
+
+        self.assertTrue(len(os.listdir(os.path.join(test_dir, 'test'))) > 0)
+        shutil.rmtree(test_dir)

mage_ai/tests/services/spark/test_spark.py

@@ -0,0 +1,52 @@
+from mage_ai.services.spark.spark import contains_same_jars, get_file_names
+from mage_ai.tests.base_test import TestCase
+
+
+class SparkTests(TestCase):
+    def test_get_file_names(self):
+        self.assertEqual(
+            get_file_names(['/home/spark-jars/aws-java-sdk-core-1.11.1026.jar']),
+            ['aws-java-sdk-core-1.11.1026.jar']
+        )
+        self.assertEqual(
+            get_file_names(['spark://************:2222/jars/aws-java-sdk-core-1.11.1026.jar']),
+            ['aws-java-sdk-core-1.11.1026.jar']
+        )
+        self.assertEqual(
+            get_file_names(['/home/spark-jars/test1.jar', '/home/spark-jars/test2.jar']),
+            ['test1.jar', 'test2.jar']
+        )
+        self.assertEqual(
+            get_file_names([
+                'spark://************:2222/jars/test1.jar',
+                'spark://************:2222/jars/test2.jar']),
+            ['test1.jar', 'test2.jar']
+        )
+
+    def test_contains_same_jars(self):
+        self.assertTrue(
+            contains_same_jars(
+                ['/home/spark-jars/aws-java-sdk-core-1.11.1026.jar'],
+                ['aws-java-sdk-core-1.11.1026.jar'])
+        )
+        self.assertTrue(
+            contains_same_jars(
+                ['spark://************:2222/jars/aws-java-sdk-core-1.11.1026.jar'],
+                ['aws-java-sdk-core-1.11.1026.jar'])
+        )
+        self.assertTrue(
+            contains_same_jars(
+                ['/home/spark-jars/aws-java-sdk-core-1.11.1026.jar'],
+                ['spark://************:2222/jars/aws-java-sdk-core-1.11.1026.jar'])
+        )
+        self.assertTrue(
+            contains_same_jars(
+                ['/home/spark-jars/test1.jar', '/home/spark-jars/test2.jar'],
+                ['spark://************:2222/jars/test2.jar',
+                 'spark://************:2222/jars/test1.jar'])
+        )
+        self.assertFalse(
+            contains_same_jars(
+                ['/home/spark-jars/test1.jar', '/home/spark-jars/test2.jar'],
+                ['spark://************:2222/jars/test1.jar'])
+        )

{mage_ai-0.9.12.dist-info → mage_ai-0.9.14.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: mage-ai
-Version: 0.9.12
+Version: 0.9.14
 Summary: Mage is a tool for building and deploying data pipelines.
 Home-page: https://github.com/mage-ai/mage-ai
 Author: Mage