mage-ai 0.9.68__py3-none-any.whl → 0.9.70__py3-none-any.whl

mage_ai/authentication/providers/okta.py

@@ -1,104 +1,43 @@
-import urllib.parse
-import uuid
-from typing import Awaitable, Dict
-
-import aiohttp
-from aiohttp import BasicAuth
+from urllib.parse import unquote, urlparse
 
 from mage_ai.authentication.oauth.constants import ProviderName
-from mage_ai.authentication.providers.oauth import OauthProvider
-from mage_ai.authentication.providers.sso import SsoProvider
-from mage_ai.authentication.providers.utils import get_base_url
+from mage_ai.authentication.providers.oidc import OidcProvider
 from mage_ai.settings import get_settings_value
 from mage_ai.settings.keys import OKTA_CLIENT_ID, OKTA_CLIENT_SECRET, OKTA_DOMAIN_URL
 
 
-class OktaProvider(SsoProvider, OauthProvider):
+class OktaProvider(
+    OidcProvider
+):  # Okta configuration uses OIDC so we can just subclass the OidcProvider
     provider = ProviderName.OKTA
 
     def __init__(self):
         self.hostname = get_settings_value(OKTA_DOMAIN_URL)
         self.client_id = get_settings_value(OKTA_CLIENT_ID)
         self.client_secret = get_settings_value(OKTA_CLIENT_SECRET)
+        self.parsed_url = urlparse(unquote(self.hostname))
+        if not self.parsed_url.scheme:
+            self.parsed_url = urlparse(unquote(f'https://{self.hostname}'))
         self.__validate()
 
-        if not self.hostname.startswith('https'):
-            self.hostname = f'https://{self.hostname}'
+        self.discovery_url = (
+            f'https://{self.parsed_url.netloc}/.well-known/openid-configuration'
+        )
+        self.discover()
 
     def __validate(self):
-        if not self.hostname:
-            raise Exception(
+        if not self.parsed_url.netloc:
+            raise ValueError(
                 'Okta hostname is empty. '
-                'Make sure the OKTA_DOMAIN_URL environment variable is set.')
+                'Make sure the OKTA_DOMAIN_URL environment variable is set.'
+            )
         if not self.client_id:
-            raise Exception(
+            raise ValueError(
                 'Okta client id is empty. '
-                'Make sure the OKTA_CLIENT_ID environment variable is set.')
+                'Make sure the OKTA_CLIENT_ID environment variable is set.'
+            )
         if not self.client_secret:
-            raise Exception(
+            raise ValueError(
                 '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) -> Dict:
-        base_url = get_base_url(redirect_uri)
-        redirect_uri_query = dict(
-            provider=self.provider,
-            redirect_uri=redirect_uri,
-        )
-        query = dict(
-            client_id=self.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,
-        )
-
-    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:
-            async with session.post(
-                f'{self.hostname}/oauth2/default/v1/token',
-                headers={
-                    'Content-Type': 'application/x-www-form-urlencoded',
-                },
-                data=dict(
-                    grant_type='authorization_code',
-                    code=code,
-                    redirect_uri=f'{base_url}/oauth',
-                ),
-                auth=BasicAuth(self.client_id, self.client_secret),
-                timeout=20,
-            ) as response:
-                data = await response.json()
-
-        return data
-
-    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:
-            async with session.get(
-                f'{self.hostname}/oauth2/default/v1/userinfo',
-                headers={
-                    'Authorization': f'Bearer {access_token}'
-                },
-                timeout=10,
-            ) as response:
-                userinfo_resp = await response.json()
-
-        return dict(
-            email=userinfo_resp.get('email'),
-            username=userinfo_resp.get('sub'),
-        )
+                'Make sure the OKTA_CLIENT_SECRET environment variable is set.'
+            )

mage_ai/cache/block_action_object/__init__.py

@@ -294,7 +294,7 @@ class BlockActionObjectCache(BaseCache):
                     )
 
                     content = None
-                    with open(block_file_absolute_path, 'r') as f:
+                    with open(block_file_absolute_path, 'r', encoding='utf-8') as f:
                         content = f.read()
 
                     mapping[OBJECT_TYPE_BLOCK_FILE][key] = dict(

mage_ai/cluster_manager/kubernetes/workload_manager.py

@@ -477,6 +477,57 @@ class WorkloadManager:
 
         return k8s_service
 
+    def patch_workload(
+        self,
+        name: str,
+        workspace_config: KubernetesWorkspaceConfig,
+        update_workspace_settings: bool = False,
+    ) -> None:
+        """
+        Update workload for k8s. Currently the only fields that can be updated are
+        container_config and workspace settings.
+
+        Args:
+            name (str): name of the workload/workspace
+            workspace_config (KubernetesWorkspaceConfig): new workspace config
+            update_workspace_settings (bool): whether to update workspace settings with
+                the current pod's environment variables.
+        """
+        container_config_yaml = workspace_config.container_config
+        container_config = dict()
+        if isinstance(container_config_yaml, str):
+            container_config = yaml.full_load(container_config_yaml)
+
+        if update_workspace_settings:
+            env_vars = self.__populate_env_vars(
+                name,
+                container_config=container_config,
+            )
+            container_config['env'] = env_vars
+
+        mage_container_config = {
+            'name': f'{name}-container',
+            **container_config,
+        }
+
+        stateful_set_template_spec = {
+            'containers': [mage_container_config],
+        }
+
+        updated_stateful_set = {
+            'spec': {
+                'template': {
+                    'spec': stateful_set_template_spec,
+                },
+            },
+        }
+
+        self.apps_client.patch_namespaced_stateful_set(
+            name,
+            namespace=self.namespace,
+            body=updated_stateful_set,
+        )
+
     def add_service_to_ingress_paths(
         self,
         ingress_name: str,
@@ -712,7 +763,7 @@ class WorkloadManager:
     def __populate_env_vars(
         self,
         name,
-        project_type: str = 'standalone',
+        project_type: str = None,
         project_uuid: str = None,
         container_config: Dict = None,
         initial_metadata: Dict = None,

mage_ai/cluster_manager/workspace/base.py

@@ -148,6 +148,12 @@ class Workspace(abc.ABC):
         """
         raise NotImplementedError('Initialize method not implemented')
 
+    def update(self, payload: Dict, **kwargs):
+        """
+        Update the workspace configuration.
+        """
+        raise NotImplementedError('Update method not implemented')
+
     @abc.abstractmethod
     @safe_db_query
     def delete(self, **kwargs):

mage_ai/cluster_manager/workspace/kubernetes.py

@@ -1,4 +1,5 @@
 import os
+from typing import Dict
 
 import yaml
 
@@ -7,7 +8,7 @@ from mage_ai.cluster_manager.constants import KUBE_NAMESPACE, ClusterType
 from mage_ai.cluster_manager.kubernetes.workload_manager import WorkloadManager
 from mage_ai.cluster_manager.workspace.base import Workspace
 from mage_ai.data_preparation.repo_manager import ProjectType, get_project_type
-from mage_ai.shared.hash import merge_dict
+from mage_ai.shared.hash import extract, merge_dict
 
 
 class KubernetesWorkspace(Workspace):
@@ -82,6 +83,26 @@ class KubernetesWorkspace(Workspace):
 
         return cls(name)
 
+    def update(self, payload: Dict, **kwargs):
+        update_workspace_settings = payload.pop('update_workspace_settings', False)
+        extracted_payload = extract(payload, [
+            'container_config',
+        ])
+        updated_config = merge_dict(
+            self.config.to_dict(),
+            extracted_payload,
+        )
+        workspace_config = KubernetesWorkspaceConfig.load(
+            config=updated_config
+        )
+        self.workload_manager.patch_workload(
+            self.name,
+            workspace_config,
+            update_workspace_settings=update_workspace_settings,
+        )
+        with open(self.config_path, 'w', encoding='utf-8') as fp:
+            yaml.dump(workspace_config.to_dict(), fp)
+
     def delete(self, **kwargs):
         try:
             self.workload_manager.delete_workload(

mage_ai/command_center/applications/utils.py

@@ -17,7 +17,7 @@ from mage_ai.version_control.models import File
 
 """
 Icons:
-Arcane Library (file browser): ChurnV3
+Text Editor (file browser): ChurnV3
 Portal Terminal (terminal):    RankingV3
 Version Control (file diffs):  ForecastV3
 
@@ -65,7 +65,7 @@ async def build_application_arcane_library() -> Dict:
     return build_generic(model_class=File, item_dict=dict(
         item_type=ItemType.DETAIL,
         object_type=ObjectType.FILE,
-        title='Arcane Library',
+        title='Text editor',
         description='browse and edit files across all projects',
         subtitle='Browser / Editor',
         applications=[

mage_ai/command_center/presenters/text.py

@@ -5,5 +5,5 @@ def application_title(uuid: ApplicationExpansionUUID) -> str:
     if ApplicationExpansionUUID.VersionControlFileDiffs == uuid:
         return 'version control file diffs'
     if ApplicationExpansionUUID.ArcaneLibrary == uuid:
-        return 'the Arcane Library'
+        return 'the Text Editor'
     return f'the {uuid} application'

mage_ai/data_preparation/executors/k8s_block_executor.py

@@ -4,10 +4,13 @@ from jinja2 import Template
 
 from mage_ai.data_preparation.executors.block_executor import BlockExecutor
 from mage_ai.data_preparation.shared.utils import get_template_vars
+from mage_ai.orchestration.db import safe_db_query
+from mage_ai.orchestration.db.models.schedules import BlockRun
 from mage_ai.services.k8s.config import K8sExecutorConfig
 from mage_ai.services.k8s.constants import DEFAULT_NAMESPACE
 from mage_ai.services.k8s.job_manager import JobManager as K8sJobManager
 from mage_ai.shared.hash import merge_dict
+from mage_ai.shared.utils import clean_name
 
 
 class K8sBlockExecutor(BlockExecutor):
@@ -15,10 +18,13 @@ class K8sBlockExecutor(BlockExecutor):
 
     def __init__(self, pipeline, block_uuid: str, execution_partition: str = None, **kwargs):
         super().__init__(pipeline, block_uuid, execution_partition=execution_partition)
-        self.executor_config = self.pipeline.repo_config.k8s_executor_config or dict()
+        self.executor_config_dict = self.pipeline.repo_config.k8s_executor_config or dict()
         if self.block.executor_config is not None:
-            self.executor_config = merge_dict(self.executor_config, self.block.executor_config)
-        self.executor_config = K8sExecutorConfig.load(config=self.executor_config)
+            self.executor_config_dict = merge_dict(
+                self.executor_config_dict,
+                self.block.executor_config,
+            )
+        self.executor_config = K8sExecutorConfig.load(config=self.executor_config_dict)
 
     def _execute(
         self,
@@ -26,10 +32,7 @@ class K8sBlockExecutor(BlockExecutor):
         global_vars: Dict = None,
         **kwargs,
     ) -> None:
-        if not self.executor_config.job_name_prefix:
-            job_name_prefix = 'data-prep'
-        else:
-            job_name_prefix = self.executor_config.job_name_prefix
+        job_name_prefix = self._get_job_name_prefix(block_run_id)
 
         if self.executor_config.namespace:
             namespace = Template(self.executor_config.namespace).render(
@@ -54,3 +57,23 @@ class K8sBlockExecutor(BlockExecutor):
             cmd,
             k8s_config=self.executor_config,
         )
+
+    @safe_db_query
+    def _get_job_name_prefix(
+        self,
+        block_run_id,
+    ):
+        if not self.executor_config.job_name_prefix:
+            job_name_prefix = 'data-prep'
+        else:
+            job_name_prefix = self.executor_config.job_name_prefix
+        if not block_run_id:
+            return job_name_prefix
+
+        if '{trigger_name}' in job_name_prefix:
+            block_run = BlockRun.query.get(block_run_id)
+            trigger = block_run.pipeline_run.pipeline_schedule
+            job_name_prefix = job_name_prefix.format(
+                trigger_name=clean_name(trigger.name).replace('_', '-'))
+
+        return job_name_prefix

mage_ai/data_preparation/executors/k8s_pipeline_executor.py

@@ -6,19 +6,25 @@ from jinja2 import Template
 from mage_ai.data_preparation.executors.pipeline_executor import PipelineExecutor
 from mage_ai.data_preparation.models.pipeline import Pipeline
 from mage_ai.data_preparation.shared.utils import get_template_vars
+from mage_ai.orchestration.db import safe_db_query
+from mage_ai.orchestration.db.models.schedules import PipelineRun
 from mage_ai.services.k8s.config import K8sExecutorConfig
 from mage_ai.services.k8s.constants import DEFAULT_NAMESPACE
 from mage_ai.services.k8s.job_manager import JobManager as K8sJobManager
 from mage_ai.shared.hash import merge_dict
+from mage_ai.shared.utils import clean_name
 
 
 class K8sPipelineExecutor(PipelineExecutor):
     def __init__(self, pipeline: Pipeline, execution_partition: str = None):
         super().__init__(pipeline, execution_partition=execution_partition)
-        self.executor_config = self.pipeline.repo_config.k8s_executor_config or dict()
+        self.executor_config_dict = self.pipeline.repo_config.k8s_executor_config or dict()
         if self.pipeline.executor_config is not None:
-            self.executor_config = merge_dict(self.executor_config, self.pipeline.executor_config)
-        self.executor_config = K8sExecutorConfig.load(config=self.executor_config)
+            self.executor_config_dict = merge_dict(
+                self.executor_config_dict,
+                self.pipeline.executor_config,
+            )
+        self.executor_config = K8sExecutorConfig.load(config=self.executor_config_dict)
 
     def cancel(
         self,
@@ -63,10 +69,7 @@ class K8sPipelineExecutor(PipelineExecutor):
     ) -> K8sJobManager:
         if global_vars is None:
             global_vars = dict()
-        if not self.executor_config.job_name_prefix:
-            job_name_prefix = 'data-prep'
-        else:
-            job_name_prefix = self.executor_config.job_name_prefix
+        job_name_prefix = self._get_job_name_prefix(pipeline_run_id)
 
         if self.executor_config.namespace:
 
@@ -83,3 +86,23 @@ class K8sPipelineExecutor(PipelineExecutor):
             logging_tags=kwargs.get('tags', dict()),
             namespace=namespace,
         )
+
+    @safe_db_query
+    def _get_job_name_prefix(
+        self,
+        pipeline_run_id,
+    ):
+        if not self.executor_config.job_name_prefix:
+            job_name_prefix = 'data-prep'
+        else:
+            job_name_prefix = self.executor_config.job_name_prefix
+        if not pipeline_run_id:
+            return job_name_prefix
+
+        if '{trigger_name}' in job_name_prefix:
+            pipeline_run = PipelineRun.query.get(pipeline_run_id)
+            trigger = pipeline_run.pipeline_schedule
+            job_name_prefix = job_name_prefix.format(
+                trigger_name=clean_name(trigger.name).replace('_', '-'))
+
+        return job_name_prefix

mage_ai/data_preparation/executors/streaming_pipeline_executor.py

@@ -2,9 +2,12 @@ import asyncio
 import copy
 import logging
 import os
+import traceback
 from contextlib import redirect_stderr, redirect_stdout
+from datetime import datetime
 from typing import Callable, Dict, List, Union
 
+import pytz
 import yaml
 from jinja2 import Template
 
@@ -12,9 +15,14 @@ from mage_ai.data_preparation.executors.pipeline_executor import PipelineExecuto
 from mage_ai.data_preparation.logging.logger import DictLogger
 from mage_ai.data_preparation.models.constants import BlockLanguage, BlockType
 from mage_ai.data_preparation.models.pipeline import Pipeline
+from mage_ai.data_preparation.shared.retry import RetryConfig
 from mage_ai.data_preparation.shared.stream import StreamToLogger
 from mage_ai.data_preparation.shared.utils import get_template_vars
+from mage_ai.orchestration.db import safe_db_query
+from mage_ai.orchestration.db.models.schedules import PipelineRun
 from mage_ai.shared.hash import merge_dict
+from mage_ai.shared.retry import retry
+from mage_ai.usage_statistics.logger import UsageStatisticLogger
 
 
 class StreamingPipelineExecutor(PipelineExecutor):
@@ -22,6 +30,7 @@ class StreamingPipelineExecutor(PipelineExecutor):
         super().__init__(pipeline, **kwargs)
         # TODO: Support custom log destination for streaming pipelines
         self.parse_and_validate_blocks()
+        self.retry_metadata = dict(attempts=0)
 
     def parse_and_validate_blocks(self):
         """
@@ -67,13 +76,16 @@ class StreamingPipelineExecutor(PipelineExecutor):
         self,
         build_block_output_stdout: Callable[..., object] = None,
         global_vars: Dict = None,
+        pipeline_run_id: int = None,
+        retry_config: Dict = None,
         **kwargs,
     ) -> None:
         # TODOs:
         # 1. Support multiple sources and sinks
         # 2. Support flink pipeline
 
-        tags = self.build_tags(**kwargs)
+        tags = self.build_tags(pipeline_run_id=pipeline_run_id, **kwargs)
+        self.logging_tags = tags
         if build_block_output_stdout:
             stdout_logger = logging.getLogger('streaming_pipeline_executor')
             self.logger = DictLogger(stdout_logger)
@@ -82,24 +94,53 @@ class StreamingPipelineExecutor(PipelineExecutor):
             self.logger = DictLogger(self.logger_manager.logger, logging_tags=tags)
             stdout = StreamToLogger(self.logger, logging_tags=tags)
         try:
-            with redirect_stdout(stdout):
-                with redirect_stderr(stdout):
-                    self.__execute_in_python(
-                        build_block_output_stdout=build_block_output_stdout,
-                        global_vars=global_vars,
-                    )
+            if retry_config is None:
+                retry_config = self.pipeline.retry_config or dict()
+            infinite_retries = False if retry_config else False
+
+            if type(retry_config) is not RetryConfig:
+                retry_config = RetryConfig.load(config=retry_config)
+
+            @retry(
+                retries=retry_config.retries,
+                delay=retry_config.delay,
+                max_delay=retry_config.max_delay,
+                exponential_backoff=retry_config.exponential_backoff,
+                logger=self.logger,
+                logging_tags=self.logging_tags,
+                retry_metadata=self.retry_metadata,
+            )
+            def __execute_with_retry():
+                with redirect_stdout(stdout):
+                    with redirect_stderr(stdout):
+                        self.__execute_in_python(
+                            build_block_output_stdout=build_block_output_stdout,
+                            global_vars=global_vars,
+                            pipeline_run_id=pipeline_run_id,
+                        )
+            __execute_with_retry()
         except Exception as e:
             if not build_block_output_stdout:
                 self.logger.exception(
                         f'Failed to execute streaming pipeline {self.pipeline.uuid}',
                         **merge_dict(dict(error=e), tags),
                     )
+            if not infinite_retries:
+                # If pipeline retry config is present, fail the pipeline after the retries
+                self.__update_pipeline_run_status(
+                    pipeline_run_id,
+                    PipelineRun.PipelineRunStatus.FAILED,
+                    error=e,
+                )
+
             raise e
 
     def __execute_in_python(
         self,
         build_block_output_stdout: Callable[..., object] = None,
-        global_vars: Dict = None
+        global_vars: Dict = None,
+        pipeline_run_id: int = None,
+
     ):
         from mage_ai.streaming.sinks.sink_factory import SinkFactory
         from mage_ai.streaming.sources.base import SourceConsumeMethod
@@ -227,6 +268,35 @@ class StreamingPipelineExecutor(PipelineExecutor):
             for sink in sinks_by_uuid.values():
                 sink.destroy()
 
+    @safe_db_query
+    def __update_pipeline_run_status(
+        self,
+        pipeline_run_id: int,
+        status: PipelineRun.PipelineRunStatus,
+        error: Exception = None,
+    ):
+        if not pipeline_run_id or not status:
+            return
+        pipeline_run = PipelineRun.query.get(pipeline_run_id)
+        pipeline_run.update(
+            status=status,
+            completed_at=datetime.now(tz=pytz.UTC),
+        )
+        if status == PipelineRun.PipelineRunStatus.FAILED:
+            asyncio.run(UsageStatisticLogger().pipeline_run_ended(pipeline_run))
+            error_msg = None
+            stacktrace = None
+            if error is not None:
+                error_msg = str(error)
+                stacktrace = traceback.format_exc()
+            notification_sender = self.pipeline.get_notification_sender()
+            notification_sender.send_pipeline_run_failure_message(
+                pipeline=self.pipeline,
+                pipeline_run=pipeline_run,
+                error=error_msg,
+                stacktrace=stacktrace,
+            )
+
     def __execute_in_flink(self):
         """
         TODO: Implement this method