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

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

@@ -0,0 +1,77 @@
+from typing import Any, Dict, List, Optional
+
+from mage_ai.data_preparation.logging.logger import DictLogger
+from mage_ai.data_preparation.models.pipeline import Pipeline
+from mage_ai.orchestration.db.models.schedules import BlockRun
+
+
+class DynamicBlockWrapperBase:
+    def __init__(
+        self,
+        block: Any,
+        block_run_id: Optional[int] = None,
+        execution_partition: Optional[str] = None,
+        logger: Optional[DictLogger] = None,
+        *args,
+        **kwargs,
+    ):
+        self.block = block
+        self.block_run_id = block_run_id
+        self.execution_partition = execution_partition
+        self.logger = logger
+
+        self._block_run = None
+        self._block_runs = None
+        self._pipeline_run = None
+
+    def __getattr__(self, name):
+        val = getattr(self.block, name)
+
+        def _missing(val=val, *args, **kwargs):
+            return val(*args, **kwargs)
+
+        if callable(val):
+            return _missing
+
+        return val
+
+    @property
+    def pipeline(self) -> Pipeline:
+        return self.block.pipeline
+
+    def block_run(self) -> BlockRun:
+        if self._block_run:
+            return self._block_run
+
+        self._block_run = BlockRun.query.get(self.block_run_id)
+
+        return self._block_run
+
+    def pipeline_run(self):
+        if self._pipeline_run:
+            return self._pipeline_run
+
+        self._pipeline_run = self.block_run().pipeline_run
+
+        return self._pipeline_run
+
+    def block_runs(self) -> List[BlockRun]:
+        if self._block_runs:
+            return self._block_runs
+
+        self._block_runs = BlockRun.query.filter(
+            BlockRun.pipeline_run_id == self.block_run().pipeline_run_id,
+        ).all()
+
+        return self._block_runs
+
+    def run_tests(self, **kwargs):
+        pass
+
+    def execute_sync(
+        self,
+        execution_partition: Optional[str] = None,
+        logging_tags: Optional[Dict] = None,
+        **kwargs,
+    ) -> List[Dict]:
+        raise Exception('Not implemented')

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

@@ -4,7 +4,7 @@ from typing import Dict, List
 from mage_ai.extensions.great_expectations.constants import (
     EXTENSION_UUID as EXTENSION_UUID_GREAT_EXPECTATIONS,
 )
-from mage_ai.shared.hash import index_by
+from mage_ai.shared.hash import ignore_keys, index_by
 
 
 def handle_run_tests(
@@ -47,3 +47,12 @@ def handle_run_tests(
             logging_tags=logging_tags,
             update_status=False,
         )
+
+
+def compare_extension(extension_a: Dict, extension_b: Dict) -> bool:
+    extension_a_blocks = (extension_a or dict()).get('blocks', [])
+    extension_b_blocks = (extension_b or dict()).get('blocks', [])
+
+    extension_a_blocks = [ignore_keys(b, 'outputs') for b in extension_a_blocks]
+    extension_b_blocks = [ignore_keys(b, 'outputs') for b in extension_b_blocks]
+    return extension_a_blocks == extension_b_blocks

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

@@ -2,6 +2,8 @@ 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
+from mage_ai.shared.path_fixer import add_absolute_path
 
 
 class GlobalDataProductBlock(Block):
@@ -12,7 +14,15 @@ class GlobalDataProductBlock(Block):
         )
 
         override_configuration = (self.configuration or {}).get('global_data_product', {})
-        global_data_product = GlobalDataProduct.get(override_configuration.get('uuid'))
+
+        project_path = self.repo_path
+        if override_configuration.get('project'):
+            project_path = add_absolute_path(override_configuration.get('project'))
+
+        global_data_product = GlobalDataProduct.get(
+            override_configuration.get('uuid'),
+            project_path,
+        )
 
         for key in [
             'outdated_after',
@@ -37,13 +47,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 []

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

@@ -429,7 +429,9 @@ class DestinationBlock(IntegrationBlock):
             from mage_ai.data_preparation.models.pipelines.integration_pipeline import (
                 IntegrationPipeline,
             )
-            integration_pipeline = IntegrationPipeline(self.pipeline.uuid)
+            integration_pipeline = IntegrationPipeline(
+                self.pipeline.uuid, repo_path=self.repo_path
+            )
             destination_state_file_path = integration_pipeline.destination_state_file_path(
                 destination_table=destination_table,
                 stream=state_stream,
@@ -482,7 +484,9 @@ class DestinationBlock(IntegrationBlock):
                 update_destination_state_bookmarks,
             )
 
-            integration_pipeline = IntegrationPipeline(self.pipeline.uuid)
+            integration_pipeline = IntegrationPipeline(
+                self.pipeline.uuid, repo_path=self.repo_path
+            )
             tap_stream_id = data.get('tap_stream_id')
             destination_table = data.get('destination_table')
             bookmark_values = data.get('bookmark_values', {})