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

mage_ai/data_preparation/git/api.py

@@ -0,0 +1,175 @@
+import asyncio
+import requests
+import subprocess
+from datetime import datetime
+from git.remote import RemoteProgress
+from git.repo.base import Repo
+from mage_ai.authentication.oauth.constants import OAUTH_PROVIDER_GITHUB
+from mage_ai.orchestration.db.models.oauth import Oauth2AccessToken, Oauth2Application, User
+from typing import Dict
+from urllib.parse import urlsplit, urlunsplit
+
+
+API_ENDPOINT = 'https://api.github.com'
+
+
+def get_access_token_for_user(user: User) -> Oauth2AccessToken:
+    oauth_client = Oauth2Application.query.filter(
+        Oauth2Application.client_id == OAUTH_PROVIDER_GITHUB,
+    ).first()
+
+    if oauth_client:
+        access_token = Oauth2AccessToken.query.filter(
+            Oauth2AccessToken.expires > datetime.utcnow(),
+            Oauth2AccessToken.oauth2_application_id == oauth_client.id,
+        ).first()
+
+        return access_token
+
+
+def fetch(remote_name: str, remote_url: str, token: str) -> RemoteProgress:
+    from mage_ai.data_preparation.git import Git
+
+    custom_progress = RemoteProgress()
+    username = get_username(token)
+
+    url = build_authenticated_remote_url(remote_url, username, token)
+    git_manager = Git.get_manager()
+
+    remote = git_manager.repo.remotes[remote_name]
+    url_original = list(remote.urls)[0]
+    remote.set_url(url)
+
+    try:
+        remote.fetch(progress=custom_progress)
+    except Exception as err:
+        raise err
+    finally:
+        try:
+            remote.set_url(url_original)
+        except Exception as err:
+            print('WARNING (mage_ai.data_preparation.git.api):')
+            print(err)
+
+    return custom_progress
+
+
+def pull(remote_name: str, remote_url: str, branch_name: str, token: str) -> RemoteProgress:
+    from mage_ai.data_preparation.git import Git
+
+    custom_progress = RemoteProgress()
+    username = get_username(token)
+
+    url = build_authenticated_remote_url(remote_url, username, token)
+    git_manager = Git.get_manager()
+
+    remote = git_manager.repo.remotes[remote_name]
+    url_original = list(remote.urls)[0]
+    remote.set_url(url)
+
+    try:
+        remote.pull(branch_name, custom_progress)
+    except Exception as err:
+        raise err
+    finally:
+        try:
+            remote.set_url(url_original)
+        except Exception as err:
+            print('WARNING (mage_ai.data_preparation.git.api):')
+            print(err)
+
+    return custom_progress
+
+
+def push(remote_name: str, remote_url: str, branch_name: str, token: str) -> RemoteProgress:
+    from mage_ai.data_preparation.git import Git
+
+    custom_progress = RemoteProgress()
+    username = get_username(token)
+
+    url = build_authenticated_remote_url(remote_url, username, token)
+    git_manager = Git.get_manager()
+
+    remote = git_manager.repo.remotes[remote_name]
+    url_original = list(remote.urls)[0]
+    remote.set_url(url)
+
+    try:
+        remote.push(branch_name, custom_progress)
+    except Exception as err:
+        raise err
+    finally:
+        try:
+            remote.set_url(url_original)
+        except Exception as err:
+            print('WARNING (mage_ai.data_preparation.git.api):')
+            print(err)
+
+    return custom_progress
+
+
+def build_authenticated_remote_url(remote_url: str, username: str, token: str) -> str:
+    # https://[username]:[token]@github.com/[remote_url]
+    url = urlsplit(remote_url)
+    url = url._replace(netloc=f'{username}:{token}@{url.netloc}')
+    return urlunsplit(url)
+
+
+def get_username(token: str) -> str:
+    return get_user(token)['login']
+
+
+def get_user(token: str) -> Dict:
+    """
+    https://docs.github.com/en/rest/users/users?apiVersion=2022-11-28#get-the-authenticated-user
+    """
+    resp = requests.get(f'{API_ENDPOINT}/user', headers={
+        'Accept': 'application/vnd.github+json',
+        'Authorization': f'Bearer {token}',
+        'X-GitHub-Api-Version': '2022-11-28',
+    })
+
+    return resp.json()
+
+
+def check_connection(repo: Repo, remote_url: str) -> None:
+    asyncio.run(validate_authentication_for_remote_url(repo, remote_url))
+
+
+async def validate_authentication_for_remote_url(repo: Repo, remote_url: str) -> None:
+    proc = repo.git.ls_remote(remote_url, as_process=True)
+
+    asyncio.run(__poll_process_with_timeout(
+        proc,
+        error_message='Error connecting to remote, make sure your access is valid.',
+    ))
+
+
+async def __poll_process_with_timeout(
+    proc: subprocess.Popen,
+    error_message: str = None,
+    timeout: int = 10,
+):
+    ct = 0
+    while ct < timeout * 2:
+        return_code = proc.poll()
+        if return_code is not None:
+            proc.kill()
+            break
+        ct += 1
+        await asyncio.sleep(0.5)
+
+    if error_message is None:
+        error_message = 'Error running Git process'
+
+    if return_code is not None and return_code != 0:
+        _, err = proc.communicate()
+        message = (
+            err.decode('UTF-8') if err
+            else error_message
+        )
+        raise ChildProcessError(message)
+
+    if return_code is None:
+        proc.kill()
+        raise TimeoutError(error_message)

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

@@ -1,11 +1,26 @@
+import os
+import re
+import shutil
+import subprocess
+import sys
+import uuid
 from contextlib import redirect_stdout
 from datetime import datetime
-from jinja2 import Template
 from logging import Logger
+from typing import Callable, Dict, List, Tuple
+
+import aiofiles
+import simplejson
+import yaml
+from jinja2 import Template
+from pandas import DataFrame
+
 from mage_ai.data_preparation.models.block import Block
+from mage_ai.data_preparation.models.block.sql import bigquery
 from mage_ai.data_preparation.models.block.sql import (
-    bigquery,
     execute_sql_code as execute_sql_code_orig,
+)
+from mage_ai.data_preparation.models.block.sql import (
     mssql,
     mysql,
     postgres,
@@ -13,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
@@ -27,18 +43,6 @@ from mage_ai.shared.hash import merge_dict
 from mage_ai.shared.parsers import encode_complex
 from mage_ai.shared.strings import remove_extension_from_filename
 from mage_ai.shared.utils import clean_name, files_in_path
-from pandas import DataFrame
-from typing import Callable, Dict, List, Tuple
-import aiofiles
-import os
-import re
-import shutil
-import simplejson
-import subprocess
-import sys
-import uuid
-import yaml
-
 
 PROFILES_FILE_NAME = 'profiles.yml'
 
@@ -594,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)
@@ -730,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
 
@@ -927,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/models/pipelines/integration_pipeline.py

@@ -212,6 +212,8 @@ class IntegrationPipeline(Pipeline):
                         error = dig(json_object, 'tags.error')
                     except Exception:
                         error = line
+                elif not error and line.startswith('CRITICAL'):
+                    error = line
             raise Exception(error)
 
     def preview_data(self, block_type: BlockType, streams: List[str] = None) -> List[str]:
@@ -286,9 +288,11 @@ class IntegrationPipeline(Pipeline):
                     except Exception:
                         error = line
             if not error:
-                raise Exception('The sample data was not able to be loaded. Please check \
-                                if the stream still exists. If it does not, click the "View and \
-                                select streams" button and confirm the valid streams.')
+                raise Exception('The sample data was not able to be loaded. Please check if the ' +
+                                'stream still exists. If it does not, click the "View and select ' +
+                                'streams" button and confirm the valid streams. If it does, ' +
+                                'loading sample data for this source may not currently ' +
+                                'be supported.')
             raise Exception(error)
 
     def count_records(self) -> List[Dict]:

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

@@ -1,7 +1,7 @@
 import os
-import shutil
 import sys
 import traceback
+import uuid
 from enum import Enum
 from typing import Dict
 
@@ -16,7 +16,6 @@ from mage_ai.data_preparation.shared.constants import (
 from mage_ai.data_preparation.templates.utils import copy_template_directory
 from mage_ai.shared.environments import is_test
 
-
 if is_test():
     DEFAULT_MAGE_DATA_DIR = './'
 else:
@@ -155,28 +154,34 @@ class RepoConfig:
             yml.dump(data, f)
 
 
-def init_repo(repo_path: str, project_type: str = ProjectType.STANDALONE) -> None:
+def init_repo(
+    repo_path: str,
+    project_type: str = ProjectType.STANDALONE,
+    cluster_type: str = None,
+    project_uuid: str = None,
+) -> None:
     """
     Initialize a repository under the current path.
     """
     if os.path.exists(repo_path):
         raise FileExistsError(f'Repository {repo_path} already exists')
 
+    new_config = dict()
     if project_type == ProjectType.MAIN:
         copy_template_directory('main', repo_path)
+        new_config.update(
+            cluster_type=cluster_type,
+        )
     elif project_type == ProjectType.SUB:
         os.makedirs(
             os.getenv(MAGE_DATA_DIR_ENV_VAR) or DEFAULT_MAGE_DATA_DIR,
             exist_ok=True,
         )
         copy_template_directory('repo', repo_path)
-        new_repo_config = get_repo_config(repo_path)
-        current_metadata = get_repo_config().metadata_path
-        new_metadata = new_repo_config.metadata_path
-        if os.path.exists(current_metadata):
-            shutil.copyfile(current_metadata, new_metadata)
-        new_repo_config.save(
+        new_config.update(
             project_type=ProjectType.SUB.value,
+            cluster_type=cluster_type,
+            project_uuid=project_uuid,
         )
     else:
         os.makedirs(
@@ -185,6 +190,11 @@ def init_repo(repo_path: str, project_type: str = ProjectType.STANDALONE) -> Non
         )
         copy_template_directory('repo', repo_path)
 
+    if not project_uuid:
+        project_uuid = uuid.uuid4().hex
+    new_config.update(project_uuid=project_uuid)
+    get_repo_config(repo_path).save(**new_config)
+
 
 def get_data_dir() -> str:
     return os.getenv(MAGE_DATA_DIR_ENV_VAR) or DEFAULT_MAGE_DATA_DIR
@@ -203,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:
@@ -213,3 +227,20 @@ def set_repo_path(repo_path: str) -> None:
 
 def get_variables_dir(repo_path: str = None) -> str:
     return get_repo_config(repo_path=repo_path).variables_dir
+
+
+config = get_repo_config()
+project_uuid = config.project_uuid
+
+
+def update_project_uuid():
+    global project_uuid
+    project_uuid = get_repo_config().project_uuid
+    if not project_uuid:
+        puuid = uuid.uuid4().hex
+        config.save(project_uuid=puuid)
+        project_uuid = puuid
+
+
+def get_project_uuid() -> str:
+    return project_uuid

mage_ai/data_preparation/shared/secrets.py

@@ -1,10 +1,9 @@
-from cryptography.fernet import Fernet, InvalidToken
-from mage_ai.data_preparation.repo_manager import (
-    get_data_dir,
-    get_repo_path
-)
-from typing import List
 import os
+from typing import List
+
+from cryptography.fernet import Fernet, InvalidToken
+
+from mage_ai.data_preparation.repo_manager import get_data_dir, get_repo_path
 
 DEFAULT_MAGE_SECRETS_DIR = 'secrets'
 

mage_ai/data_preparation/sync/__init__.py

@@ -17,10 +17,11 @@ class AuthType(str, Enum):
 
 @dataclass
 class GitConfig(BaseConfig):
-    remote_repo_link: str
+    remote_repo_link: str = None
     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)
 

mage_ai/orchestration/db/models/oauth.py

@@ -15,7 +15,7 @@ from sqlalchemy import (
 )
 from sqlalchemy.orm import relationship, validates
 
-from mage_ai.data_preparation.repo_manager import get_repo_path
+from mage_ai.data_preparation.repo_manager import get_project_uuid, get_repo_path
 from mage_ai.orchestration.db import db_connection, safe_db_query
 from mage_ai.orchestration.db.errors import ValidationError
 from mage_ai.orchestration.db.models.base import BaseModel
@@ -70,7 +70,7 @@ class User(BaseModel):
 
     @property
     def project_access(self) -> int:
-        return self.get_access(Permission.Entity.PROJECT, get_repo_path())
+        return self.get_access(Permission.Entity.PROJECT, get_project_uuid())
 
     def get_access(
         self,
@@ -195,11 +195,13 @@ class Role(BaseModel):
 
     def get_access(
         self,
-        entity: Union['Permission.Entity', None] = None,
+        entity: Union['Permission.Entity', None],
         entity_id: Union[str, None] = None,
     ) -> int:
         permissions = []
         if entity is None:
+            return 0
+        elif entity == Permission.Entity.ANY:
             permissions.extend(self.permissions)
         else:
             entity_permissions = list(filter(
@@ -225,7 +227,7 @@ class Role(BaseModel):
         we will go up the entity chain to see if there are permissions for parent entities.
         '''
         if entity == Permission.Entity.PIPELINE:
-            return self.get_access(Permission.Entity.PROJECT, get_repo_path())
+            return self.get_access(Permission.Entity.PROJECT, get_project_uuid())
         elif entity == Permission.Entity.PROJECT:
             return self.get_access(Permission.Entity.GLOBAL)
         else:
@@ -239,6 +241,9 @@ class UserRole(BaseModel):
 
 class Permission(BaseModel):
     class Entity(str, enum.Enum):
+        # Permissions saved to the DB should not have the "ANY" entity. It should only be used
+        # when evaluating permissions.
+        ANY = 'any'
         GLOBAL = 'global'
         PROJECT = 'project'
         PIPELINE = 'pipeline'
@@ -260,6 +265,19 @@ class Permission(BaseModel):
 
     role = relationship(Role, back_populates='permissions')
 
+    @validates('entity')
+    def validate_entity(self, key, value):
+        if value == Permission.Entity.ANY:
+            raise ValidationError(
+                'Permission entity cannot be ANY. Please select a specific entity.',
+                metadata=dict(
+                    key=key,
+                    value=value,
+                ),
+            )
+
+        return value
+
     @classmethod
     @safe_db_query
     def create_default_permissions(

mage_ai/orchestration/pipeline_scheduler.py

@@ -76,7 +76,12 @@ class PipelineScheduler:
         )
         self.logger = DictLogger(self.logger_manager.logger)
         self.notification_sender = NotificationSender(
-            NotificationConfig.load(config=self.pipeline.repo_config.notification_config),
+            NotificationConfig.load(
+                config=merge_dict(
+                    self.pipeline.repo_config.notification_config,
+                    self.pipeline.notification_config,
+                )
+            )
         )
 
         self.allow_blocks_to_fail = (
@@ -85,12 +90,13 @@ class PipelineScheduler:
         )
 
     def start(self, should_schedule: bool = True) -> None:
-        if get_preferences().sync_config:
+        preferences = get_preferences()
+        if preferences.sync_config:
             tags = dict(
                 pipeline_run_id=self.pipeline_run.id,
                 pipeline_uuid=self.pipeline.uuid,
             )
-            sync_config = GitConfig.load(config=get_preferences().sync_config)
+            sync_config = GitConfig.load(config=preferences.sync_config)
             if sync_config.sync_on_pipeline_run:
                 sync = GitSync(sync_config)
                 try:
@@ -1060,10 +1066,6 @@ def check_sla():
     pipeline_runs = PipelineRun.in_progress_runs(pipeline_schedules)
 
     if pipeline_runs:
-        notification_sender = NotificationSender(
-             NotificationConfig.load(config=get_repo_config(get_repo_path()).notification_config),
-        )
-
         current_time = datetime.now(tz=pytz.UTC)
         # TODO: combine all SLA alerts in one notification
         for pipeline_run in pipeline_runs:
@@ -1076,8 +1078,17 @@ def check_sla():
                 else pipeline_run.created_at
             if compare(start_date, current_time - timedelta(seconds=sla)) == 1:
                 # passed SLA for pipeline_run
+                pipeline = Pipeline.get(pipeline_run.pipeline_schedule.pipeline_uuid)
+                notification_sender = NotificationSender(
+                    NotificationConfig.load(
+                        config=merge_dict(
+                            pipeline.repo_config.notification_config,
+                            pipeline.notification_config,
+                        ),
+                    ),
+                )
                 notification_sender.send_pipeline_run_sla_passed_message(
-                    Pipeline.get(pipeline_run.pipeline_schedule.pipeline_uuid),
+                    pipeline,
                     pipeline_run,
                 )