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

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

@@ -84,7 +84,7 @@ class StreamingPipelineExecutor(PipelineExecutor):
         # 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')
@@ -96,7 +96,7 @@ class StreamingPipelineExecutor(PipelineExecutor):
         try:
             if retry_config is None:
                 retry_config = self.pipeline.retry_config or dict()
-            infinite_retries = False if retry_config else False
+            infinite_retries = False if retry_config else True
 
             if type(retry_config) is not RetryConfig:
                 retry_config = RetryConfig.load(config=retry_config)

mage_ai/data_preparation/git/__init__.py

@@ -19,9 +19,11 @@ from mage_ai.data_preparation.git.utils import (
     check_connection_async,
     create_ssh_keys,
     get_access_token,
+    get_oauth_access_token_for_user,
     get_provider_from_remote_url,
     poll_process_with_timeout,
     run_command,
+    validate_authentication_for_remote_url,
 )
 from mage_ai.data_preparation.preferences import get_preferences
 from mage_ai.data_preparation.sync import AuthType, GitConfig
@@ -40,22 +42,32 @@ class Git:
     def __init__(
         self,
         auth_type: AuthType = None,
+        context_data: Dict = None,
         git_config: GitConfig = None,
+        repo_path: str = None,
         setup_repo: bool = False,
+        user: User = None,
     ) -> None:
         import git
         import git.exc
 
         self.auth_type = auth_type if auth_type else AuthType.SSH
+        self.context_data = context_data if context_data is not None else dict()
         self.git_config = git_config
+        self.user = user
         self.origin = None
         self.remote_repo_link = None
         self.repo = None
 
-        self.repo_path = os.getcwd()
+        self.project_repo_path = repo_path or get_repo_path(user=user)
 
+        self.repo_path = os.getcwd()
         if project_platform_activated():
-            git_dict = git_settings()
+            git_dict = git_settings(
+                context_data=self.context_data,
+                repo_path=self.project_repo_path,
+                user=user,
+            )
             if git_dict and git_dict.get('path'):
                 self.repo_path = git_dict.get('path')
 
@@ -79,8 +91,7 @@ class Git:
             if setup_repo:
                 self.__setup_repo()
 
-        if self.repo and self.git_config:
-            self.__set_git_config()
+        self.__set_git_config()
 
         if self.remote_repo_link and self.repo:
             try:
@@ -148,12 +159,19 @@ class Git:
     @classmethod
     def get_manager(
         self,
-        user: User = None,
-        setup_repo: bool = False,
         auth_type: str = None,
         config_overwrite: Dict = None,
+        context_data: Dict = None,
+        preferences=None,
+        repo_path: str = None,
+        setup_repo: bool = False,
+        user: User = None,
     ) -> 'Git':
-        preferences = get_preferences(user=user)
+        if preferences is None:
+            preferences = get_preferences(
+                repo_path=repo_path,
+                user=user,
+            )
         git_config = None
         config = preferences.sync_config if preferences and preferences.sync_config else {}
 
@@ -162,10 +180,16 @@ class Git:
 
         git_config = GitConfig.load(config=config)
 
+        if auth_type is None:
+            auth_type = git_config.auth_type
+
         return Git(
             auth_type=auth_type,
+            context_data=context_data,
             git_config=git_config,
+            repo_path=repo_path,
             setup_repo=setup_repo,
+            user=user,
         )
 
     @property
@@ -182,8 +206,11 @@ class Git:
 
         return [branch.name for branch in self.repo.branches]
 
-    async def check_connection(self) -> None:
-        await check_connection_async(self.repo.git, self.origin.name)
+    async def check_connection(self, remote_url: str = None) -> None:
+        if remote_url:
+            await validate_authentication_for_remote_url(self.repo.git, remote_url)
+        else:
+            await check_connection_async(self.repo.git, self.origin.name)
 
     def _remote_command(func: Callable) -> None:
         """
@@ -224,19 +251,19 @@ class Git:
             elif self.auth_type == AuthType.HTTPS:
                 token = self.get_access_token()
                 url_original = list(self.origin.urls)[0]
+                remote_repo_link = self.remote_repo_link
                 if self.git_config and self.remote_repo_link and token:
-                    self.remote_repo_link = build_authenticated_remote_url(
+                    remote_repo_link = build_authenticated_remote_url(
                         self.remote_repo_link,
                         self.git_config.username,
                         token,
                     )
-                    self.origin.set_url(self.remote_repo_link)
+                    self.origin.set_url(remote_repo_link)
                 try:
-                    asyncio.run(self.check_connection())
+                    asyncio.run(self.check_connection(remote_url=remote_repo_link))
                     return func(self, *args, **kwargs)
                 finally:
                     self.origin.set_url(url_original)
-                    self.remote_repo_link = url_original
         return wrapper
 
     def add_remote(self, name: str, url: str) -> None:
@@ -432,6 +459,10 @@ class Git:
         self.repo.git.reset('--hard', f'{remote.name}/{branch}')
         self.__pip_install()
 
+    @_remote_command
+    def fetch(self) -> None:
+        self.origin.fetch()
+
     @_remote_command
     def push(self) -> None:
         self.repo.git.push(
@@ -564,7 +595,7 @@ class Git:
                 if len(remote_refs) == 0 and user:
                     from mage_ai.data_preparation.git import api
 
-                    access_token = api.get_access_token_for_user(user)
+                    access_token = get_oauth_access_token_for_user(user)
                     if access_token:
 
                         if remote_exists:
@@ -685,15 +716,22 @@ class Git:
         elif remote:
             # For remote branches, switch to the local branch if it exists. Otherwise create a new
             # branch using the remote branch as the starting point.
-            if branch.startswith(remote):
-                branch = branch[len(remote) + 1:]
-            if branch in self.repo.heads:
-                self.repo.git.switch(branch)
-            else:
-                self.repo.git.switch('-c', branch, f'{remote}/{branch}')
+            self._switch_remote_branch(branch, remote)
         else:
             self.repo.git.switch('-c', branch)
 
+    @_remote_command
+    def _switch_remote_branch(self, branch: str, remote: str) -> None:
+        self.switch_remote_branch(branch, remote)
+
+    def switch_remote_branch(self, branch: str, remote: str) -> None:
+        if branch.startswith(remote):
+            branch = branch[len(remote) + 1:]
+        if branch in self.repo.heads:
+            self.repo.git.switch(branch)
+        else:
+            self.repo.git.switch('-c', branch, f'{remote}/{branch}')
+
     @_remote_command
     def clone(self, sync_submodules: bool = False) -> None:
         from git import Repo
@@ -729,13 +767,23 @@ class Git:
             shutil.rmtree(tmp_path)
 
     def __set_git_config(self):
-        if self.git_config.username:
-            self.repo.config_writer().set_value(
-                'user', 'name', self.git_config.username).release()
-        if self.git_config.email:
-            self.repo.config_writer().set_value(
-                'user', 'email', self.git_config.email).release()
-        self.repo.config_writer('global').set_value(
+        from git.config import GitConfigParser
+        if self.git_config:
+            if self.git_config.username:
+                self.repo.config_writer().set_value(
+                    'user', 'name', self.git_config.username).release()
+            if self.git_config.email:
+                self.repo.config_writer().set_value(
+                    'user', 'email', self.git_config.email).release()
+
+        # This GitConfigParser is created in the same way that GitPython creates a
+        # global config parser. We create it this way to avoid needing to create a
+        # git.Repo object.
+        global_config = GitConfigParser(
+            os.path.normpath(os.path.expanduser("~/.gitconfig")),
+            read_only=False,
+        )
+        global_config.set_value(
             'safe', 'directory', Path(self.repo_path).as_posix()).release()
 
     def __pip_install(self) -> None:
@@ -756,10 +804,10 @@ class Git:
                 pass
 
     def __create_ssh_keys(self, overwrite: bool = False) -> str:
-        return create_ssh_keys(self.git_config, get_repo_path(), overwrite=overwrite)
+        return create_ssh_keys(self.git_config, self.project_repo_path, overwrite=overwrite)
 
     def __add_host_to_known_hosts(self):
         return add_host_to_known_hosts_util(self.git_config.remote_repo_link)
 
     def get_access_token(self) -> str:
-        return get_access_token(self.git_config)
+        return get_access_token(self.git_config, self.project_repo_path)

mage_ai/data_preparation/git/api.py

@@ -3,9 +3,6 @@ import shutil
 import uuid
 from typing import Dict
 
-from git.remote import RemoteProgress
-from git.repo.base import Repo
-
 from mage_ai.authentication.oauth.constants import ProviderName, get_ghe_hostname
 from mage_ai.authentication.oauth.utils import access_tokens_for_client
 from mage_ai.data_preparation.git.clients.base import Client as GitClient
@@ -33,13 +30,53 @@ def get_access_token_for_user(user: User, provider: str = None) -> Oauth2AccessT
         return access_tokens[0]
 
 
+def switch_branch(
+    remote_name: str,
+    remote_url: str,
+    branch_name: str,
+    token: str,
+    config_overwrite: Dict = None,
+    user: User = None,
+    **kwargs,
+):
+    from mage_ai.data_preparation.git import Git
+    provider = get_provider_from_remote_url(remote_url)
+    username = get_username(token, user=user, provider=provider)
+
+    url = build_authenticated_remote_url(remote_url, username, token)
+    git_manager = Git.get_manager(user=user, config_overwrite=config_overwrite)
+
+    remote = git_manager.repo.remotes[remote_name]
+    url_original = list(remote.urls)[0]
+    remote.set_url(url)
+
+    try:
+        git_manager.switch_remote_branch(branch_name, remote_name)
+    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)
+
+
 def fetch(
     remote_name: str,
     remote_url: str,
     token: str,
     user: User = None,
     config_overwrite: Dict = None,
-) -> RemoteProgress:
+    **kwargs,
+):
+    """
+    Returns:
+        git.remote.RemoteProgress: Custom progress object that can be used to monitor the
+            fetch progress.
+    """
+    from git.remote import RemoteProgress
+
     from mage_ai.data_preparation.git import Git
 
     custom_progress = RemoteProgress()
@@ -54,7 +91,7 @@ def fetch(
     remote.set_url(url)
 
     try:
-        remote.fetch(progress=custom_progress)
+        remote.fetch(progress=custom_progress, **kwargs)
     except Exception as err:
         raise err
     finally:
@@ -74,7 +111,15 @@ def pull(
     token: str,
     user: User = None,
     config_overwrite: Dict = None,
-) -> RemoteProgress:
+    **kwargs,
+):
+    """
+    Returns:
+        git.remote.RemoteProgress: Custom progress object that can be used to monitor the
+            pull progress.
+    """
+    from git.remote import RemoteProgress
+
     from mage_ai.data_preparation.git import Git
 
     custom_progress = RemoteProgress()
@@ -110,7 +155,13 @@ def push_raw(
     token: str,
     user: User = None,
     **kwargs,
-) -> RemoteProgress:
+):
+    """
+    Returns:
+        git.remote.RemoteProgress: Custom progress object that can be used to monitor the
+            push progress.
+    """
+    from git.remote import RemoteProgress
     custom_progress = RemoteProgress()
     provider = get_provider_from_remote_url(remote_url)
     username = get_username(token, user=user, provider=provider)
@@ -142,7 +193,13 @@ def push(
     token: str,
     user: User = None,
     config_overwrite: Dict = None,
-) -> RemoteProgress:
+    **kwargs,
+):
+    """
+    Returns:
+        git.remote.RemoteProgress: Custom progress object that can be used to monitor the
+            push progress.
+    """
     from mage_ai.data_preparation.git import Git
 
     git_manager = Git.get_manager(user=user, config_overwrite=config_overwrite)
@@ -165,6 +222,7 @@ def reset_hard(
     token: str,
     user: User = None,
     config_overwrite: Dict = None,
+    **kwargs,
 ) -> None:
     from mage_ai.data_preparation.git import Git
 
@@ -195,6 +253,7 @@ def clone(
     token: str,
     user: User = None,
     config_overwrite: Dict = None,
+    **kwargs,
 ) -> None:
     from mage_ai.data_preparation.git import Git
 
@@ -213,6 +272,8 @@ def clone(
     tmp_path = f'{git_manager.repo_path}_{uuid.uuid4().hex}'
     os.mkdir(tmp_path)
     try:
+        from git.repo.base import Repo
+
         # Clone to a tmp folder first, then copy the folder to the actual repo path. Git
         # won't allow you to clone to a folder that is not empty.
         Repo.clone_from(