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

mage_ai/api/resources/SeedResource.py

@@ -6,6 +6,7 @@ class SeedResource(GenericResource):
     @classmethod
     async def create(self, payload, user, **kwargs):
         if payload.get('roles') and payload.get('permissions'):
-            await bootstrap_permissions()
+            policy_names = payload.get('policy_names')
+            await bootstrap_permissions(policy_names=policy_names)
 
         return self({}, user, **kwargs)

mage_ai/api/resources/SessionResource.py

@@ -11,9 +11,10 @@ from mage_ai.orchestration.db.models.oauth import Role, User
 from mage_ai.settings import (
     AUTHENTICATION_MODE,
     OAUTH_DEFAULT_ACCESS,
+    get_bool_value,
     get_settings_value,
 )
-from mage_ai.settings.keys import LDAP_DEFAULT_ACCESS
+from mage_ai.settings.keys import LDAP_DEFAULT_ACCESS, UPDATE_ROLES_ON_LOGIN
 from mage_ai.usage_statistics.logger import UsageStatisticLogger
 
 
@@ -29,6 +30,9 @@ class SessionResource(BaseResource):
 
         oauth_client = kwargs.get('oauth_client')
 
+        update_roles_on_login = get_bool_value(
+            get_settings_value(UPDATE_ROLES_ON_LOGIN, default='False'))
+
         # Oauth sign in
         if token and provider:
             roles = []
@@ -57,6 +61,9 @@ class SessionResource(BaseResource):
                         email=email,
                         roles_new=roles,
                     )
+                elif update_roles_on_login and roles:
+                    user.update(roles_new=roles)
+
                 oauth_token = generate_access_token(user, oauth_client)
                 return self(oauth_token, user, **kwargs)
 
@@ -109,6 +116,11 @@ class SessionResource(BaseResource):
                     roles_new=roles,
                     username=email,
                 )
+            elif update_roles_on_login:
+                role_names = conn.get_user_roles(user_attributes)
+                if role_names:
+                    roles = Role.query.filter(Role.name.in_(role_names)).all()
+                    user.update(roles_new=roles)
 
             oauth_token = generate_access_token(user, kwargs['oauth_client'])
             return self(oauth_token, user, **kwargs)

mage_ai/authentication/permissions/constants.py

@@ -27,6 +27,7 @@ class EntityName(str, Enum):
     CustomTemplate = 'CustomTemplate'
     DataProvider = 'DataProvider'
     Database = 'Database'
+    Download = 'Download'
     EventMatcher = 'EventMatcher'
     EventRule = 'EventRule'
     ExecutionState = 'ExecutionState'
@@ -68,6 +69,7 @@ class EntityName(str, Enum):
     Scheduler = 'Scheduler'
     SearchResult = 'SearchResult'
     Secret = 'Secret'
+    Seed = 'Seed'
     Session = 'Session'
     SparkApplication = 'SparkApplication'
     SparkEnvironment = 'SparkEnvironment'

mage_ai/authentication/permissions/seed.py

@@ -171,13 +171,17 @@ async def action_rules(policy_class):
             if scope not in operations_mapping:
                 operations_mapping[scope] = {}
 
-            condition = config2.get('condition')
-            key = await get_key(policy_class, condition)
+            if not isinstance(config2, list):
+                config2 = [config2]
 
-            if key not in operations_mapping[scope]:
-                operations_mapping[scope][key] = []
+            for c in config2:
+                condition = c.get('condition')
+                key = await get_key(policy_class, condition)
+
+                if key not in operations_mapping[scope]:
+                    operations_mapping[scope][key] = []
 
-            operations_mapping[scope][key].append(operation)
+                operations_mapping[scope][key].append(operation)
 
     return operations_mapping
 
@@ -202,34 +206,40 @@ async def attribute_rules(policy_class, rules):
                 if config2 is None:
                     continue
 
-                condition = config2.get('condition') if config2 else None
-                key = await get_key(policy_class, condition)
+                if not isinstance(config2, list):
+                    config2 = [config2]
+
+                for c in config2:
+                    condition = c.get('condition')
+                    key = await get_key(policy_class, condition)
 
-                if key not in mapping[scope]:
-                    mapping[scope][key] = {}
+                    if key not in mapping[scope]:
+                        mapping[scope][key] = {}
 
-                if operation not in mapping[scope][key]:
-                    mapping[scope][key][operation] = []
+                    if operation not in mapping[scope][key]:
+                        mapping[scope][key][operation] = []
 
-                mapping[scope][key][operation].append(attribute)
+                    mapping[scope][key][operation].append(attribute)
 
     return mapping
 
 
-async def bootstrap_permissions():
+async def bootstrap_permissions(policy_names: str = None):
     action_rules_mapping = {}
     query_rules_mapping = {}
     read_rules_mapping = {}
     write_rules_mapping = {}
 
-    policy_names = []
-    for n in os.listdir(policies.__path__[0]):
-        if n.endswith('Policy.py') and n not in [
-            'BasePolicy.py',
-            'SeedPolicy.py',
-            'UserPolicy.py',
-        ]:
-            policy_names.append(n.replace('Policy.py', ''))
+    if policy_names is None:
+        policy_names = []
+        for n in os.listdir(policies.__path__[0]):
+            if n.endswith('Policy.py') and n not in [
+                'AsyncBasePolicy.py',
+                'BasePolicy.py',
+                'SeedPolicy.py',
+                'UserPolicy.py',
+            ]:
+                policy_names.append(n.replace('Policy.py', ''))
 
     for policy_name in policy_names:
         policy_class = getattr(
@@ -238,6 +248,7 @@ async def bootstrap_permissions():
         )
 
         model_name = policy_class.model_name()
+        print(f'Processing {model_name}...')
         action_rules_mapping[model_name] = await action_rules(policy_class)
         query_rules_mapping[model_name] = await attribute_rules(
             policy_class,

mage_ai/authentication/providers/active_directory.py

@@ -34,10 +34,11 @@ class ADProvider(SsoProvider, OauthProvider):
         self.client_secret = get_settings_value(ACTIVE_DIRECTORY_CLIENT_SECRET)
         self.__validate()
 
-        self.roles_mapping = get_settings_value(ACTIVE_DIRECTORY_ROLES_MAPPING)
-        if self.roles_mapping:
+        self.roles_mapping = {}
+        roles_mapping = get_settings_value(ACTIVE_DIRECTORY_ROLES_MAPPING)
+        if roles_mapping:
             try:
-                self.roles_mapping = json.loads(self.roles_mapping)
+                self.roles_mapping = json.loads(roles_mapping)
             except Exception:
                 logger.exception('Failed to parse roles mapping.')
 

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/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,6 +76,8 @@ 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:
@@ -74,6 +85,7 @@ class StreamingPipelineExecutor(PipelineExecutor):
         # 2. Support flink pipeline
 
         tags = self.build_tags(**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

mage_ai/data_preparation/models/block/__init__.py

@@ -907,10 +907,11 @@ class Block(DataIntegrationMixin, SparkBlock, ProjectPlatformAccessible):
         parts = get_path_parts(file_path)
 
         if parts and len(parts) >= 3:
+            from mage_ai.data_preparation.models.block.block_factory import BlockFactory
+
             # If file_path == transformers/test4.py
             # parts ==
             # ('/home/src/default_repo/default_platform2/project3', 'transformers', 'test4.py')
-
             # If project platform platform activated, then parts ==
             # ('/home/src', 'default_repo', 'data_loaders/astral_violet.py')
 
@@ -929,7 +930,7 @@ class Block(DataIntegrationMixin, SparkBlock, ProjectPlatformAccessible):
             configuration = dict(file_path=file_path, file_source=dict(path=file_path))
             language = FILE_EXTENSION_TO_BLOCK_LANGUAGE.get(extension)
 
-            return self.get_block(
+            return BlockFactory.get_block(
                 block_uuid,
                 block_uuid,
                 block_type,
@@ -3473,12 +3474,15 @@ df = get_variable('{self.pipeline.uuid}', '{self.uuid}', 'df')
 
             cache = BlockCache()
             if detach:
+                from mage_ai.data_preparation.models.block.block_factory import (
+                    BlockFactory,
+                )
                 """"
                 New block added to pipeline, so it must be added to the block cache.
                 Old block no longer in pipeline, so it must be removed from block cache.
                 """
                 cache.add_pipeline(self, self.pipeline)
-                old_block = self.get_block(
+                old_block = BlockFactory.get_block(
                     old_uuid,
                     old_uuid,
                     self.type,

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

@@ -202,18 +202,29 @@ class DataIntegrationMixin:
             with open(catalog_full_path, mode='w') as f:
                 f.write(json.dumps(catalog))
 
-    def is_data_integration(self) -> bool:
+    def is_data_integration(self, pipeline_project: Project = None) -> bool:
         """
         Check if the block is a data integration block.
         If the data_integration_in_batch_pipeline feature is not enabled, return False.
 
+        Args:
+            pipeline_project (Project, optional): A cached Project value to avoid
+                looking it up many times when called inside loops. Defaults to None.
+
         Returns:
             bool: True if it's a data integration block, False otherwise.
         """
-        if not self.pipeline or not \
-                Project(self.pipeline.repo_config).is_feature_enabled(
-                    FeatureUUID.DATA_INTEGRATION_IN_BATCH_PIPELINE,
-                ):
+        if not self.pipeline:
+
+            return False
+
+        actual_project: Project = pipeline_project
+        if not actual_project:
+            actual_project = Project(self.pipeline.repo_config)
+
+        if not actual_project.is_feature_enabled(
+                        FeatureUUID.DATA_INTEGRATION_IN_BATCH_PIPELINE,
+                    ):
 
             return False
 

mage_ai/data_preparation/models/block/dynamic/child.py

@@ -229,3 +229,6 @@ class DynamicChildController:
             block_runs.append(block_run)
 
         return block_runs
+
+    def run_tests(self, **kwargs):
+        pass

mage_ai/data_preparation/models/block/dynamic/variables.py

@@ -6,7 +6,7 @@ from typing import Any, Dict, List, Tuple, Union
 
 import pandas as pd
 
-from mage_ai.data_preparation.models.constants import BlockLanguage
+from mage_ai.data_preparation.models.constants import BlockLanguage, BlockType
 from mage_ai.data_preparation.models.variable import Variable
 from mage_ai.shared.memory import get_memory_usage, get_memory_usage_async
 from mage_ai.shared.strings import to_ordinal_integers
@@ -470,7 +470,7 @@ def fetch_input_variables_for_dynamic_upstream_blocks(
 
             # If dynamic child should reduce its output (which means it passes the entire
             # output to its downstream blocks):
-            if should_reduce_output(upstream_block):
+            if should_reduce_output(upstream_block) and block.type != BlockType.EXTENSION:
                 child_data = []
                 metadata = {}
                 for lazy_variable_set in lazy_variable_controller:

mage_ai/data_preparation/models/block/extension/utils.py

@@ -45,4 +45,5 @@ def handle_run_tests(
             global_vars=global_vars,
             logger=logger,
             logging_tags=logging_tags,
+            update_status=False,
         )

mage_ai/data_preparation/models/block/integration/__init__.py

@@ -155,7 +155,7 @@ class IntegrationBlock(Block):
                 proc = subprocess.Popen(args, stdout=subprocess.PIPE, stderr=subprocess.STDOUT)
 
                 for line in proc.stdout:
-                    f.write(line.decode()),
+                    f.write(line.decode())
                     print_log_from_line(
                         line,
                         config=config,

mage_ai/data_preparation/models/block/sql/__init__.py

@@ -337,7 +337,7 @@ def execute_sql_code(
             if should_query:
                 return [
                     loader.load(
-                        f'SELECT * FROM {table_name}',
+                        f'SELECT * FROM {schema}.{table_name}',
                         verbose=False,
                     ),
                 ]