mage-ai 0.9.57__py3-none-any.whl → 0.9.59__py3-none-any.whl

mage_ai/data_preparation/executors/ecs_block_executor.py

@@ -8,7 +8,7 @@ from mage_ai.shared.hash import merge_dict
 class EcsBlockExecutor(BlockExecutor):
     RETRYABLE = False
 
-    def __init__(self, pipeline, block_uuid: str, execution_partition: str = None):
+    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.ecs_config or dict()
         if self.block.executor_config is not None:
@@ -25,4 +25,4 @@ class EcsBlockExecutor(BlockExecutor):
             global_vars=global_vars,
             **kwargs,
         )
-        ecs.run_task(' '.join(cmd), ecs_config=self.executor_config)
+        ecs.run_task(cmd, ecs_config=self.executor_config)

mage_ai/data_preparation/executors/executor_factory.py

@@ -96,6 +96,7 @@ class ExecutorFactory:
         self,
         pipeline: Pipeline,
         block_uuid: str,
+        block_run_id: int = None,
         execution_partition: Union[str, None] = None,
         executor_type: Union[ExecutorType, str, None] = None,
     ) -> BlockExecutor:
@@ -119,6 +120,7 @@ class ExecutorFactory:
         """
         executor_kwargs = dict(
             pipeline=pipeline,
+            block_run_id=block_run_id,
             block_uuid=block_uuid,
             execution_partition=execution_partition,
         )

mage_ai/data_preparation/executors/gcp_cloud_run_block_executor.py

@@ -9,7 +9,7 @@ from mage_ai.shared.hash import merge_dict
 class GcpCloudRunBlockExecutor(BlockExecutor):
     RETRYABLE = False
 
-    def __init__(self, pipeline, block_uuid: str, execution_partition: str = None):
+    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.gcp_cloud_run_config or dict()
         if os.getenv('GCP_REGION'):

mage_ai/data_preparation/executors/k8s_block_executor.py

@@ -2,15 +2,15 @@ from typing import Dict
 
 from mage_ai.data_preparation.executors.block_executor import BlockExecutor
 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.services.k8s.constants import DEFAULT_NAMESPACE
 
 
 class K8sBlockExecutor(BlockExecutor):
     RETRYABLE = False
 
-    def __init__(self, pipeline, block_uuid: str, execution_partition: str = None):
+    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()
         if self.block.executor_config is not None:

mage_ai/data_preparation/executors/pyspark_block_executor.py

@@ -14,7 +14,13 @@ from mage_ai.shared.hash import merge_dict
 
 
 class PySparkBlockExecutor(BlockExecutor):
-    def __init__(self, pipeline: Pipeline, block_uuid: str, execution_partition: str = None):
+    def __init__(
+        self,
+        pipeline: Pipeline,
+        block_uuid: str,
+        execution_partition: str = None,
+        **kwargs,
+    ):
         super().__init__(pipeline, block_uuid, execution_partition=execution_partition)
         self.resource_manager = EmrResourceManager(
             pipeline.repo_config.s3_bucket,

mage_ai/data_preparation/git/__init__.py

@@ -45,7 +45,6 @@ class Git:
 
         self.repo_path = os.getcwd()
 
-        print(project_platform_activated())
         if project_platform_activated():
             git_dict = git_settings()
             if git_dict and git_dict.get('path'):
@@ -101,11 +100,17 @@ class Git:
         user: User = None,
         setup_repo: bool = True,
         auth_type: str = None,
+        config_overwrite: Dict = None,
     ) -> 'Git':
         preferences = get_preferences(user=user)
         git_config = None
-        if preferences and preferences.sync_config:
-            git_config = GitConfig.load(config=preferences.sync_config)
+        config = preferences.sync_config if preferences and preferences.sync_config else {}
+
+        if config_overwrite:
+            config.update(**config_overwrite)
+
+        git_config = GitConfig.load(config=config)
+
         return Git(
             auth_type=auth_type,
             git_config=git_config,
@@ -280,15 +285,18 @@ class Git:
             repo = self.repo
             repo_path = self.repo_path
 
-        parser = GitConfigParser(
-            os.path.join(repo_path, '.gitmodules'),
-            read_only=True,
-        )
+        gitmodules_path = os.path.join(repo_path, '.gitmodules')
+
+        def update_gitmodules(section: str, option: str, value: Any):
+            gitmodules_parser = GitConfigParser(gitmodules_path, read_only=False)
+            gitmodules_parser.set(section, option, value)
+            gitmodules_parser.release()
+
+        parser = GitConfigParser(gitmodules_path)
         sections = parser.sections()
         for section in sections:
             path = parser.get(section, 'path', fallback=None)
             submodule_url = parser.get(section, 'url', fallback=None)
-            parser.release()
             if path and submodule_url:
                 submodule_full_path = os.path.join(repo_path, path)
                 tmp_full_path = f'{submodule_full_path}-{str(uuid.uuid4())}'
@@ -311,8 +319,7 @@ class Git:
                         url = url._replace(netloc=f'{user}:{token}@{url.netloc}')
                         url = urlunsplit(url)
                         # Overwrite the submodule URL with git credentials.
-                        repo.config_writer().set_value(
-                            f'submodule.{path}', 'url', url).release()
+                        update_gitmodules(section, 'url', url)
 
                     subprocess.run(
                         [
@@ -335,12 +342,16 @@ class Git:
                 else:
                     print(f'{section} updated!')
                 finally:
+                    # Restore the submodule URL.
+                    update_gitmodules(section, 'url', submodule_url)
                     if os.path.exists(tmp_full_path):
                         shutil.rmtree(tmp_full_path)
                     repo_config_writer = repo.config_writer()
-                    repo_config_writer.remove_section(f'submodule.{path}')
+                    repo_config_writer.remove_section(section)
                     repo_config_writer.release()
 
+        parser.release()
+
     @_remote_command
     def reset_hard(self, branch: str = None, remote_name: str = None) -> None:
         if branch is None:
@@ -451,6 +462,13 @@ class Git:
         if message:
             self.repo.index.commit(message)
 
+    def update_config(self, settings: Dict) -> None:
+        for key, value in settings.items():
+            self.repo.config_writer().set_value(
+                *key.split('.'),
+                value,
+            ).release()
+
     def remotes(self, limit: int = 40, user: User = None) -> List[Dict]:
         arr = []
 

mage_ai/data_preparation/git/api.py

@@ -9,11 +9,7 @@ from urllib.parse import urlsplit, urlunsplit
 from git.remote import RemoteProgress
 from git.repo.base import Repo
 
-from mage_ai.authentication.oauth.constants import (
-    OAUTH_PROVIDER_GHE,
-    OAUTH_PROVIDER_GITHUB,
-    get_ghe_hostname,
-)
+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
 from mage_ai.data_preparation.git.utils import get_provider_from_remote_url
@@ -29,7 +25,7 @@ def get_oauth_client_id(provider: str) -> str:
 
 def get_access_token_for_user(user: User, provider: str = None) -> Oauth2AccessToken:
     if not provider:
-        provider = OAUTH_PROVIDER_GHE if get_ghe_hostname() else OAUTH_PROVIDER_GITHUB
+        provider = ProviderName.GHE if get_ghe_hostname() else ProviderName.GITHUB
     access_tokens = access_tokens_for_client(get_oauth_client_id(provider), user=user)
     if access_tokens:
         return access_tokens[0]
@@ -40,6 +36,7 @@ def fetch(
     remote_url: str,
     token: str,
     user: User = None,
+    config_overwrite: Dict = None,
 ) -> RemoteProgress:
     from mage_ai.data_preparation.git import Git
 
@@ -48,7 +45,7 @@ def fetch(
     username = get_username(token, user=user, provider=provider)
 
     url = build_authenticated_remote_url(remote_url, username, token)
-    git_manager = Git.get_manager(user=user)
+    git_manager = Git.get_manager(user=user, config_overwrite=config_overwrite)
 
     remote = git_manager.repo.remotes[remote_name]
     url_original = list(remote.urls)[0]
@@ -74,6 +71,7 @@ def pull(
     branch_name: str,
     token: str,
     user: User = None,
+    config_overwrite: Dict = None,
 ) -> RemoteProgress:
     from mage_ai.data_preparation.git import Git
 
@@ -82,7 +80,7 @@ def pull(
     username = get_username(token, user=user, provider=provider)
 
     url = build_authenticated_remote_url(remote_url, username, token)
-    git_manager = Git.get_manager(user=user)
+    git_manager = Git.get_manager(user=user, config_overwrite=config_overwrite)
 
     remote = git_manager.repo.remotes[remote_name]
     url_original = list(remote.urls)[0]
@@ -108,6 +106,7 @@ def push(
     branch_name: str,
     token: str,
     user: User = None,
+    config_overwrite: Dict = None,
 ) -> RemoteProgress:
     from mage_ai.data_preparation.git import Git
 
@@ -116,7 +115,7 @@ def push(
     username = get_username(token, user=user, provider=provider)
 
     url = build_authenticated_remote_url(remote_url, username, token)
-    git_manager = Git.get_manager(user=user)
+    git_manager = Git.get_manager(user=user, config_overwrite=config_overwrite)
 
     remote = git_manager.repo.remotes[remote_name]
     url_original = list(remote.urls)[0]
@@ -142,6 +141,7 @@ def reset_hard(
     branch_name: str,
     token: str,
     user: User = None,
+    config_overwrite: Dict = None,
 ) -> None:
     from mage_ai.data_preparation.git import Git
 
@@ -149,7 +149,7 @@ def reset_hard(
     username = get_username(token, user=user, provider=provider)
 
     url = build_authenticated_remote_url(remote_url, username, token)
-    git_manager = Git.get_manager(user=user)
+    git_manager = Git.get_manager(user=user, config_overwrite=config_overwrite)
 
     remote = git_manager.repo.remotes[remote_name]
     url_original = list(remote.urls)[0]
@@ -171,6 +171,7 @@ def clone(
     remote_url: str,
     token: str,
     user: User = None,
+    config_overwrite: Dict = None,
 ) -> None:
     from mage_ai.data_preparation.git import Git
 
@@ -178,7 +179,7 @@ def clone(
     username = get_username(token, user=user, provider=provider)
 
     url = build_authenticated_remote_url(remote_url, username, token)
-    git_manager = Git.get_manager(user=user)
+    git_manager = Git.get_manager(user=user, config_overwrite=config_overwrite)
 
     remote = git_manager.repo.remotes[remote_name]
     url_original = list(remote.urls)[0]
@@ -230,7 +231,7 @@ def build_authenticated_remote_url(remote_url: str, username: str, token: str) -
     return urlunsplit(url)
 
 
-def get_username(token: str, user: User = None, provider: str = OAUTH_PROVIDER_GITHUB) -> str:
+def get_username(token: str, user: User = None, provider: str = ProviderName.GITHUB) -> str:
     resp = get_user(token, provider=provider)
     if resp.get('username') is None:
         from mage_ai.data_preparation.git import Git
@@ -240,7 +241,7 @@ def get_username(token: str, user: User = None, provider: str = OAUTH_PROVIDER_G
         return resp['username']
 
 
-def get_user(token: str, provider: str = OAUTH_PROVIDER_GITHUB) -> Dict:
+def get_user(token: str, provider: str = ProviderName.GITHUB) -> Dict:
     """
     https://docs.github.com/en/rest/users/users?apiVersion=2022-11-28#get-the-authenticated-user
     """

mage_ai/data_preparation/git/clients/base.py

@@ -1,10 +1,6 @@
 from abc import ABC, abstractmethod
 
-from mage_ai.authentication.oauth.constants import (
-    OAUTH_PROVIDER_BITBUCKET,
-    OAUTH_PROVIDER_GHE,
-    OAUTH_PROVIDER_GITHUB,
-)
+from mage_ai.authentication.oauth.constants import ProviderName
 
 
 class Client(ABC):
@@ -13,15 +9,15 @@ class Client(ABC):
 
     @classmethod
     def get_client_for_provider(cls, provider: str):
-        if provider == OAUTH_PROVIDER_GITHUB:
+        if not provider or provider == ProviderName.GITHUB:
             from mage_ai.data_preparation.git.clients.github import GitHubClient
 
             return GitHubClient
-        elif provider == OAUTH_PROVIDER_BITBUCKET:
+        elif provider == ProviderName.BITBUCKET:
             from mage_ai.data_preparation.git.clients.bitbucket import BitbucketClient
 
             return BitbucketClient
-        elif provider == OAUTH_PROVIDER_GHE:
+        elif provider == ProviderName.GHE:
             from mage_ai.data_preparation.git.clients.ghe import GHEClient
 
             return GHEClient

mage_ai/data_preparation/git/clients/github.py

@@ -27,6 +27,7 @@ class GitHubClient(Client):
 
         return dict(
             username=data.get('login'),
+            email=data.get('email'),
         )
 
     def get_branches(self, repository_name: str):

mage_ai/data_preparation/git/utils.py

@@ -1,8 +1,6 @@
 from mage_ai.authentication.oauth.constants import (
     BITBUCKET_HOST,
-    OAUTH_PROVIDER_BITBUCKET,
-    OAUTH_PROVIDER_GHE,
-    OAUTH_PROVIDER_GITHUB,
+    ProviderName,
     get_ghe_hostname,
 )
 
@@ -11,11 +9,11 @@ def get_provider_from_remote_url(remote_url: str) -> str:
     ghe_hostname = get_ghe_hostname()
 
     if not remote_url:
-        return OAUTH_PROVIDER_GITHUB
+        return ProviderName.GITHUB
 
     if BITBUCKET_HOST and BITBUCKET_HOST in remote_url or 'bitbucket.org' in remote_url:
-        return OAUTH_PROVIDER_BITBUCKET
+        return ProviderName.BITBUCKET
     elif ghe_hostname and ghe_hostname in remote_url:
-        return OAUTH_PROVIDER_GHE
+        return ProviderName.GHE
     else:
-        return OAUTH_PROVIDER_GITHUB
+        return ProviderName.GITHUB