mage-ai 0.9.44__py3-none-any.whl → 0.9.46__py3-none-any.whl

mage_ai/authentication/providers/sso.py

@@ -0,0 +1,20 @@
+from abc import ABC, abstractmethod
+from typing import Awaitable, Dict
+
+
+class SsoProvider(ABC):
+    """
+    Base class for single sign on providers. Can be used in conjunction with OauthProvider.
+    """
+    @abstractmethod
+    async def get_user_info(self, **kwargs) -> Awaitable[Dict]:
+        """
+        This method should return a dictionary containing the user's information.
+
+        Returns:
+            Dict: dictionary with the following parameters
+                email (str): user's email
+                username (str): username
+                user_roles (List[Role], optional): list of user roles for the user
+        """
+        pass

mage_ai/authentication/providers/utils.py

@@ -0,0 +1,12 @@
+from urllib.parse import unquote, urlparse
+
+from mage_ai.settings import ROUTES_BASE_PATH
+
+
+def get_base_url(url: str) -> str:
+    parsed_url = urlparse(unquote(url))
+    base_url = parsed_url.scheme + '://' + parsed_url.netloc
+    if ROUTES_BASE_PATH:
+        base_url += f'/{ROUTES_BASE_PATH}'
+
+    return base_url

mage_ai/cluster_manager/aws/emr_cluster_manager.py

@@ -51,6 +51,7 @@ class EmrClusterManager(ClusterManager):
 
     def set_active_cluster(
         self,
+        auto_creation: bool = True,
         auto_selection: bool = False,
         cluster_id=None,
         emr_config: Dict = None,
@@ -59,7 +60,7 @@ class EmrClusterManager(ClusterManager):
             clusters = self.list_clusters()
             if len(clusters) > 0:
                 cluster_id = clusters[0]['id']
-            else:
+            elif auto_creation:
                 self.create_cluster(emr_config=emr_config)
         if cluster_id is None:
             return

mage_ai/cluster_manager/workspace/base.py

@@ -10,6 +10,7 @@ from mage_ai.data_preparation.repo_manager import ProjectType, get_project_type
 from mage_ai.orchestration.constants import Entity
 from mage_ai.orchestration.db.models.oauth import User
 from mage_ai.settings.repo import get_repo_path
+from mage_ai.shared.hash import merge_dict
 
 
 class classproperty(property):
@@ -25,7 +26,7 @@ class Workspace(abc.ABC):
         self._config = None
 
     @classproperty
-    def project_folder(cls) -> str:
+    def project_folder(self) -> str:
         return os.path.join(get_repo_path(), 'projects')
 
     @property
@@ -48,7 +49,7 @@ class Workspace(abc.ABC):
 
     @property
     def project_uuid(self) -> Optional[str]:
-        self.config.get('project_uuid')
+        return self.config.project_uuid
 
     def get_access(self, user: User) -> int:
         return user.get_access(Entity.PROJECT, self.project_uuid)
@@ -114,8 +115,11 @@ class Workspace(abc.ABC):
             data['project_uuid'] = project_uuid
 
         workspace_class = cls.workspace_class_from_type(cluster_type)
+
         try:
-            return workspace_class.initialize(name, config_path, **payload, **data)
+            return workspace_class.initialize(
+                name, config_path, **merge_dict(payload, data)
+            )
         except Exception:
             if config_path and os.path.exists(config_path):
                 os.remove(config_path)

mage_ai/data_integrations/destinations/constants.py

@@ -2,9 +2,11 @@ DESTINATIONS = [
     dict(name='Amazon S3'),
     dict(name='BigQuery'),
     dict(name='Clickhouse'),
+    dict(name='Delta Lake Azure'),
     dict(name='Delta Lake S3'),
     dict(name='Elasticsearch'),
     dict(name='Google Cloud Storage'),
+    dict(name='Kafka'),
     dict(name='MongoDB'),
     dict(
         module_name='MSSQL',

mage_ai/data_preparation/executors/block_executor.py

@@ -557,6 +557,7 @@ class BlockExecutor:
                         )
                     self._execute_callback(
                         'on_failure',
+                        callback_kwargs=dict(__error=error),
                         dynamic_block_index=dynamic_block_index,
                         dynamic_upstream_block_uuids=dynamic_upstream_block_uuids,
                         global_vars=global_vars,
@@ -1043,6 +1044,7 @@ class BlockExecutor:
         global_vars: Dict,
         logging_tags: Dict,
         pipeline_run: PipelineRun,
+        callback_kwargs: Dict = None,
         dynamic_block_index: Union[int, None] = None,
         dynamic_upstream_block_uuids: Union[List[str], None] = None,
     ):
@@ -1068,6 +1070,7 @@ class BlockExecutor:
             try:
                 callback_block.execute_callback(
                     callback,
+                    callback_kwargs=callback_kwargs,
                     dynamic_block_index=dynamic_block_index,
                     dynamic_upstream_block_uuids=dynamic_upstream_block_uuids,
                     execution_partition=self.execution_partition,

mage_ai/data_preparation/executors/pipeline_executor.py

@@ -1,6 +1,9 @@
 import asyncio
+from datetime import datetime
 from typing import Dict, List
 
+import pytz
+
 from mage_ai.data_preparation.executors.block_executor import BlockExecutor
 from mage_ai.data_preparation.logging.logger import DictLogger
 from mage_ai.data_preparation.logging.logger_manager_factory import LoggerManagerFactory
@@ -93,6 +96,12 @@ class PipelineExecutor:
                     block_uuid=block_run.block_uuid,
                     execution_partition=self.execution_partition,
                 )
+                block_run_data = dict(status=BlockRun.BlockRunStatus.RUNNING)
+                if not block_run.started_at or \
+                        (block_run.metrics and not block_run.metrics.get('controller')):
+                    block_run_data['started_at'] = datetime.now(tz=pytz.UTC)
+
+                block_run.update(**block_run_data)
                 BlockExecutor(**executor_kwargs).execute(
                     block_run_id=block_run.id,
                     global_vars=global_vars,

mage_ai/data_preparation/git/__init__.py

@@ -59,7 +59,7 @@ class Git:
             if self.remote_repo_link:
                 url = urlsplit(self.remote_repo_link)
 
-            token = self.__get_access_token()
+            token = self.get_access_token()
 
             if self.git_config and url:
                 user = self.git_config.username
@@ -299,7 +299,7 @@ class Git:
                     url = urlsplit(submodule_url)
                     if self.auth_type == AuthType.HTTPS:
                         user = self.git_config.username
-                        token = self.__get_access_token()
+                        token = self.get_access_token()
                         url = url._replace(netloc=f'{user}:{token}@{url.netloc}')
                         url = urlunsplit(url)
                         # Overwrite the submodule URL with git credentials.
@@ -780,7 +780,7 @@ class Git:
             proc.kill()
             raise TimeoutError
 
-    def __get_access_token(self) -> str:
+    def get_access_token(self) -> str:
         token = None
         if os.getenv(GIT_ACCESS_TOKEN_VAR):
             token = os.getenv(GIT_ACCESS_TOKEN_VAR)

mage_ai/data_preparation/git/api.py

@@ -27,10 +27,7 @@ def get_oauth_client_id(provider: str) -> str:
     return f'{provider}_{get_project_uuid()}'
 
 
-def get_access_token_for_user(
-    user: User,
-    provider: str = None
-) -> Oauth2AccessToken:
+def get_access_token_for_user(user: User, provider: str = None) -> Oauth2AccessToken:
     if not provider:
         provider = OAUTH_PROVIDER_GHE if get_ghe_hostname() else OAUTH_PROVIDER_GITHUB
     access_tokens = access_tokens_for_client(get_oauth_client_id(provider), user=user)
@@ -38,14 +35,19 @@ def get_access_token_for_user(
         return access_tokens[0]
 
 
-def fetch(remote_name: str, remote_url: str, token: str) -> RemoteProgress:
+def fetch(
+    remote_name: str,
+    remote_url: str,
+    token: str,
+    user: User = None,
+) -> RemoteProgress:
     from mage_ai.data_preparation.git import Git
 
     custom_progress = RemoteProgress()
-    username = get_username(token)
+    username = get_username(token, user=user)
 
     url = build_authenticated_remote_url(remote_url, username, token)
-    git_manager = Git.get_manager()
+    git_manager = Git.get_manager(user=user)
 
     remote = git_manager.repo.remotes[remote_name]
     url_original = list(remote.urls)[0]
@@ -65,14 +67,20 @@ def fetch(remote_name: str, remote_url: str, token: str) -> RemoteProgress:
     return custom_progress
 
 
-def pull(remote_name: str, remote_url: str, branch_name: str, token: str) -> RemoteProgress:
+def pull(
+    remote_name: str,
+    remote_url: str,
+    branch_name: str,
+    token: str,
+    user: User = None,
+) -> RemoteProgress:
     from mage_ai.data_preparation.git import Git
 
     custom_progress = RemoteProgress()
-    username = get_username(token)
+    username = get_username(token, user=user)
 
     url = build_authenticated_remote_url(remote_url, username, token)
-    git_manager = Git.get_manager()
+    git_manager = Git.get_manager(user=user)
 
     remote = git_manager.repo.remotes[remote_name]
     url_original = list(remote.urls)[0]
@@ -92,14 +100,20 @@ def pull(remote_name: str, remote_url: str, branch_name: str, token: str) -> Rem
     return custom_progress
 
 
-def push(remote_name: str, remote_url: str, branch_name: str, token: str) -> RemoteProgress:
+def push(
+    remote_name: str,
+    remote_url: str,
+    branch_name: str,
+    token: str,
+    user: User = None,
+) -> RemoteProgress:
     from mage_ai.data_preparation.git import Git
 
     custom_progress = RemoteProgress()
-    username = get_username(token)
+    username = get_username(token, user=user)
 
     url = build_authenticated_remote_url(remote_url, username, token)
-    git_manager = Git.get_manager()
+    git_manager = Git.get_manager(user=user)
 
     remote = git_manager.repo.remotes[remote_name]
     url_original = list(remote.urls)[0]
@@ -119,13 +133,19 @@ def push(remote_name: str, remote_url: str, branch_name: str, token: str) -> Rem
     return custom_progress
 
 
-def reset_hard(remote_name: str, remote_url: str, branch_name: str, token: str) -> None:
+def reset_hard(
+    remote_name: str,
+    remote_url: str,
+    branch_name: str,
+    token: str,
+    user: User = None,
+) -> None:
     from mage_ai.data_preparation.git import Git
 
-    username = get_username(token)
+    username = get_username(token, user=user)
 
     url = build_authenticated_remote_url(remote_url, username, token)
-    git_manager = Git.get_manager()
+    git_manager = Git.get_manager(user=user)
 
     remote = git_manager.repo.remotes[remote_name]
     url_original = list(remote.urls)[0]
@@ -142,13 +162,18 @@ def reset_hard(remote_name: str, remote_url: str, branch_name: str, token: str)
             print(err)
 
 
-def clone(remote_name: str, remote_url: str, token: str) -> None:
+def clone(
+    remote_name: str,
+    remote_url: str,
+    token: str,
+    user: User = None,
+) -> None:
     from mage_ai.data_preparation.git import Git
 
-    username = get_username(token)
+    username = get_username(token, user=user)
 
     url = build_authenticated_remote_url(remote_url, username, token)
-    git_manager = Git.get_manager()
+    git_manager = Git.get_manager(user=user)
 
     remote = git_manager.repo.remotes[remote_name]
     url_original = list(remote.urls)[0]
@@ -172,7 +197,7 @@ def clone(remote_name: str, remote_url: str, token: str) -> None:
             tmp_path,
             git_manager.repo_path,
             dirs_exist_ok=True,
-            ignore=lambda x, y: ['.preferences.yaml']
+            ignore=lambda x, y: ['.preferences.yaml'],
         )
         Git.get_manager().repo.git.clean('-fd', exclude='.preferences.yaml')
     finally:
@@ -200,8 +225,14 @@ def build_authenticated_remote_url(remote_url: str, username: str, token: str) -
     return urlunsplit(url)
 
 
-def get_username(token: str) -> str:
-    return get_user(token)['login']
+def get_username(token: str, user: User = None) -> str:
+    resp = get_user(token)
+    if resp.get('login') is None:
+        from mage_ai.data_preparation.git import Git
+        git_manager = Git.get_manager(user=user, setup_repo=False)
+        return git_manager.git_config.username
+    else:
+        return resp['login']
 
 
 def get_user(token: str) -> Dict:
@@ -212,11 +243,15 @@ def get_user(token: str) -> Dict:
     endpoint = f'{API_ENDPOINT}/user'
     if ghe_hostname and ghe_hostname != DEFAULT_GITHUB_HOSTNAME:
         endpoint = f'{ghe_hostname}/api/v3/user'
-    resp = requests.get(endpoint, headers={
-        'Accept': 'application/vnd.github+json',
-        'Authorization': f'Bearer {token}',
-        'X-GitHub-Api-Version': '2022-11-28',
-    })
+    resp = requests.get(
+        endpoint,
+        headers={
+            'Accept': 'application/vnd.github+json',
+            'Authorization': f'Bearer {token}',
+            'X-GitHub-Api-Version': '2022-11-28',
+        },
+        timeout=10,
+    )
 
     return resp.json()
 
@@ -228,10 +263,12 @@ def check_connection(repo: Repo, remote_url: str) -> None:
 async def validate_authentication_for_remote_url(repo: Repo, remote_url: str) -> None:
     proc = repo.git.ls_remote(remote_url, as_process=True)
 
-    asyncio.run(__poll_process_with_timeout(
-        proc,
-        error_message='Error connecting to remote, make sure your access is valid.',
-    ))
+    asyncio.run(
+        __poll_process_with_timeout(
+            proc,
+            error_message='Error connecting to remote, make sure your access is valid.',
+        )
+    )
 
 
 async def __poll_process_with_timeout(
@@ -253,10 +290,7 @@ async def __poll_process_with_timeout(
 
     if return_code is not None and return_code != 0:
         _, err = proc.communicate()
-        message = (
-            err.decode('UTF-8') if err
-            else error_message
-        )
+        message = err.decode('UTF-8') if err else error_message
         raise ChildProcessError(message)
 
     if return_code is None:

mage_ai/data_preparation/models/block/__init__.py

@@ -1,6 +1,7 @@
 import asyncio
 import functools
 import importlib.util
+import inspect
 import json
 import logging
 import os
@@ -188,6 +189,7 @@ def run_blocks_sync(
     run_sensors: bool = True,
     run_tests: bool = True,
     selected_blocks: Set[str] = None,
+    update_status: bool = True,
 ) -> None:
     tries_by_block_uuid = {}
     tasks = dict()
@@ -242,6 +244,7 @@ def run_blocks_sync(
                 from_notebook=not run_sensors,
                 global_vars=global_vars,
                 run_all_blocks=True,
+                update_status=update_status,
             )
 
             if run_tests:
@@ -277,6 +280,7 @@ class Block(DataIntegrationMixin, SparkBlock):
         language: BlockLanguage = BlockLanguage.PYTHON,
         configuration: Dict = None,
         has_callback: bool = False,
+        repo_config=None,
         timeout: int = None,
     ) -> None:
         if configuration is None:
@@ -338,12 +342,26 @@ class Block(DataIntegrationMixin, SparkBlock):
 
         self.execution_timestamp_start = None
         self.execution_timestamp_end = None
+        self.execution_uuid = None
+        self._compute_service_uuid = None
+        self._repo_config = repo_config
         self._spark_session_current = None
+        self.global_vars = None
 
     @property
     def uuid(self) -> str:
         return self._uuid
 
+    @property
+    def repo_config(self):
+        if self._repo_config:
+            return self._repo_config
+
+        if self.pipeline:
+            self._repo_config = self.pipeline.repo_config
+
+        return self._repo_config
+
     @property
     def uuid_replicated(self) -> str:
         if self.replicated_block:
@@ -873,6 +891,7 @@ class Block(DataIntegrationMixin, SparkBlock):
         global_vars: Dict = None,
         logger: Logger = None,
         logging_tags: Dict = None,
+        execution_uuid: str = None,
         from_notebook: bool = False,
         **kwargs
     ) -> Dict:
@@ -882,6 +901,9 @@ class Block(DataIntegrationMixin, SparkBlock):
         websocket as a way to test the code in the callback. To run a block in a pipeline
         run, use a BlockExecutor.
         """
+        if execution_uuid:
+            self.execution_uuid = execution_uuid
+
         if logging_tags is None:
             logging_tags = dict()
 
@@ -921,17 +943,18 @@ class Block(DataIntegrationMixin, SparkBlock):
                 from_notebook=from_notebook,
                 **kwargs
             )
-        except Exception as e:
+        except Exception as error:
             for callback_block in callback_arr:
                 callback_block.execute_callback(
                     'on_failure',
+                    callback_kwargs=dict(__error=error),
+                    from_notebook=from_notebook,
                     global_vars=global_vars,
                     logger=logger,
                     logging_tags=logging_tags,
                     parent_block=self,
-                    from_notebook=from_notebook,
                 )
-            raise e
+            raise
 
         for callback_block in callback_arr:
             callback_block.execute_callback(
@@ -1281,44 +1304,6 @@ class Block(DataIntegrationMixin, SparkBlock):
                 **kwargs,
             )
 
-            # If block has downstream dbt block, then materialize output
-            if (
-                BlockType.DBT != self.type and
-                self.language in [BlockLanguage.SQL, BlockLanguage.PYTHON, BlockLanguage.R] and
-                any(BlockType.DBT == block.type for block in self.downstream_blocks) and
-                len(outputs) > 0
-            ):
-                from mage_ai.data_preparation.models.block.dbt import DBTBlock
-
-                # normalize output
-                output = outputs[0]
-                if isinstance(output, pd.DataFrame):
-                    df = output
-                elif isinstance(output, dict):
-                    df = pd.DataFrame([output])
-                elif isinstance(output, list):
-                    df = pd.DataFrame(output)
-                else:
-                    df = pd.DataFrame()
-
-                if df.empty:
-                    if logger:
-                        logger.info('No data for dbt to materialize.')
-                else:
-                    DBTBlock.materialize_df(
-                        df=df,
-                        pipeline_uuid=self.pipeline.uuid,
-                        block_uuid=self.uuid,
-                        targets=list(set(
-                            (block.project_path, block.target(variables=global_vars))
-                            for _uuid, block in self.pipeline.blocks_by_uuid.items()
-                            if isinstance(block, DBTBlock)
-                        )),
-                        logger=logger,
-                        global_vars=global_vars_copy,
-                        runtime_arguments=runtime_arguments,
-                    )
-
         return dict(output=outputs)
 
     @contextmanager
@@ -1442,8 +1427,7 @@ class Block(DataIntegrationMixin, SparkBlock):
             if logger and 'logger' not in global_vars:
                 global_vars['logger'] = logger
 
-            track_spark = from_notebook and self.is_using_spark() and \
-                self.compute_management_enabled()
+            track_spark = from_notebook and self.should_track_spark()
 
             if track_spark:
                 self.clear_spark_jobs_cache()
@@ -1526,7 +1510,8 @@ class Block(DataIntegrationMixin, SparkBlock):
             self.module = module
 
             return block_function_updated
-        except Exception:
+        except Exception as err:
+            print(f'[WARNING] Block.initialize_decorator_modules: {err}')
             print('Falling back to default block execution...')
 
         return block_function
@@ -2511,6 +2496,31 @@ df = get_variable('{self.pipeline.uuid}', '{block_uuid}', 'df')
         return variable_mapping
 
     def __enrich_global_vars(self, global_vars: Dict = None) -> Dict:
+        """
+        Enriches the provided global variables dictionary with additional context, Spark session,
+        environment, configuration, and an empty context dictionary.
+
+        Args:
+            global_vars (Optional[Dict]): A dictionary of global variables to be enriched.
+                If not provided, an empty dictionary is created.
+
+        Returns:
+            Dict: The enriched global variables dictionary.
+
+        This method checks if the pipeline type is DATABRICKS or if the environment is a Spark
+        environment. If true, it adds the Spark session to the global variables if not already
+        present.
+
+        If 'env' is not in the global variables, it adds the environment information using the
+        'get_env()' function.
+
+        Adds the block configuration to the global variables.
+
+        If 'context' is not in global_vars, it adds an empty context dictionary.
+
+        The final enriched global variables dictionary is assigned to the object's 'global_vars'
+        attribute and returned.
+        """
         if global_vars is None:
             global_vars = dict()
         if ((self.pipeline is not None and self.pipeline.type == PipelineType.DATABRICKS) or
@@ -2521,16 +2531,19 @@ df = get_variable('{self.pipeline.uuid}', '{block_uuid}', 'df')
                     global_vars['spark'] = spark
         if 'env' not in global_vars:
             global_vars['env'] = get_env()
-        if 'configuration' not in global_vars:
-            global_vars['configuration'] = self.configuration
+        global_vars['configuration'] = self.configuration
         if 'context' not in global_vars:
             global_vars['context'] = dict()
+
+        self.global_vars = global_vars
+
         return global_vars
 
     def get_spark_session(self):
         if self.spark_init and (not self.pipeline or
                                 not self.pipeline.spark_config):
             return self.spark
+
         try:
             if self.pipeline and self.pipeline.spark_config:
                 spark_config = SparkConfig.load(
@@ -2543,6 +2556,9 @@ df = get_variable('{self.pipeline.uuid}', '{block_uuid}', 'df')
         except Exception:
             self.spark = None
 
+        if not self.spark and self.global_vars and self.global_vars.get('spark'):
+            self.spark = self.global_vars.get('spark')
+
         self.spark_init = True
         return self.spark
 
@@ -3095,6 +3111,7 @@ class CallbackBlock(AddonBlock):
     def execute_callback(
         self,
         callback: str,
+        callback_kwargs: Dict = None,
         dynamic_block_index: Union[int, None] = None,
         dynamic_upstream_block_uuids: Union[List[str], None] = None,
         execution_partition: str = None,
@@ -3109,6 +3126,8 @@ class CallbackBlock(AddonBlock):
             logger=logger,
             logging_tags=logging_tags,
         ):
+            if callback_kwargs is None:
+                callback_kwargs = dict()
             global_vars = self._create_global_vars(
                 global_vars,
                 parent_block,
@@ -3160,23 +3179,39 @@ class CallbackBlock(AddonBlock):
             for kwargs_var in kwargs_vars:
                 global_vars_copy.update(kwargs_var)
 
+            callback_kwargs = merge_dict(
+                callback_kwargs,
+                global_vars_copy,
+            )
+
             for callback_function in callback_functions_legacy:
-                callback_function(**global_vars_copy)
+                callback_function(**callback_kwargs)
 
             for callback_function in callback_functions:
-                try:
+                inner_function = callback_function(callback_status)
+                if inner_function is not None:
+                    sig = inspect.signature(inner_function)
                     # As of version 0.8.81, callback functions have access to the parent block’s
-                    # data output.
-                    callback_function(callback_status, *input_vars, **global_vars_copy)
-                except TypeError:
-                    # This try except block will make the above code backwards compatible in case
+                    # data output. If the callback function has any positional arguments, we will
+                    # pass in the input variables as positional arguments.
+                    if not input_vars or any(
+                        [
+                            p.kind not in [p.KEYWORD_ONLY, p.VAR_KEYWORD]
+                            for p in sig.parameters.values()
+                        ]
+                    ):
+                        inner_function(
+                            *input_vars,
+                            **callback_kwargs,
+                        )
+                    # This else block will make the above code backwards compatible in case
                     # a user has already written callback functions with only keyword arguments.
-                    callback_function(
-                        callback_status,
-                        **merge_dict(global_vars_copy, dict(
-                            __input=outputs_from_input_vars,
-                        )),
-                    )
+                    else:
+                        inner_function(
+                            **merge_dict(callback_kwargs, dict(
+                                __input=outputs_from_input_vars,
+                            )),
+                        )
 
     def update_content(self, content, widget=False) -> 'CallbackBlock':
         if not self.file.exists():
@@ -3211,9 +3246,9 @@ class CallbackBlock(AddonBlock):
                 )
 
             def inner(function):
-                def func(callback_status_inner, *args, **kwargs):
+                def func(callback_status_inner):
                     if callback_status_inner == callback_status:
-                        return function(*args, **kwargs)
+                        return function
                 decorated_functions.append(func)
 
             return inner