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

mage_ai/authentication/providers/oauth.py

@@ -1,5 +1,5 @@
 from abc import ABC, abstractmethod
-from typing import Dict
+from typing import Awaitable, Dict
 
 
 class OauthProvider(ABC):
@@ -38,7 +38,7 @@ class OauthProvider(ABC):
         pass
 
     @abstractmethod
-    async def get_access_token_response(self, code: str, **kwargs) -> Dict:
+    async def get_access_token_response(self, code: str, **kwargs) -> Awaitable[Dict]:
         """
         This method should call the oauth provider with the authorization code and return
         a dictionary containing an access token for the user.

mage_ai/authentication/providers/okta.py

@@ -1,6 +1,6 @@
 import urllib.parse
 import uuid
-from typing import Dict, Optional
+from typing import Awaitable, Dict
 
 import aiohttp
 from aiohttp import BasicAuth
@@ -17,11 +17,11 @@ class OktaProvider(SsoProvider, OauthProvider):
 
     def __init__(self):
         self.hostname = OKTA_DOMAIN_URL
+        self.__validate()
+
         if not self.hostname.startswith('https'):
             self.hostname = f'https://{self.hostname}'
 
-        self.__validate()
-
     def __validate(self):
         if not self.hostname:
             raise Exception(
@@ -36,33 +36,32 @@ class OktaProvider(SsoProvider, OauthProvider):
                 'Okta client secret is empty. '
                 'Make sure the OKTA_CLIENT_SECRET environment variable is set.')
 
-    def get_auth_url_response(self, redirect_uri: str = None, **kwargs) -> Optional[Dict]:
-        if OKTA_CLIENT_ID:
-            base_url = get_base_url(redirect_uri)
-            redirect_uri_query = dict(
-                provider=self.provider,
-                redirect_uri=redirect_uri,
-            )
-            query = dict(
-                client_id=OKTA_CLIENT_ID,
-                redirect_uri=urllib.parse.quote_plus(
-                    f'{base_url}/oauth',
-                ),
-                response_mode='query',
-                response_type='code',
-                scope='openid email profile',
-                state=uuid.uuid4().hex,
-            )
-            query_strings = []
-            for k, v in query.items():
-                query_strings.append(f'{k}={v}')
+    def get_auth_url_response(self, redirect_uri: str = None, **kwargs) -> Dict:
+        base_url = get_base_url(redirect_uri)
+        redirect_uri_query = dict(
+            provider=self.provider,
+            redirect_uri=redirect_uri,
+        )
+        query = dict(
+            client_id=OKTA_CLIENT_ID,
+            redirect_uri=urllib.parse.quote_plus(
+                f'{base_url}/oauth',
+            ),
+            response_mode='query',
+            response_type='code',
+            scope='openid email profile',
+            state=uuid.uuid4().hex,
+        )
+        query_strings = []
+        for k, v in query.items():
+            query_strings.append(f'{k}={v}')
 
-            return dict(
-                url=f"{self.hostname}/oauth2/default/v1/authorize?{'&'.join(query_strings)}",
-                redirect_query_params=redirect_uri_query,
-            )
+        return dict(
+            url=f"{self.hostname}/oauth2/default/v1/authorize?{'&'.join(query_strings)}",
+            redirect_query_params=redirect_uri_query,
+        )
 
-    async def get_access_token_response(self, code: str, **kwargs) -> Dict:
+    async def get_access_token_response(self, code: str, **kwargs) -> Awaitable[Dict]:
         base_url = get_base_url(kwargs.get('redirect_uri'))
         data = dict()
         async with aiohttp.ClientSession() as session:
@@ -83,7 +82,7 @@ class OktaProvider(SsoProvider, OauthProvider):
 
         return data
 
-    async def get_user_info(self, access_token: str = None, **kwargs) -> Dict:
+    async def get_user_info(self, access_token: str = None, **kwargs) -> Awaitable[Dict]:
         if access_token is None:
             raise Exception('Access token is required to fetch user info.')
         async with aiohttp.ClientSession() as session:

mage_ai/authentication/providers/sso.py

@@ -1,5 +1,5 @@
 from abc import ABC, abstractmethod
-from typing import Dict
+from typing import Awaitable, Dict
 
 
 class SsoProvider(ABC):
@@ -7,7 +7,7 @@ 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) -> Dict:
+    async def get_user_info(self, **kwargs) -> Awaitable[Dict]:
         """
         This method should return a dictionary containing the user's information.
 

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

@@ -6,6 +6,7 @@ DESTINATIONS = [
     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/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/models/block/__init__.py

@@ -189,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()
@@ -243,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:
@@ -278,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:
@@ -339,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:
@@ -874,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:
@@ -883,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()
 
@@ -1406,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()
@@ -1490,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
@@ -2475,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
@@ -2485,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(
@@ -2507,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
 

mage_ai/data_preparation/models/block/data_integration/data.py

@@ -75,7 +75,7 @@ def convert_dataframe_to_output(
             if log_message:
                 log_message(
                     f'Writing {len(records_in_batch)} records from stream {stream} to {file_path} '
-                    f'for batch index {index} ({index + 1 } out of {batches} batch(s)).',
+                    f'for batch index {index} ({index + 1} out of {batches} batch(s)).',
                 )
 
             output_file_paths.append(file_path)

mage_ai/data_preparation/models/block/spark/mixins.py

@@ -12,6 +12,7 @@ from mage_ai.data_preparation.models.block.spark.constants import (
 from mage_ai.data_preparation.models.project import Project
 from mage_ai.data_preparation.models.project.constants import FeatureUUID
 from mage_ai.services.spark.api.service import API
+from mage_ai.services.spark.constants import ComputeServiceUUID
 from mage_ai.services.spark.models.applications import Application
 from mage_ai.services.spark.models.jobs import Job
 from mage_ai.services.spark.models.sqls import Sql
@@ -30,6 +31,16 @@ class SparkBlock:
 
         return self._spark_session_current
 
+    @property
+    def compute_service_uuid(self) -> ComputeServiceUUID:
+        if self._compute_service_uuid:
+            return self._compute_service_uuid
+        self._compute_service_uuid = get_compute_service(
+            ignore_active_kernel=True,
+            repo_config=self.repo_config,
+        )
+        return self._compute_service_uuid
+
     def spark_session_application(self) -> Application:
         if not self.spark_session:
             return
@@ -40,7 +51,7 @@ class SparkBlock:
         if value_tup:
             application_id = value_tup[1]
 
-            return Application(
+            return Application.load(
                 id=application_id,
                 spark_ui_url=self.spark_session.sparkContext.uiWebUrl,
             )
@@ -51,7 +62,13 @@ class SparkBlock:
         )
 
     def is_using_spark(self) -> bool:
-        return get_compute_service()
+        return self.compute_service_uuid in [
+            ComputeServiceUUID.AWS_EMR,
+            ComputeServiceUUID.STANDALONE_CLUSTER,
+        ]
+
+    def should_track_spark(self) -> bool:
+        return self.is_using_spark() and self.compute_management_enabled()
 
     def execution_states(self, cache: bool = False) -> Dict:
         jobs_cache = self.__load_cache()
@@ -82,33 +99,40 @@ class SparkBlock:
     def jobs_during_execution(self) -> List[Job]:
         self.__load_spark_job_submission_timestamps()
 
-        if self.execution_timestamp_start:
-            jobs = self.__get_jobs(application=self.execution_start_application)
-
-            def _filter(
-                job: Job,
-                execution_timestamp_start: float = self.execution_timestamp_start,
-                execution_timestamp_end: float = self.execution_timestamp_end,
-            ) -> bool:
-                if not job.submission_time:
-                    return False
-
-                if isinstance(job.submission_time, str):
-                    submission_timestamp = dateutil.parser.parse(job.submission_time).timestamp()
-                elif isinstance(job.submission_time, float) or isinstance(job.submission_time, int):
-                    submission_timestamp = datetime.fromtimestamp(job.submission_time)
-
-                return execution_timestamp_start and \
-                    execution_timestamp_end and \
-                    submission_timestamp >= execution_timestamp_start and \
-                    (
-                        not execution_timestamp_end or
-                        submission_timestamp <= execution_timestamp_end
-                    )
+        if not self.execution_timestamp_start and self.execution_uuid_start:
+            return []
+
+        def _filter(
+            job: Job,
+            block_uuid=self.uuid,
+            compute_service_uuid=self.compute_service_uuid,
+            execution_timestamp_end: float = self.execution_timestamp_end,
+            execution_timestamp_start: float = self.execution_timestamp_start,
+            execution_uuid_start=self.execution_uuid_start,
+        ) -> bool:
+            if ComputeServiceUUID.AWS_EMR == compute_service_uuid:
+                key = self.execution_uuid_start or self.execution_timestamp_start
+                return job.name == f'{block_uuid}:{key}'
+
+            if not job.submission_time:
+                return False
+
+            if isinstance(job.submission_time, str):
+                submission_timestamp = dateutil.parser.parse(job.submission_time).timestamp()
+            elif isinstance(job.submission_time, float) or isinstance(job.submission_time, int):
+                submission_timestamp = datetime.fromtimestamp(job.submission_time)
+
+            return execution_timestamp_start and \
+                execution_timestamp_end and \
+                submission_timestamp >= execution_timestamp_start and \
+                (
+                    not execution_timestamp_end or
+                    submission_timestamp <= execution_timestamp_end
+                )
 
-            return list(filter(_filter, jobs))
+        jobs = self.__get_jobs(application=self.execution_start_application)
 
-        return []
+        return list(filter(_filter, jobs or []))
 
     def stages_during_execution(self, jobs: List[Job]):
         if not jobs:
@@ -122,7 +146,7 @@ class SparkBlock:
 
         stages = self.__get_stages(application=self.execution_start_application)
 
-        return list(filter(_filter, stages))
+        return list(filter(_filter, stages or []))
 
     def sqls_during_execution(self, jobs: List[Job]):
         if not jobs:
@@ -140,7 +164,7 @@ class SparkBlock:
 
         sqls = self.__get_sqls(application=self.execution_start_application)
 
-        return list(filter(_filter, sqls))
+        return list(filter(_filter, sqls or []))
 
     def clear_spark_jobs_cache(self) -> None:
         if os.path.exists(self.spark_jobs_full_path):
@@ -155,13 +179,24 @@ class SparkBlock:
         for application in applications:
             Application.cache_application(application)
 
-    def set_spark_job_execution_start(self) -> None:
+    def set_spark_job_execution_start(self, execution_uuid: str = None) -> None:
         self.execution_timestamp_start = datetime.utcnow().timestamp()
         application = self.spark_session_application()
 
+        if execution_uuid:
+            self.execution_uuid = execution_uuid
+
+        if self.spark_session and self.spark_session.sparkContext:
+            key = f'{self.uuid}:{self.execution_uuid or self.execution_timestamp_start}'
+            # For jobs
+            self.spark_session.sparkContext.setLocalProperty('callSite.short', key)
+            # For stages
+            self.spark_session.sparkContext.setLocalProperty('callSite.long', key)
+
         self.__update_spark_jobs_cache(
             dict(
                 application=application.to_dict() if application else None,
+                execution_uuid=self.execution_uuid,
                 submission_timestamp=self.execution_timestamp_start,
             ),
             'before',
@@ -187,7 +222,7 @@ class SparkBlock:
 
         if application:
             build_options.update(dict(
-                application_id=application.id,
+                application_id=application.calculated_id(),
                 application_spark_ui_url=application.spark_ui_url,
             ))
         else:
@@ -207,7 +242,7 @@ class SparkBlock:
         else:
             applications = api.applications_sync()
             if applications:
-                jobs = api.jobs_sync(applications[0].id)
+                jobs = api.jobs_sync(applications[0].calculated_id())
 
         return sorted(
             jobs,
@@ -231,7 +266,7 @@ class SparkBlock:
             applications = api.applications_sync()
 
             if applications:
-                stages = api.stages_sync(applications[0].id, dict(
+                stages = api.stages_sync(applications[0].calculated_id(), dict(
                     quantiles='0.01,0.25,0.5,0.75,0.99',
                     withSummaries=True,
                 ))
@@ -253,7 +288,7 @@ class SparkBlock:
 
             sqls = []
             if applications:
-                sqls = api.sqls_sync(applications[0].id, dict(
+                sqls = api.sqls_sync(applications[0].calculated_id(), dict(
                     length=9999,
                 ))
 
@@ -304,6 +339,13 @@ class SparkBlock:
             f.write(json.dumps(data))
 
     def __load_spark_job_submission_timestamps(self) -> None:
+        self.execution_end_application = None
+        self.execution_start_application = None
+        self.execution_timestamp_end = None
+        self.execution_timestamp_start = None
+        self.execution_uuid_end = None
+        self.execution_uuid_start = None
+
         jobs_cache = self.__load_cache()
         if jobs_cache:
             before = jobs_cache.get('before')
@@ -311,6 +353,9 @@ class SparkBlock:
                 self.execution_timestamp_start = (before or {}).get(
                     'submission_timestamp',
                 )
+                self.execution_uuid_start = (before or {}).get(
+                    'execution_uuid',
+                )
                 self.execution_start_application = (before or {}).get(
                     'application',
                 )
@@ -324,6 +369,9 @@ class SparkBlock:
                 self.execution_timestamp_end = (after or {}).get(
                     'submission_timestamp',
                 )
+                self.execution_uuid_end = (after or {}).get(
+                    'execution_uuid',
+                )
                 self.execution_end_application = (after or {}).get(
                     'application',
                 )

mage_ai/data_preparation/models/download/__init__.py

@@ -0,0 +1,8 @@
+class Download():
+    def __init__(self, token):
+        self.token = token
+
+    def to_dict(self, **kwargs):
+        data = {}
+        data['token'] = self.token
+        return data

mage_ai/data_preparation/models/global_hooks/constants.py

@@ -0,0 +1,44 @@
+from enum import Enum
+
+from mage_ai.authentication.permissions.constants import EntityName
+
+DISABLED_RESOURCE_TYPES = [
+    EntityName.ALL,
+    EntityName.ALL_EXCEPT_RESERVED,
+]
+
+RESTRICTED_RESOURCE_TYPES = [
+    EntityName.File,
+    EntityName.FileContent,
+    EntityName.Folder,
+    EntityName.GlobalHook,
+    EntityName.Oauth,
+    EntityName.OauthAccessToken,
+    EntityName.OauthApplication,
+    EntityName.Permission,
+    EntityName.Role,
+    EntityName.RolePermission,
+    EntityName.Secret,
+    EntityName.Session,
+    EntityName.User,
+    EntityName.UserRole,
+    EntityName.Workspace,
+]
+
+RESOURCE_TYPES = [en for en in EntityName if en not in DISABLED_RESOURCE_TYPES]
+
+
+class HookOutputKey(str, Enum):
+    ERROR = 'error'
+    META = 'meta'
+    METADATA = 'metadata'
+    PAYLOAD = 'payload'
+    QUERY = 'query'
+    RESOURCE = 'resource'
+    RESOURCES = 'resources'
+
+
+VALID_KEYS_FOR_INPUT_OUTPUT_DATA_RESTRICTED = [key.value for key in HookOutputKey]
+VALID_KEYS_FOR_INPUT_OUTPUT_DATA_UNRESTRICTED = ['hook']
+VALID_KEYS_FOR_INPUT_OUTPUT_DATA_ALL = \
+    VALID_KEYS_FOR_INPUT_OUTPUT_DATA_RESTRICTED + VALID_KEYS_FOR_INPUT_OUTPUT_DATA_UNRESTRICTED