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

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')

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
@@ -79,8 +81,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:
@@ -162,6 +163,9 @@ 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,
             git_config=git_config,
@@ -182,8 +186,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 +231,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 +439,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 +575,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 +696,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 +747,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:

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,51 @@ 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,
+):
+    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:
+):
+    """
+    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()
@@ -74,7 +109,14 @@ def pull(
     token: str,
     user: User = None,
     config_overwrite: Dict = None,
-) -> RemoteProgress:
+):
+    """
+    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 +152,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 +190,12 @@ def push(
     token: str,
     user: User = None,
     config_overwrite: Dict = None,
-) -> RemoteProgress:
+):
+    """
+    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)
@@ -213,6 +266,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(

mage_ai/data_preparation/git/utils.py

@@ -6,16 +6,20 @@ from typing import Callable
 from urllib.parse import urlparse, urlsplit, urlunsplit
 
 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.constants import (
     DEFAULT_KNOWN_HOSTS_FILE,
     DEFAULT_SSH_KEY_DIRECTORY,
 )
+from mage_ai.data_preparation.repo_manager import get_project_uuid
 from mage_ai.data_preparation.shared.secrets import get_secret_value
 from mage_ai.data_preparation.sync import AuthType, GitConfig
-from mage_ai.settings import get_settings_value
+from mage_ai.orchestration.db.models.oauth import Oauth2AccessToken, User
+from mage_ai.settings import get_bool_value, get_settings_value
 from mage_ai.settings.keys import (
     BITBUCKET_HOST,
     GIT_ACCESS_TOKEN,
+    GIT_OVERWRITE_WITH_PROJECT_SETTINGS,
     GIT_SSH_PRIVATE_KEY,
     GIT_SSH_PUBLIC_KEY,
     GITLAB_HOST,
@@ -52,19 +56,23 @@ def get_provider_from_remote_url(remote_url: str) -> str:
         return ProviderName.GITHUB
 
 
-def create_ssh_keys(git_config: GitConfig, repo_path: str, overwrite: bool = False) -> str:
+def create_ssh_keys(
+    git_config: GitConfig, repo_path: str, overwrite: bool = False
+) -> str:
     if not os.path.exists(DEFAULT_SSH_KEY_DIRECTORY):
         os.mkdir(DEFAULT_SSH_KEY_DIRECTORY, 0o700)
     pubk_secret_name = git_config.ssh_public_key_secret_name
+    overwrite_with_project_settings = get_bool_value(
+        get_settings_value(GIT_OVERWRITE_WITH_PROJECT_SETTINGS)
+    )
     if pubk_secret_name:
         public_key_file = os.path.join(
-            DEFAULT_SSH_KEY_DIRECTORY,
-            f'id_rsa_{pubk_secret_name}.pub'
+            DEFAULT_SSH_KEY_DIRECTORY, f'id_rsa_{pubk_secret_name}.pub'
         )
         if not os.path.exists(public_key_file) or overwrite:
             try:
                 public_key = get_settings_value(GIT_SSH_PUBLIC_KEY)
-                if not public_key:
+                if not public_key or overwrite_with_project_settings:
                     public_key = get_secret_value(
                         pubk_secret_name,
                         repo_name=repo_path,
@@ -80,13 +88,12 @@ def create_ssh_keys(git_config: GitConfig, repo_path: str, overwrite: bool = Fal
     private_key_file = os.path.join(DEFAULT_SSH_KEY_DIRECTORY, 'id_rsa')
     if pk_secret_name:
         custom_private_key_file = os.path.join(
-            DEFAULT_SSH_KEY_DIRECTORY,
-            f'id_rsa_{pk_secret_name}'
+            DEFAULT_SSH_KEY_DIRECTORY, f'id_rsa_{pk_secret_name}'
         )
         if not os.path.exists(custom_private_key_file) or overwrite:
             try:
                 private_key = get_settings_value(GIT_SSH_PRIVATE_KEY)
-                if not private_key:
+                if not private_key or overwrite_with_project_settings:
                     private_key = get_secret_value(
                         pk_secret_name,
                         repo_name=repo_path,
@@ -142,11 +149,16 @@ def add_host_to_known_hosts(remote_repo_link: str):
 
 def get_access_token(git_config, repo_path: str = None) -> str:
     token = get_settings_value(GIT_ACCESS_TOKEN)
-    if not token and git_config and git_config.access_token_secret_name:
-        token = get_secret_value(
+    overwrite_with_project_settings = get_bool_value(
+        get_settings_value(GIT_OVERWRITE_WITH_PROJECT_SETTINGS)
+    )
+    if git_config and git_config.access_token_secret_name:
+        token_from_secrets = get_secret_value(
             git_config.access_token_secret_name,
             repo_name=repo_path or get_repo_path(),
         )
+        if not token or overwrite_with_project_settings:
+            token = token_from_secrets
 
     return token
 
@@ -179,10 +191,7 @@ async def poll_process_with_timeout(
     if return_code is not None:
         out, err = proc.communicate()
         if return_code != 0:
-            message = (
-                err.decode('UTF-8') if err
-                else error_message
-            )
+            message = err.decode('UTF-8') if err else error_message
             raise ChildProcessError(message)
         else:
             return out.decode('UTF-8') if out else None
@@ -205,17 +214,16 @@ async def check_connection_async(git, remote_name: str) -> None:
             'Error connecting to remote, make sure your access token or SSH key is '
             'set up properly.'
         ),
+        timeout=5,
     )
 
 
 async def validate_authentication_for_remote_url(git, remote_url: str) -> None:
     proc = 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.',
-        )
+    await poll_process_with_timeout(
+        proc,
+        error_message='Error connecting to remote, make sure your access is valid.',
     )
 
 
@@ -258,14 +266,30 @@ def execute_on_remote_branch(func: Callable, branch) -> None:
                     if hostname:
                         await __add_host_to_known_hosts(remote_name)
                     else:
-                        raise TimeoutError("""
+                        raise TimeoutError(
+                            """
 Connecting to remote timed out, make sure your SSH key is set up properly
 and your repository host is added as a known host. More information
 here: https://docs.mage.ai/developing-in-the-cloud/setting-up-git#5-add-github-com-to-known-hosts
-""")
+"""
+                        )
             return await func(*args, **kwargs)
         else:
             await check_connection_async(git, remote_name)
             return await func(*args, **kwargs)
 
     return wrapper
+
+
+def get_oauth_client_id(provider: str) -> str:
+    return f'{provider}_{get_project_uuid()}'
+
+
+def get_oauth_access_token_for_user(
+    user: User, provider: str = None
+) -> Oauth2AccessToken:
+    if not provider:
+        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]

mage_ai/data_preparation/models/block/__init__.py

@@ -1122,6 +1122,7 @@ class Block(DataIntegrationMixin, SparkBlock, ProjectPlatformAccessible):
         data_integration_runtime_settings: Dict = None,
         execution_partition_previous: str = None,
         metadata: Dict = None,
+        override_outputs: bool = True,
         **kwargs,
     ) -> Dict:
         if logging_tags is None:
@@ -1213,7 +1214,7 @@ class Block(DataIntegrationMixin, SparkBlock, ProjectPlatformAccessible):
                         DX_PRINTER.critical(
                             block=self,
                             execution_partition=execution_partition,
-                            override_outputs=True,
+                            override_outputs=override_outputs,
                             dynamic_block_uuid=dynamic_block_uuid,
                             __uuid='store_variables',
                         )
@@ -1221,7 +1222,7 @@ class Block(DataIntegrationMixin, SparkBlock, ProjectPlatformAccessible):
                         self.store_variables(
                             variable_mapping,
                             execution_partition=execution_partition,
-                            override_outputs=True,
+                            override_outputs=override_outputs,
                             spark=self.__get_spark_session_from_global_vars(
                                 global_vars=global_vars,
                             ),
@@ -1616,7 +1617,7 @@ class Block(DataIntegrationMixin, SparkBlock, ProjectPlatformAccessible):
 
         block_function = self._validate_execution(decorated_functions, input_vars)
         if block_function is not None:
-            if logger and 'logger' not in global_vars:
+            if logger:
                 global_vars['logger'] = logger
 
             track_spark = from_notebook and self.should_track_spark()
@@ -1780,6 +1781,7 @@ class Block(DataIntegrationMixin, SparkBlock, ProjectPlatformAccessible):
             global_vars=global_vars,
             metadata=metadata,
             upstream_block_uuids_override=upstream_block_uuids_override,
+            current_block=self,
         )
 
         return variables
@@ -2386,8 +2388,17 @@ df = get_variable('{self.pipeline.uuid}', '{self.uuid}', 'df')
 
     def get_executor_type(self) -> str:
         if self.executor_type:
-            return Template(self.executor_type).render(**get_template_vars())
-        return self.executor_type
+            block_executor_type = Template(self.executor_type).render(**get_template_vars())
+        else:
+            block_executor_type = None
+        if not block_executor_type or block_executor_type == ExecutorType.LOCAL_PYTHON:
+            # If block executor_type is not set, fall back to pipeline level executor_type
+            if self.pipeline:
+                pipeline_executor_type = self.pipeline.get_executor_type()
+            else:
+                pipeline_executor_type = None
+            block_executor_type = pipeline_executor_type or block_executor_type
+        return block_executor_type
 
     def get_pipelines_from_cache(self, block_cache: BlockCache = None) -> List[Dict]:
         if block_cache is None:
@@ -2994,6 +3005,8 @@ df = get_variable('{self.pipeline.uuid}', '{self.uuid}', 'df')
         global_vars: Dict = None,
         dynamic_block_index: int = None,
     ) -> Dict:
+        from mage_ai.data_preparation.models.block.remote.models import RemoteBlock
+
         """
         Enriches the provided global variables dictionary with additional context, Spark session,
         environment, configuration, and an empty context dictionary.
@@ -3040,6 +3053,12 @@ df = get_variable('{self.pipeline.uuid}', '{self.uuid}', 'df')
         if dynamic_block_index is not None:
             global_vars['dynamic_block_index'] = dynamic_block_index
 
+        # Remote blocks
+        if global_vars.get('remote_blocks'):
+            global_vars['remote_blocks'] = [RemoteBlock.load(
+                **remote_block_dict,
+            ).get_outputs() for remote_block_dict in global_vars['remote_blocks']]
+
         self.global_vars = global_vars
 
         return global_vars

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

@@ -9,7 +9,10 @@ import pandas as pd
 from mage_ai.data_preparation.models.block.dynamic.variables import (
     get_outputs_for_dynamic_block,
 )
-from mage_ai.data_preparation.models.constants import DATAFRAME_ANALYSIS_MAX_COLUMNS
+from mage_ai.data_preparation.models.constants import (
+    DATAFRAME_ANALYSIS_MAX_COLUMNS,
+    DATAFRAME_SAMPLE_COUNT_PREVIEW,
+)
 from mage_ai.server.kernel_output_parser import DataType
 from mage_ai.shared.array import find
 from mage_ai.shared.custom_logger import DX_PRINTER
@@ -245,7 +248,9 @@ def transform_dataframe_for_display(dataframe: pd.DataFrame) -> Dict:
         data = dict(
             columns=columns_to_display,
             rows=json.loads(
-                dataframe[columns_to_display].to_json(orient='split')
+                dataframe[columns_to_display][:DATAFRAME_SAMPLE_COUNT_PREVIEW].to_json(
+                    orient='split',
+                )
             )['data'],
             index=list(dataframe.index),
             shape=[row_count, column_count],
@@ -253,7 +258,7 @@ def transform_dataframe_for_display(dataframe: pd.DataFrame) -> Dict:
     else:
         data = dict(
             columns=['col0'],
-            rows=[[dataframe]],
+            rows=[[dataframe[:DATAFRAME_SAMPLE_COUNT_PREVIEW]]],
             index=[0],
             shape=[1, 1],
         )
@@ -413,7 +418,7 @@ def transform_output_for_display_dynamic_child(
             df = pd.concat([df, df_inner], axis=1)
 
     df = limit_output(df, sample_count)
-    if 1 == len(set(df.columns)):
+    if isinstance(df, pd.DataFrame) and 1 == len(set(df.columns)):
         df.columns = [f'{col}_{idx}' for idx, col in enumerate(df.columns)]
 
     return transform_dataframe_for_display(df)

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

@@ -2,6 +2,7 @@ from logging import Logger
 from typing import Dict, List
 
 from mage_ai.data_preparation.models.block import Block
+from mage_ai.data_preparation.models.block.remote.models import RemoteBlock
 
 
 class GlobalDataProductBlock(Block):
@@ -37,13 +38,35 @@ class GlobalDataProductBlock(Block):
         # Avoid circular dependency of Pipeline class
         from mage_ai.orchestration.triggers import global_data_product as trigger
 
-        trigger.trigger_and_check_status(
-            self.get_global_data_product(),
+        global_data_product = self.get_global_data_product()
+
+        pipeline_run = trigger.trigger_and_check_status(
+            global_data_product,
             block=self,
             from_notebook=from_notebook,
             logger=logger,
             logging_tags=logging_tags,
+            poll_interval=30,
+            remote_blocks=global_vars.get('remote_blocks'),
             variables=global_vars.get('variables') if global_vars else None,
         )
 
+        if pipeline_run and pipeline_run.pipeline:
+            pipeline = pipeline_run.pipeline
+
+            arr = []
+
+            for block in global_data_product.get_blocks():
+                if block and block.uuid:
+                    arr.append(RemoteBlock.load(
+                        block_uuid=block.uuid,
+                        execution_partition=pipeline_run.execution_partition,
+                        pipeline_uuid=pipeline.uuid,
+                        repo_path=pipeline.repo_path,
+                    ))
+
+            return [
+                dict(remote_blocks=arr),
+            ]
+
         return []