mage-ai 0.9.46__py3-none-any.whl → 0.9.47__py3-none-any.whl

mage_ai/api/operations/base.py

@@ -61,6 +61,7 @@ class BaseOperation():
         self.oauth_token = kwargs.get('oauth_token')
         self.options = kwargs.get('options', {})
         self.payload = kwargs.get('payload') or {}
+        self.payload_mutated = None
         self._query = kwargs.get('query', {}) or {}
         self.resource = kwargs.get('resource')
         self.resource_parent = kwargs.get('resource_parent')
@@ -96,6 +97,7 @@ class BaseOperation():
                 condition=HookCondition.SUCCESS,
                 metadata=metadata,
                 operation_resource=result,
+                payload=self.payload_mutated,
                 **{
                     resource_key: presented_init,
                 },
@@ -221,6 +223,7 @@ class BaseOperation():
                 error=self.__present_error(err),
                 metadata=metadata,
                 operation_resource=result,
+                payload=self.payload_mutated,
                 **{
                     resource_key: presented_init,
                 },
@@ -269,8 +272,14 @@ class BaseOperation():
         operation_type: HookOperation,
         stage: HookStage,
         condition: HookCondition = None,
+        error: Dict = None,
+        meta: Dict = None,
+        metadata: Dict = None,
         operation_resource: Union[BaseResource, Dict, List[BaseResource]] = None,
-        **kwargs,
+        payload: Dict = None,
+        query: Dict = None,
+        resource: Dict = None,
+        resources: List[Dict] = None,
     ) -> List[Hook]:
         project = Project()
         if not project.is_feature_enabled(FeatureUUID.GLOBAL_HOOKS):
@@ -283,12 +292,21 @@ class BaseOperation():
         try:
             global_hooks = GlobalHooks.load_from_file()
             hooks = global_hooks.get_and_run_hooks(
+                operation_types,
+                EntityName(self.__classified_class()),
+                stage,
                 conditions=[condition] if condition else None,
+                error=error,
+                meta=meta,
+                metadata=metadata,
                 operation_resource=operation_resource,
-                operation_types=operation_types,
-                resource_type=EntityName(self.__classified_class()),
-                stage=stage,
-                **kwargs,
+                payload=payload,
+                query=query,
+                resource=resource,
+                resource_id=self.pk,
+                resource_parent_id=self.resource_parent_id,
+                resources=resources,
+                user=dict(id=self.user.id) if self.user else None,
             )
         except Exception:
             hooks = []
@@ -309,6 +327,7 @@ class BaseOperation():
         error: Dict = None,
         metadata: Dict = None,
         operation_resource: Union[BaseResource, List[BaseResource]] = None,
+        payload: Dict = None,
         resource: Dict = None,
         resources: List[Dict] = None,
     ) -> Union[Dict, List[Dict]]:
@@ -319,6 +338,7 @@ class BaseOperation():
             metadata=metadata,
             operation_resource=operation_resource,
             operation_type=HookOperation(self.action),
+            payload=payload,
             query=self.query,
             resource=resource,
             resources=resources,
@@ -399,6 +419,8 @@ class BaseOperation():
                 if isinstance(value, dict):
                     self.query = merge_dict(self.query, value)
 
+        self.payload_mutated = payload
+
         return payload
 
     async def __executed_result(self) -> Union[BaseResource, Dict, List[BaseResource]]:

mage_ai/api/policies/GlobalHookPolicy.py

@@ -8,7 +8,7 @@ WRITEABLE_ATTRIBUTES = [
     'operation_type',
     'outputs',
     'pipeline',
-    'predicates',
+    'predicate',
     'resource_type',
     'run_settings',
     'stages',

mage_ai/api/policies/IntegrationSourcePolicy.py

@@ -14,6 +14,19 @@ IntegrationSourcePolicy.allow_actions([
     OauthScope.CLIENT_PRIVATE,
 ], condition=lambda policy: policy.has_at_least_viewer_role())
 
+
+IntegrationSourcePolicy.allow_actions(
+    [
+        constants.DETAIL,
+    ],
+    scopes=[
+        OauthScope.CLIENT_PRIVATE,
+    ],
+    condition=lambda policy: policy.has_at_least_viewer_role(),
+    override_permission_condition=lambda _policy: True,
+)
+
+
 IntegrationSourcePolicy.allow_actions([
     constants.CREATE,
     constants.UPDATE,
@@ -21,6 +34,7 @@ IntegrationSourcePolicy.allow_actions([
     OauthScope.CLIENT_PRIVATE,
 ], condition=lambda policy: policy.has_at_least_editor_role())
 
+
 IntegrationSourcePolicy.allow_read(IntegrationSourcePresenter.default_attributes + [
     'docs'
 ], scopes=[
@@ -29,6 +43,20 @@ IntegrationSourcePolicy.allow_read(IntegrationSourcePresenter.default_attributes
     constants.LIST,
 ], condition=lambda policy: policy.has_at_least_viewer_role())
 
+
+IntegrationSourcePolicy.allow_read(
+    IntegrationSourcePresenter.default_attributes,
+    scopes=[
+        OauthScope.CLIENT_PRIVATE,
+    ],
+    on_action=[
+        constants.DETAIL,
+    ],
+    condition=lambda policy: policy.has_at_least_viewer_role(),
+    override_permission_condition=lambda _policy: True,
+)
+
+
 IntegrationSourcePolicy.allow_read([
     'error_message',
     'streams',
@@ -39,6 +67,7 @@ IntegrationSourcePolicy.allow_read([
     constants.CREATE,
 ], condition=lambda policy: policy.has_at_least_editor_role())
 
+
 IntegrationSourcePolicy.allow_read([
     'selected_streams',
     'streams',
@@ -49,6 +78,7 @@ IntegrationSourcePolicy.allow_read([
     constants.UPDATE,
 ], condition=lambda policy: policy.has_at_least_editor_role())
 
+
 IntegrationSourcePolicy.allow_write([
     'action_type',
     'block_uuid',
@@ -61,6 +91,7 @@ IntegrationSourcePolicy.allow_write([
     constants.CREATE,
 ], condition=lambda policy: policy.has_at_least_editor_role())
 
+
 IntegrationSourcePolicy.allow_write([
     'block_uuid',
     'streams',
@@ -69,3 +100,34 @@ IntegrationSourcePolicy.allow_write([
 ], on_action=[
     constants.UPDATE,
 ], condition=lambda policy: policy.has_at_least_editor_role())
+
+
+IntegrationSourcePolicy.allow_query(
+    [
+        'block_uuid[]',
+    ],
+    scopes=[
+        OauthScope.CLIENT_PRIVATE,
+    ],
+    on_action=[
+        constants.LIST,
+    ],
+    condition=lambda policy: policy.has_at_least_viewer_role(),
+    override_permission_condition=lambda _policy: True,
+)
+
+
+IntegrationSourcePolicy.allow_query(
+    [
+        'pipeline_schedule_id',
+        'stream',
+    ],
+    scopes=[
+        OauthScope.CLIENT_PRIVATE,
+    ],
+    on_action=[
+        constants.DETAIL,
+    ],
+    condition=lambda policy: policy.has_at_least_viewer_role(),
+    override_permission_condition=lambda _policy: True,
+)

mage_ai/api/policies/OauthPolicy.py

@@ -46,12 +46,12 @@ class OauthPolicy(BasePolicy, UserPermissionMixIn):
             ),
             {
                 OauthScope.CLIENT_PUBLIC: {
-                    OperationType.DETAIL: config[OauthScope.CLIENT_PUBLIC][
-                        OperationType.DETAIL
-                    ],
-                    OperationType.LIST: config[OauthScope.CLIENT_PUBLIC][
-                        OperationType.LIST
-                    ],
+                    OperationType.DETAIL: (config.get(
+                        OauthScope.CLIENT_PUBLIC,
+                    ) or {}).get(OperationType.DETAIL) or {},
+                    OperationType.LIST: (config.get(
+                        OauthScope.CLIENT_PUBLIC,
+                    ) or {}).get(OperationType.LIST) or {},
                 },
             },
         )

mage_ai/api/presenters/GlobalHookPresenter.py

@@ -13,7 +13,7 @@ class GlobalHookPresenter(BasePresenter):
         'operation_type',
         'outputs',
         'pipeline',
-        'predicates',
+        'predicate',
         'resource_type',
         'run_settings',
         'stages',

mage_ai/api/presenters/IntegrationSourcePresenter.py

@@ -9,13 +9,20 @@ from mage_ai.data_preparation.models.block.data_integration.utils import (
 
 class IntegrationSourcePresenter(BasePresenter):
     default_attributes = [
+        'block',
         'name',
+        'partition',
+        'pipeline_run',
+        'pipeline_schedule',
+        'streams',
         'templates',
         'uuid',
     ]
 
-    def present(self, **kwargs):
-        if constants.UPDATE == kwargs['format']:
+    async def prepare_present(self, **kwargs):
+        display_format = kwargs['format']
+
+        if constants.UPDATE == display_format:
             payload = (kwargs.get('payload') or {}).get('integration_source') or {}
 
             catalog = None

mage_ai/api/presenters/PipelinePresenter.py

@@ -20,6 +20,7 @@ class PipelinePresenter(BasePresenter):
         'executor_type',
         'name',
         'notification_config',
+        'remote_variables_dir',
         'retry_config',
         'run_pipeline_in_one_process',
         'settings',
@@ -29,6 +30,7 @@ class PipelinePresenter(BasePresenter):
         'updated_at',
         'uuid',
         'variables',
+        'variables_dir',
         'widgets',
     ]
 

mage_ai/api/presenters/PipelineRunPresenter.py

@@ -21,7 +21,7 @@ class PipelineRunPresenter(BasePresenter):
         'variables',
     ]
 
-    async def present(self, **kwargs):
+    async def prepare_present(self, **kwargs):
         display_format = kwargs.get('format')
         data_to_display = self.model
 
@@ -92,6 +92,8 @@ PipelineRunPresenter.register_format(
         'pipeline_schedule_name',
         'pipeline_schedule_token',
         'pipeline_schedule_type',
+        'pipeline_tags',
+        'pipeline_type',
     ],
 )
 
@@ -102,8 +104,11 @@ PipelineRunPresenter.register_format(
     ],
 )
 
-PipelineRunPresenter.register_format(
-    'with_basic_details',
+PipelineRunPresenter.register_formats(
+    [
+        f'integration_source/{constants.DETAIL}',
+        'with_basic_details',
+    ],
     [
         'execution_date',
         'id',

mage_ai/api/presenters/PipelineSchedulePresenter.py

@@ -22,7 +22,7 @@ class PipelineSchedulePresenter(BasePresenter):
         'variables',
     ]
 
-    async def present(self, **kwargs):
+    async def prepare_present(self, **kwargs):
         display_format = kwargs['format']
         data = self.model.to_dict()
         next_execution_date = self.model.next_execution_date()
@@ -59,8 +59,10 @@ PipelineSchedulePresenter.register_format(
     PipelineSchedulePresenter.default_attributes + [
         'event_matchers',
         'last_pipeline_run_status',
+        'next_pipeline_run_date',
         'pipeline_in_progress_runs_count',
         'pipeline_runs_count',
+        'tags',
     ],
 )
 
@@ -69,6 +71,8 @@ PipelineSchedulePresenter.register_formats([
     constants.UPDATE,
 ], PipelineSchedulePresenter.default_attributes + [
         'event_matchers',
+        'next_pipeline_run_date',
+        'tags',
     ],
 )
 
@@ -76,6 +80,23 @@ PipelineSchedulePresenter.register_formats([
 PipelineSchedulePresenter.register_formats([
     'with_runtime_average',
 ], PipelineSchedulePresenter.default_attributes + [
+        'next_pipeline_run_date',
         'runtime_average',
+        'tags',
+    ],
+)
+
+
+PipelineSchedulePresenter.register_formats(
+    [
+        f'integration_source/{constants.DETAIL}',
+    ],
+    [
+        'id',
+        'name',
+        'schedule_interval',
+        'schedule_type',
+        'settings',
+        'variables',
     ],
 )

mage_ai/api/resources/BlockResource.py

@@ -267,6 +267,11 @@ class BlockResource(GenericResource):
                                     ':'.join(parts_new),
                                 )
 
+                    if metrics.get('dynamic_upstream_block_uuids'):
+                        block_mapping[block_run_block_uuid]['uuids'].extend(
+                            metrics.get('dynamic_upstream_block_uuids') or [],
+                        )
+
                     # If it should reduce, then all the children have 1 downstream.
                     if should_reduce_output(block):
                         for db in block.downstream_blocks:

mage_ai/api/resources/DataProviderResource.py

@@ -8,6 +8,7 @@ from mage_ai.settings.repo import get_repo_path
 
 DATA_PROVIDERS = [
     DataSource.BIGQUERY,
+    DataSource.CHROMA,
     DataSource.CLICKHOUSE,
     DataSource.DRUID,
     DataSource.DUCKDB,
@@ -21,6 +22,7 @@ DATA_PROVIDERS = [
 ]
 DATA_PROVIDERS_NAME = {
     DataSource.BIGQUERY: 'BigQuery',
+    DataSource.CHROMA: 'Chroma',
     DataSource.CLICKHOUSE: 'ClickHouse',
     DataSource.DUCKDB: 'DuckDB',
     DataSource.DRUID: 'Druid',

mage_ai/api/resources/IntegrationSourceResource.py

@@ -1,22 +1,131 @@
 import traceback
+import urllib.parse
+from typing import Dict, List
 
 from mage_ai.api.resources.GenericResource import GenericResource
+from mage_ai.api.resources.PipelineRunResource import PipelineRunResource
+from mage_ai.api.resources.PipelineScheduleResource import PipelineScheduleResource
 from mage_ai.data_integrations.sources.constants import SOURCES
-from mage_ai.data_preparation.models.block.data_integration.utils import test_connection
+from mage_ai.data_preparation.models.block.data_integration.constants import (
+    KEY_REPLICATION_METHOD,
+    REPLICATION_METHOD_INCREMENTAL,
+)
+from mage_ai.data_preparation.models.block.data_integration.utils import (
+    get_state_data,
+    test_connection,
+)
 from mage_ai.data_preparation.models.constants import BlockType
 from mage_ai.data_preparation.models.pipeline import Pipeline
 from mage_ai.data_preparation.models.pipelines.integration_pipeline import (
     IntegrationPipeline,
 )
 from mage_ai.orchestration.db import safe_db_query
+from mage_ai.orchestration.db.models.schedules import PipelineRun
 from mage_ai.server.api.integration_sources import get_collection
 
 
+def get_state_data_for_blocks(
+    pipeline: Pipeline,
+    block_uuids: List[str],
+    pipeline_schedule_id: int = None,
+    stream: str = None,
+) -> List[Dict]:
+    arr = []
+
+    for block_uuid in block_uuids:
+        block = pipeline.get_block(block_uuid)
+
+        data_integration_uuid = None
+        execution_partition_previous = None
+        pipeline_run = None
+        pipeline_schedule = None
+        state_data_by_stream = {}
+
+        pipeline_runs = PipelineRun.recently_completed_pipeline_runs(
+            pipeline.uuid,
+            pipeline_schedule_id=pipeline_schedule_id,
+            sample_size=1,
+        )
+        if pipeline_runs:
+            pipeline_run = pipeline_runs[0]
+            execution_partition_previous = pipeline_run.execution_partition
+            pipeline_schedule = pipeline_run.pipeline_schedule
+
+        if execution_partition_previous:
+            settings = block.get_data_integration_settings(
+                partition=execution_partition_previous,
+            )
+            catalog = settings.get('catalog')
+
+            if catalog and catalog.get('streams'):
+                for stream_dict in (catalog.get('streams') or []):
+                    stream_id = stream_dict.get('tap_stream_id') or stream_dict.get('stream')
+                    if stream and stream != stream_id:
+                        continue
+
+                    if REPLICATION_METHOD_INCREMENTAL != stream_dict.get(
+                        KEY_REPLICATION_METHOD,
+                    ):
+                        continue
+
+                    data_integration_uuid = settings.get('data_integration_uuid')
+                    state_data, record = get_state_data(
+                        block,
+                        catalog,
+                        data_integration_uuid=data_integration_uuid,
+                        include_record=True,
+                        partition=execution_partition_previous,
+                        stream_id=stream_id,
+                    )
+                    state_data_by_stream[stream_id] = dict(
+                        record=record,
+                        state=state_data,
+                    )
+
+        arr.append(dict(
+            block=block.to_dict(),
+            partition=execution_partition_previous,
+            pipeline_run=pipeline_run,
+            pipeline_schedule=pipeline_schedule,
+            streams=state_data_by_stream,
+            uuid=data_integration_uuid,
+        ))
+
+    return arr
+
+
 class IntegrationSourceResource(GenericResource):
     @classmethod
     @safe_db_query
     async def collection(self, query, meta, user, **kwargs):
-        collection = get_collection('sources', SOURCES)
+        collection = []
+
+        parent_model = kwargs.get('parent_model')
+        if parent_model and isinstance(parent_model, Pipeline):
+            block_uuids = query.get('block_uuid[]', [])
+            if block_uuids:
+                block_uuids = block_uuids[0]
+            if block_uuids:
+                block_uuids = block_uuids.split(',')
+
+            for state_data in get_state_data_for_blocks(parent_model, block_uuids):
+                if state_data.get('pipeline_run'):
+                    state_data['pipeline_run'] = PipelineRunResource(
+                        state_data['pipeline_run'],
+                        user,
+                        **kwargs,
+                    )
+
+                if state_data.get('pipeline_schedule'):
+                    state_data['pipeline_schedule'] = PipelineScheduleResource(
+                        state_data['pipeline_schedule'],
+                        user,
+                        **kwargs,
+                    )
+
+                collection.append(state_data)
+        else:
+            collection = get_collection('sources', SOURCES)
 
         return self.build_result_set(
             collection,
@@ -75,4 +184,42 @@ class IntegrationSourceResource(GenericResource):
     @classmethod
     @safe_db_query
     def member(self, pk, user, **kwargs):
+        parent_model = kwargs.get('parent_model')
+        query = kwargs.get('query', {})
+
+        pipeline_schedule_id = query.get('pipeline_schedule_id', [None])
+        if pipeline_schedule_id:
+            pipeline_schedule_id = pipeline_schedule_id[0]
+
+        stream = query.get('stream', [None])
+        if stream:
+            stream = stream[0]
+
+        if parent_model and isinstance(parent_model, Pipeline):
+            block_uuid = urllib.parse.unquote(pk)
+            state_data_arr = get_state_data_for_blocks(
+                parent_model,
+                [block_uuid],
+                pipeline_schedule_id=pipeline_schedule_id,
+                stream=stream,
+            )
+
+            state_data = state_data_arr[0]
+
+            if state_data.get('pipeline_run'):
+                state_data['pipeline_run'] = PipelineRunResource(
+                    state_data['pipeline_run'],
+                    user,
+                    **kwargs,
+                )
+
+            if state_data.get('pipeline_schedule'):
+                state_data['pipeline_schedule'] = PipelineScheduleResource(
+                    state_data['pipeline_schedule'],
+                    user,
+                    **kwargs,
+                )
+
+            return self(state_data, user, **kwargs)
+
         return self(IntegrationPipeline.get(pk), user, **kwargs)

mage_ai/data_integrations/sources/constants.py

@@ -1,3 +1,6 @@
+from mage_ai.data_integrations.utils.settings import get_uuid
+from mage_ai.shared.hash import index_by
+
 SQL_SOURCES = [
     dict(name='BigQuery'),
     dict(
@@ -11,6 +14,8 @@ SQL_SOURCES = [
     dict(name='Snowflake'),
 ]
 
+SQL_SOURCES_MAPPING = index_by(get_uuid, SQL_SOURCES)
+
 SOURCES = sorted([
     dict(name='Amazon S3'),
     dict(name='Amplitude'),