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

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

@@ -907,10 +907,11 @@ class Block(DataIntegrationMixin, SparkBlock, ProjectPlatformAccessible):
         parts = get_path_parts(file_path)
 
         if parts and len(parts) >= 3:
+            from mage_ai.data_preparation.models.block.block_factory import BlockFactory
+
             # If file_path == transformers/test4.py
             # parts ==
             # ('/home/src/default_repo/default_platform2/project3', 'transformers', 'test4.py')
-
             # If project platform platform activated, then parts ==
             # ('/home/src', 'default_repo', 'data_loaders/astral_violet.py')
 
@@ -929,7 +930,7 @@ class Block(DataIntegrationMixin, SparkBlock, ProjectPlatformAccessible):
             configuration = dict(file_path=file_path, file_source=dict(path=file_path))
             language = FILE_EXTENSION_TO_BLOCK_LANGUAGE.get(extension)
 
-            return self.get_block(
+            return BlockFactory.get_block(
                 block_uuid,
                 block_uuid,
                 block_type,
@@ -1121,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:
@@ -1212,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',
                         )
@@ -1220,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,
                             ),
@@ -1615,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()
@@ -1779,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
@@ -2385,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:
@@ -2993,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.
@@ -3039,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
@@ -3473,12 +3493,15 @@ df = get_variable('{self.pipeline.uuid}', '{self.uuid}', 'df')
 
             cache = BlockCache()
             if detach:
+                from mage_ai.data_preparation.models.block.block_factory import (
+                    BlockFactory,
+                )
                 """"
                 New block added to pipeline, so it must be added to the block cache.
                 Old block no longer in pipeline, so it must be removed from block cache.
                 """
                 cache.add_pipeline(self, self.pipeline)
-                old_block = self.get_block(
+                old_block = BlockFactory.get_block(
                     old_uuid,
                     old_uuid,
                     self.type,

mage_ai/data_preparation/models/block/data_integration/mixins.py

@@ -202,18 +202,29 @@ class DataIntegrationMixin:
             with open(catalog_full_path, mode='w') as f:
                 f.write(json.dumps(catalog))
 
-    def is_data_integration(self) -> bool:
+    def is_data_integration(self, pipeline_project: Project = None) -> bool:
         """
         Check if the block is a data integration block.
         If the data_integration_in_batch_pipeline feature is not enabled, return False.
 
+        Args:
+            pipeline_project (Project, optional): A cached Project value to avoid
+                looking it up many times when called inside loops. Defaults to None.
+
         Returns:
             bool: True if it's a data integration block, False otherwise.
         """
-        if not self.pipeline or not \
-                Project(self.pipeline.repo_config).is_feature_enabled(
-                    FeatureUUID.DATA_INTEGRATION_IN_BATCH_PIPELINE,
-                ):
+        if not self.pipeline:
+
+            return False
+
+        actual_project: Project = pipeline_project
+        if not actual_project:
+            actual_project = Project(self.pipeline.repo_config)
+
+        if not actual_project.is_feature_enabled(
+                        FeatureUUID.DATA_INTEGRATION_IN_BATCH_PIPELINE,
+                    ):
 
             return False
 

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

@@ -229,3 +229,6 @@ class DynamicChildController:
             block_runs.append(block_run)
 
         return block_runs
+
+    def run_tests(self, **kwargs):
+        pass

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/dynamic/variables.py

@@ -6,7 +6,7 @@ from typing import Any, Dict, List, Tuple, Union
 
 import pandas as pd
 
-from mage_ai.data_preparation.models.constants import BlockLanguage
+from mage_ai.data_preparation.models.constants import BlockLanguage, BlockType
 from mage_ai.data_preparation.models.variable import Variable
 from mage_ai.shared.memory import get_memory_usage, get_memory_usage_async
 from mage_ai.shared.strings import to_ordinal_integers
@@ -470,7 +470,7 @@ def fetch_input_variables_for_dynamic_upstream_blocks(
 
             # If dynamic child should reduce its output (which means it passes the entire
             # output to its downstream blocks):
-            if should_reduce_output(upstream_block):
+            if should_reduce_output(upstream_block) and block.type != BlockType.EXTENSION:
                 child_data = []
                 metadata = {}
                 for lazy_variable_set in lazy_variable_controller:

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

@@ -45,4 +45,5 @@ def handle_run_tests(
             global_vars=global_vars,
             logger=logger,
             logging_tags=logging_tags,
+            update_status=False,
         )