mage-ai 0.8.98__py3-none-any.whl → 0.8.99__py3-none-any.whl

mage_ai/api/policies/OauthPolicy.py

@@ -10,9 +10,16 @@ class OauthPolicy(BasePolicy):
 
 OauthPolicy.allow_actions([
     constants.CREATE,
+], scopes=[
+    OauthScope.CLIENT_PRIVATE,
+], condition=lambda policy: policy.has_at_least_viewer_role())
+
+
+OauthPolicy.allow_actions([
     constants.DETAIL,
 ], scopes=[
     OauthScope.CLIENT_PRIVATE,
+    OauthScope.CLIENT_PUBLIC,
 ], condition=lambda policy: policy.has_at_least_viewer_role())
 
 
@@ -20,6 +27,13 @@ OauthPolicy.allow_read(OauthPresenter.default_attributes, scopes=[
     OauthScope.CLIENT_PRIVATE,
 ], on_action=[
     constants.CREATE,
+], condition=lambda policy: policy.has_at_least_viewer_role())
+
+
+OauthPolicy.allow_read(OauthPresenter.default_attributes, scopes=[
+    OauthScope.CLIENT_PRIVATE,
+    OauthScope.CLIENT_PUBLIC,
+], on_action=[
     constants.DETAIL,
 ], condition=lambda policy: policy.has_at_least_viewer_role())
 
@@ -40,4 +54,5 @@ OauthPolicy.allow_query([
     constants.DETAIL,
 ], scopes=[
     OauthScope.CLIENT_PRIVATE,
+    OauthScope.CLIENT_PUBLIC,
 ], condition=lambda policy: policy.has_at_least_viewer_role())

mage_ai/api/policies/SessionPolicy.py

@@ -46,6 +46,8 @@ SessionPolicy.allow_write([
     'email',
     'password',
     'username',
+    'token',
+    'provider',
 ], scopes=[
     OauthScope.CLIENT_PUBLIC,
 ], on_action=[

mage_ai/api/presenters/PipelinePresenter.py

@@ -12,6 +12,7 @@ class PipelinePresenter(BasePresenter):
         'executor_count',
         'executor_type',
         'name',
+        'notification_config',
         'spark_config',
         'type',
         'updated_at',

mage_ai/api/presenters/SyncPresenter.py

@@ -12,6 +12,7 @@ class SyncPresenter(BasePresenter):
         'ssh_private_key',
         'ssh_public_key',
         'sync_on_pipeline_run',
+        'sync_on_start',
         'type',
         'user_git_settings',
         'username',

mage_ai/api/resources/OauthResource.py

@@ -1,3 +1,4 @@
+import json
 import urllib.parse
 from datetime import datetime, timedelta
 
@@ -5,13 +6,17 @@ from mage_ai.api.errors import ApiError
 from mage_ai.api.resources.GenericResource import GenericResource
 from mage_ai.authentication.oauth2 import generate_access_token
 from mage_ai.authentication.oauth.constants import (
+    ACTIVE_DIRECTORY_CLIENT_ID,
     GITHUB_CLIENT_ID,
     GITHUB_STATE,
+    OAUTH_PROVIDER_ACTIVE_DIRECTORY,
     OAUTH_PROVIDER_GITHUB,
+    VALID_OAUTH_PROVIDERS,
 )
 from mage_ai.authentication.oauth.utils import access_tokens_for_provider
 from mage_ai.orchestration.db import safe_db_query
 from mage_ai.orchestration.db.models.oauth import Oauth2AccessToken, Oauth2Application
+from mage_ai.settings import ACTIVE_DIRECTORY_DIRECTORY_ID
 
 
 class OauthResource(GenericResource):
@@ -23,7 +28,7 @@ class OauthResource(GenericResource):
         provider = payload.get('provider')
         token = payload.get('token')
 
-        if not provider or provider not in [OAUTH_PROVIDER_GITHUB]:
+        if not provider or provider not in VALID_OAUTH_PROVIDERS:
             error.update(dict(message='Invalid provider.'))
             raise ApiError(error)
 
@@ -66,7 +71,7 @@ class OauthResource(GenericResource):
         model = dict(provider=pk)
 
         error = ApiError.RESOURCE_INVALID.copy()
-        if pk not in [OAUTH_PROVIDER_GITHUB]:
+        if pk not in VALID_OAUTH_PROVIDERS:
             error.update(dict(message='Invalid provider.'))
             raise ApiError(error)
 
@@ -95,5 +100,26 @@ class OauthResource(GenericResource):
                     query_strings.append(f'{k}={v}')
 
                 model['url'] = f"https://github.com/login/oauth/authorize?{'&'.join(query_strings)}"
+            elif OAUTH_PROVIDER_ACTIVE_DIRECTORY == pk:
+                ad_directory_id = ACTIVE_DIRECTORY_DIRECTORY_ID
+                if ad_directory_id:
+                    mage_redirect_uri = '?'.join([
+                        redirect_uri,
+                        f'provider={pk}',
+                    ])
+                    query = dict(
+                        client_id=ACTIVE_DIRECTORY_CLIENT_ID,
+                        redirect_uri=f'https://api.mage.ai/v1/oauth/{pk}',
+                        response_type='code',
+                        scope='User.Read',
+                        state=urllib.parse.quote_plus(json.dumps(dict(
+                            redirect_uri=mage_redirect_uri,
+                            tenant_id=ad_directory_id,
+                        ))),
+                    )
+                    query_strings = []
+                    for k, v in query.items():
+                        query_strings.append(f'{k}={v}')
+                    model['url'] = f"https://login.microsoftonline.com/{ad_directory_id}/oauth2/v2.0/authorize?{'&'.join(query_strings)}"  # noqa: E501
 
         return self(model, user, **kwargs)

mage_ai/api/resources/PipelineResource.py

@@ -1,5 +1,6 @@
 import asyncio
 
+from sqlalchemy import or_
 from sqlalchemy.orm import aliased
 
 from mage_ai.api.operations.constants import DELETE
@@ -81,7 +82,10 @@ class PipelineResource(BaseResource):
                 ]).
                 filter(
                     a.pipeline_uuid.in_(pipeline_uuids),
-                    a.repo_path == get_repo_path(),
+                    or_(
+                        a.repo_path == get_repo_path(),
+                        a.repo_path.is_(None),
+                    )
                 )
             ).all()
             return group_by(lambda x: x.pipeline_uuid, result)

mage_ai/api/resources/SessionResource.py

@@ -4,6 +4,8 @@ from mage_ai.api.errors import ApiError
 from mage_ai.api.resources.BaseResource import BaseResource
 from mage_ai.authentication.ldap import new_ldap_connection
 from mage_ai.authentication.oauth2 import encode_token, generate_access_token
+from mage_ai.authentication.oauth.active_directory import get_user_info
+from mage_ai.authentication.oauth.constants import OAUTH_PROVIDER_ACTIVE_DIRECTORY
 from mage_ai.authentication.passwords import verify_password
 from mage_ai.orchestration.db import safe_db_query
 from mage_ai.orchestration.db.models.oauth import Role, User
@@ -18,6 +20,22 @@ class SessionResource(BaseResource):
         email = payload.get('email')
         password = payload.get('password')
         username = payload.get('username')
+        token = payload.get('token')
+        provider = payload.get('provider')
+
+        if token and provider:
+            if provider == OAUTH_PROVIDER_ACTIVE_DIRECTORY:
+                user_info = get_user_info(token)
+                principal_name = user_info.get('userPrincipalName')
+                user = User.query.filter(User.email == principal_name).first()
+                if not user:  # noqa: E712
+                    print('first user login, creating user.')
+                    user = User.create(
+                        username=principal_name,
+                        email=principal_name,
+                    )
+                oauth_token = generate_access_token(user, kwargs['oauth_client'])
+                return self(oauth_token, user, **kwargs)
 
         error = ApiError.RESOURCE_NOT_FOUND
         error.update({'message': 'Email/username and/or password invalid.'})

mage_ai/api/resources/WorkspaceResource.py

@@ -46,7 +46,8 @@ class WorkspaceResource(GenericResource):
         query_user = None
         if user_id:
             user_id = user_id[0]
-            query_user = User.query.get(user_id)
+            if user_id:
+                query_user = User.query.get(user_id)
 
         instances = self.get_instances(cluster_type)
         instance_map = {
@@ -313,7 +314,7 @@ class WorkspaceResource(GenericResource):
             namespace = os.getenv(KUBE_NAMESPACE)
             workload_manager = WorkloadManager(namespace)
 
-            instances = workload_manager.list_services()
+            instances = workload_manager.list_workloads()
         elif cluster_type == ClusterType.ECS:
             from mage_ai.cluster_manager.aws.ecs_task_manager import EcsTaskManager
 

mage_ai/authentication/oauth/active_directory.py

@@ -0,0 +1,17 @@
+from typing import Dict
+
+import requests
+
+
+def get_user_info(token: str) -> Dict:
+    url = 'https://graph.microsoft.com/v1.0/me'
+
+    headers = {
+        'Content-Type': 'application\\json',
+        'Authorization': 'Bearer {}'.format(token)
+    }
+
+    resp = requests.get(url, headers=headers)
+    result = resp.json()
+
+    return result

mage_ai/authentication/oauth/constants.py

@@ -1,3 +1,12 @@
+ACTIVE_DIRECTORY_CLIENT_ID = '51aec820-9d49-40a9-b046-17c1f28f620d'
+
 GITHUB_CLIENT_ID = '8577f13ddc81e2848b07'
 GITHUB_STATE = '1337'
+
+OAUTH_PROVIDER_ACTIVE_DIRECTORY = 'active_directory'
 OAUTH_PROVIDER_GITHUB = 'github'
+
+VALID_OAUTH_PROVIDERS = [
+    OAUTH_PROVIDER_ACTIVE_DIRECTORY,
+    OAUTH_PROVIDER_GITHUB,
+]

mage_ai/cli/main.py

@@ -10,6 +10,7 @@ from typer.core import TyperGroup
 
 from mage_ai.cli.utils import parse_runtime_variables
 from mage_ai.data_preparation.repo_manager import ProjectType
+from mage_ai.services.newrelic import initialize_new_relic
 from mage_ai.shared.constants import InstanceType
 
 
@@ -191,6 +192,9 @@ def run(
     project_path = os.path.abspath(project_path)
     set_repo_path(project_path)
 
+    from contextlib import nullcontext
+
+    import newrelic.agent
     import sentry_sdk
 
     from mage_ai.data_preparation.executors.executor_factory import ExecutorFactory
@@ -207,55 +211,58 @@ def run(
             sentry_dsn,
             traces_sample_rate=SENTRY_TRACES_SAMPLE_RATE,
         )
-
-    runtime_variables = dict()
-    if runtime_vars is not None:
-        runtime_variables = parse_runtime_variables(runtime_vars)
-
-    sys.path.append(os.path.dirname(project_path))
-    pipeline = Pipeline.get(pipeline_uuid, repo_path=project_path)
-
-    db_connection.start_session()
-
-    if pipeline_run_id is None:
-        default_variables = get_global_variables(pipeline_uuid)
-        global_vars = merge_dict(default_variables, runtime_variables)
-    else:
-        pipeline_run = PipelineRun.query.get(pipeline_run_id)
-        global_vars = pipeline_run.get_variables(extra_variables=runtime_variables)
-
-    if template_runtime_configuration is not None:
-        template_runtime_configuration = json.loads(template_runtime_configuration)
-
-    if block_uuid is None:
-        ExecutorFactory.get_pipeline_executor(
-            pipeline,
-            execution_partition=execution_partition,
-            executor_type=executor_type,
-        ).execute(
-            analyze_outputs=False,
-            global_vars=global_vars,
-            pipeline_run_id=pipeline_run_id,
-            run_sensors=not skip_sensors,
-            run_tests=test,
-            update_status=False,
-        )
-    else:
-        ExecutorFactory.get_block_executor(
-            pipeline,
-            block_uuid,
-            execution_partition=execution_partition,
-            executor_type=executor_type,
-        ).execute(
-            analyze_outputs=False,
-            block_run_id=block_run_id,
-            callback_url=callback_url,
-            global_vars=global_vars,
-            pipeline_run_id=pipeline_run_id,
-            template_runtime_configuration=template_runtime_configuration,
-            update_status=False,
-        )
-    print('Pipeline run completed.')
+    (enable_new_relic, application) = initialize_new_relic()
+
+    with newrelic.agent.BackgroundTask(application, name="mage-run", group='Task') \
+         if enable_new_relic else nullcontext():
+        runtime_variables = dict()
+        if runtime_vars is not None:
+            runtime_variables = parse_runtime_variables(runtime_vars)
+
+        sys.path.append(os.path.dirname(project_path))
+        pipeline = Pipeline.get(pipeline_uuid, repo_path=project_path)
+
+        db_connection.start_session()
+
+        if pipeline_run_id is None:
+            default_variables = get_global_variables(pipeline_uuid)
+            global_vars = merge_dict(default_variables, runtime_variables)
+        else:
+            pipeline_run = PipelineRun.query.get(pipeline_run_id)
+            global_vars = pipeline_run.get_variables(extra_variables=runtime_variables)
+
+        if template_runtime_configuration is not None:
+            template_runtime_configuration = json.loads(template_runtime_configuration)
+
+        if block_uuid is None:
+            ExecutorFactory.get_pipeline_executor(
+                pipeline,
+                execution_partition=execution_partition,
+                executor_type=executor_type,
+            ).execute(
+                analyze_outputs=False,
+                global_vars=global_vars,
+                pipeline_run_id=pipeline_run_id,
+                run_sensors=not skip_sensors,
+                run_tests=test,
+                update_status=False,
+            )
+        else:
+            ExecutorFactory.get_block_executor(
+                pipeline,
+                block_uuid,
+                execution_partition=execution_partition,
+                executor_type=executor_type,
+            ).execute(
+                analyze_outputs=False,
+                block_run_id=block_run_id,
+                callback_url=callback_url,
+                global_vars=global_vars,
+                pipeline_run_id=pipeline_run_id,
+                template_runtime_configuration=template_runtime_configuration,
+                update_status=False,
+            )
+        print('Pipeline run completed.')
 
 
 @app.command()

mage_ai/cluster_manager/kubernetes/workload_manager.py

@@ -26,6 +26,7 @@ from mage_ai.orchestration.constants import (
     DB_USER,
 )
 from mage_ai.settings import MAGE_SETTINGS_ENVIRONMENT_VARIABLES
+from mage_ai.shared.array import find
 
 
 class WorkloadManager:
@@ -53,24 +54,57 @@ class WorkloadManager:
 
         return False
 
-    def list_services(self):
+    def list_workloads(self):
         services = self.core_client.list_namespaced_service(self.namespace).items
-        services_list = []
+        workloads_list = []
+
+        pods = self.core_client.list_namespaced_pod(self.namespace).items
+        pod_map = dict()
+        for pod in pods:
+            try:
+                name = pod.metadata.labels.get('app')
+                pod_map[name] = pod
+            except Exception:
+                pass
         for service in services:
             try:
                 labels = service.metadata.labels
                 if not labels.get('dev-instance'):
                     continue
-                conditions = service.status.conditions or list()
-                services_list.append(dict(
+                name = labels.get('app')
+                pod = pod_map[name]
+                service_type = service.spec.type
+                workload = dict(
                     name=labels.get('app'),
-                    status='RUNNING' if len(conditions) == 0 else conditions[0].status,
-                    type='kubernetes',
-                ))
+                    type=service_type,
+                )
+                if pod:
+                    status = pod.status.phase
+                    workload['status'] = status.upper()
+
+                    node_name = pod.spec.node_name
+                    ip = None
+                    if service_type == 'NodePort':
+                        try:
+                            if node_name:
+                                items = self.core_client.list_node(
+                                    field_selector=f'metadata.name={node_name}').items
+                                node = items[0]
+                                ip = find(
+                                    lambda a: a.type == 'ExternalIP',
+                                    node.status.addresses
+                                ).address
+                                if ip:
+                                    node_port = service.spec.ports[0].node_port
+                                    workload['ip'] = f'{ip}:{node_port}'
+                        except Exception:
+                            pass
+
+                workloads_list.append(workload)
             except Exception:
                 pass
 
-        return services_list
+        return workloads_list
 
     def create_workload(
         self,
@@ -88,7 +122,10 @@ class WorkloadManager:
             'service_account_name',
             os.getenv(KUBE_SERVICE_ACCOUNT_NAME),
         )
-        storage_class_name = kwargs.get('storage_class_name', os.getenv(KUBE_STORAGE_CLASS_NAME))
+        storage_class_name = kwargs.get(
+            'storage_class_name',
+            os.getenv(KUBE_STORAGE_CLASS_NAME, 'default'),
+        )
         storage_access_mode = kwargs.get('storage_access_mode', 'ReadWriteOnce')
         storage_request_size = kwargs.get('storage_request_size', 2)
 

mage_ai/data_preparation/git/__init__.py

@@ -486,7 +486,7 @@ class Git:
                     self._run_command(cmd)
                 print(f'Installing {requirements_file} completed successfully.')
             except Exception as err:
-                print(f'Skip installing {requirements_file} due to error: {err}')
+                print(f'Skip installing {requirements_file} due to error: {str(err)}')
                 pass
 
     def __create_ssh_keys(self) -> str:

mage_ai/data_preparation/models/block/dbt/utils/__init__.py

@@ -28,6 +28,7 @@ from mage_ai.data_preparation.models.block.sql import (
     snowflake,
     spark,
     trino,
+    clickhouse,
 )
 from mage_ai.data_preparation.models.constants import BlockLanguage, BlockType
 from mage_ai.data_preparation.repo_manager import get_repo_path
@@ -597,6 +598,23 @@ def config_file_loader_and_configuration(block, profile_target: str) -> Dict:
             data_provider_schema=schema,
             export_write_policy=ExportWritePolicy.REPLACE,
         )
+    elif DataSource.CLICKHOUSE == profile_type:
+        database = profile.get('schema')
+        interface = profile.get('driver')
+
+        config_file_loader = ConfigFileLoader(config=dict(
+            CLICKHOUSE_DATABASE=database,
+            CLICKHOUSE_HOST=profile.get('host'),
+            CLICKHOUSE_INTERFACE=interface,
+            CLICKHOUSE_PASSWORD=profile.get('password'),
+            CLICKHOUSE_PORT=profile.get('port'),
+            CLICKHOUSE_USERNAME=profile.get('user'),
+        ))
+        configuration = dict(
+            data_provider=profile_type,
+            data_provider_database=database,
+            export_write_policy=ExportWritePolicy.REPLACE,
+        )
 
     if not config_file_loader or not configuration:
         attr = parse_attributes(block)
@@ -733,6 +751,15 @@ def create_upstream_tables(
                 block,
                 **kwargs_shared,
             )
+    elif DataSource.CLICKHOUSE == data_provider:
+        from mage_ai.io.clickhouse import ClickHouse
+
+        loader = ClickHouse.with_config(config_file_loader)
+        clickhouse.create_upstream_block_tables(
+            loader,
+            block,
+            **kwargs_shared,
+        )
 
     block.upstream_blocks = upstream_blocks_init
 
@@ -930,6 +957,11 @@ def execute_query(
 
         with Trino.with_config(config_file_loader) as loader:
             return loader.load(query_string, **shared_kwargs)
+    elif DataSource.CLICKHOUSE == data_provider:
+        from mage_ai.io.clickhouse import ClickHouse
+
+        loader = ClickHouse.with_config(config_file_loader)
+        return loader.load(query_string, **shared_kwargs)
 
 
 def query_from_compiled_sql(block, profile_target: str, limit: int = None) -> DataFrame:

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

@@ -137,8 +137,9 @@ def execute_sql_code(
                 NotFound: 404 Not found: Table database:schema.table_name
                     was not found in location XX
                 """
+                total_retries = 5
                 tries = 0
-                while tries < 10:
+                while tries < total_retries:
                     sleep(tries)
                     tries += 1
                     try:
@@ -149,7 +150,7 @@ def execute_sql_code(
                         )
                         return [result]
                     except Exception as err:
-                        if '404' not in str(err):
+                        if '404' not in str(err) or tries == total_retries:
                             raise err
     elif DataSource.CLICKHOUSE.value == data_provider:
         from mage_ai.io.clickhouse import ClickHouse

mage_ai/data_preparation/models/pipeline.py

@@ -53,6 +53,7 @@ class Pipeline:
         self.executor_type = None
         self.executor_config = dict()
         self.name = None
+        self.notification_config = dict()
         self.repo_path = repo_path or get_repo_path()
         self.schedules = []
         self.uuid = uuid
@@ -425,7 +426,8 @@ class Pipeline:
         self.callback_configs = config.get('callbacks') or []
         self.conditional_configs = config.get('conditionals') or []
         self.executor_type = config.get('executor_type')
-        self.executor_config = config.get('executor_confid') or dict()
+        self.executor_config = config.get('executor_config') or dict()
+        self.notification_config = config.get('notification_config') or dict()
         self.spark_config = config.get('spark_config') or dict()
         self.widget_configs = config.get('widgets') or []
 
@@ -543,6 +545,7 @@ class Pipeline:
             executor_count=self.executor_count,
             executor_type=self.executor_type,
             name=self.name,
+            notification_config=self.notification_config,
             type=self.type.value if type(self.type) is not str else self.type,
             updated_at=self.updated_at,
             uuid=self.uuid,

mage_ai/data_preparation/preferences.py

@@ -16,7 +16,8 @@ GIT_USERNAME_VAR = 'GIT_USERNAME'
 GIT_EMAIL_VAR = 'GIT_EMAIL'
 GIT_AUTH_TYPE_VAR = 'GIT_AUTH_TYPE'
 GIT_BRANCH_VAR = 'GIT_BRANCH'
-GIT_SYNC_ON_PIPELINE_RUN_TYPE = 'GIT_SYNC_ON_PIPELINE_RUN'
+GIT_SYNC_ON_PIPELINE_RUN_VAR = 'GIT_SYNC_ON_PIPELINE_RUN'
+GIT_SYNC_ON_START_VAR = 'GIT_SYNC_ON_PIPELINE_RUN'
 
 
 class Preferences:
@@ -52,7 +53,8 @@ class Preferences:
                 username=os.getenv(GIT_USERNAME_VAR),
                 email=os.getenv(GIT_EMAIL_VAR),
                 branch=os.getenv(GIT_BRANCH_VAR),
-                sync_on_pipeline_run=bool(int(os.getenv(GIT_SYNC_ON_PIPELINE_RUN_TYPE) or 0)),
+                sync_on_pipeline_run=bool(int(os.getenv(GIT_SYNC_ON_PIPELINE_RUN_VAR) or 0)),
+                sync_on_start=bool(int(os.getenv(GIT_SYNC_ON_START_VAR) or 0)),
             )
         else:
             project_sync_config = project_preferences.get('sync_config', dict())

mage_ai/data_preparation/repo_manager.py

@@ -213,7 +213,11 @@ def get_repo_config(repo_path=None) -> RepoConfig:
 
 
 def get_project_type(repo_path=None) -> ProjectType:
-    return get_repo_config(repo_path=repo_path).project_type
+    try:
+        return get_repo_config(repo_path=repo_path).project_type
+    except Exception:
+        # default to standalone project type
+        return ProjectType.STANDALONE
 
 
 def set_repo_path(repo_path: str) -> None:

mage_ai/data_preparation/sync/__init__.py

@@ -21,6 +21,7 @@ class GitConfig(BaseConfig):
     repo_path: str = os.getcwd()
     branch: str = 'main'
     sync_on_pipeline_run: bool = False
+    sync_on_start: bool = False
     auth_type: AuthType = AuthType.SSH
     # User settings moved to UserGitConfig, these will be used for Git syncs
     username: str = ''

mage_ai/data_preparation/sync/git_sync.py

@@ -11,12 +11,9 @@ class GitSync(BaseSync):
         self.git_manager = Git(sync_config)
 
     def sync_data(self):
-        with VerboseFunctionExec(
-            f'Syncing data with remote repo {self.remote_repo_link}',
-            verbose=True,
-        ):
-            self.git_manager.reset(self.branch)
+        self.git_manager.reset(self.branch)
 
+    # Reset git sync by cloning the remote repo
     def reset(self):
         with VerboseFunctionExec(
             f'Attempting to clone from remote repo {self.remote_repo_link}',

mage_ai/data_preparation/templates/utils.py

@@ -51,6 +51,8 @@ def write_template(template_source: str, dest_path: str) -> None:
         template_source (str): Template source code to write to file
         dest_path (str): Destination file to write template source code to.
     """
+    os.makedirs(os.path.dirname(dest_path), exist_ok=True)
+
     with open(dest_path, 'w') as foutput:
         foutput.write(template_source)