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

mage_ai/api/resources/PipelineResource.py

@@ -62,29 +62,38 @@ from mage_ai.usage_statistics.logger import UsageStatisticLogger
 
 
 @safe_db_query
-def query_pipeline_schedules(pipeline_uuids: List[str]):
+def query_pipeline_schedules(
+    pipeline_uuids: List[str],
+    context_data: Dict = None,
+    repo_path: str = None,
+):
     a = aliased(PipelineSchedule, name='a')
     result = (
-        PipelineSchedule.
-        select(*[
-            a.created_at,
-            a.id,
-            a.name,
-            a.pipeline_uuid,
-            a.schedule_interval,
-            a.schedule_type,
-            a.start_time,
-            a.status,
-            a.updated_at,
-        ]).
-        filter(
+        PipelineSchedule.select(
+            *[
+                a.created_at,
+                a.id,
+                a.name,
+                a.pipeline_uuid,
+                a.schedule_interval,
+                a.schedule_type,
+                a.start_time,
+                a.status,
+                a.updated_at,
+            ]
+        ).filter(
             a.pipeline_uuid.in_(pipeline_uuids),
             or_(
-                a.repo_path.in_(Project().repo_path_for_database_query(
-                    'pipeline_schedules',
-                )),
+                a.repo_path.in_(
+                    Project(
+                        context_data=context_data,
+                        repo_path=repo_path,
+                    ).repo_path_for_database_query(
+                        'pipeline_schedules',
+                    )
+                ),
                 a.repo_path.is_(None),
-            )
+            ),
         )
     ).all()
     return group_by(lambda x: x.pipeline_uuid, result)
@@ -94,6 +103,7 @@ class PipelineResource(BaseResource):
     @classmethod
     @safe_db_query
     async def collection(self, query, meta, user, **kwargs):
+        context_data = kwargs.get('context_data')
         limit = (meta or {}).get(META_KEY_LIMIT, None)
         if limit is not None:
             limit = int(limit)
@@ -123,7 +133,11 @@ class PipelineResource(BaseResource):
         if repo_path:
             repo_path = repo_path[0]
         if not repo_path:
-            repo_path = get_repo_path(root_project=False)
+            repo_path = get_repo_path(
+                context_data=context_data,
+                root_project=False,
+                user=user,
+            )
 
         search_query = query.get('search', [None])
         if search_query:
@@ -154,9 +168,12 @@ class PipelineResource(BaseResource):
 
         history_by_pipeline_uuid = {}
         if from_history_days is not None and is_number(from_history_days):
-            timestamp_start = (datetime.utcnow() - timedelta(
-                hours=24 * int(from_history_days),
-            )).timestamp()
+            timestamp_start = (
+                datetime.utcnow()
+                - timedelta(
+                    hours=24 * int(from_history_days),
+                )
+            ).timestamp()
             history = load_pipelines_detail(timestamp_start=timestamp_start)
             history = sorted(
                 history,
@@ -180,21 +197,29 @@ class PipelineResource(BaseResource):
 
             if NO_TAGS_QUERY in tags:
                 pipeline_uuids_with_tags = set(cache.get_all_pipeline_uuids_with_tags())
-                all_pipeline_uuids = set(Pipeline.get_all_pipelines(repo_path))
-                pipeline_uuids_without_tags = list(all_pipeline_uuids - pipeline_uuids_with_tags)
+                all_pipeline_uuids = set(
+                    Pipeline.get_all_pipelines(repo_path=repo_path)
+                )
+                pipeline_uuids_without_tags = list(
+                    all_pipeline_uuids - pipeline_uuids_with_tags
+                )
                 pipeline_uuids = pipeline_uuids + pipeline_uuids_without_tags
         else:
-            pipeline_uuids = Pipeline.get_all_pipelines(repo_path)
+            pipeline_uuids = Pipeline.get_all_pipelines(repo_path=repo_path)
 
         if search_query:
-            pipeline_uuids = list(filter(
-                lambda x: search_query.lower() in x.lower() or
-                search_query.lower() in x.lower().split(' '),
-                pipeline_uuids,
-            ))
+            pipeline_uuids = list(
+                filter(
+                    lambda x: search_query.lower() in x.lower()
+                    or search_query.lower() in x.lower().split(' '),
+                    pipeline_uuids,
+                )
+            )
 
         total_count = len(pipeline_uuids)
-        await UsageStatisticLogger().pipelines_impression(lambda: total_count)
+        await UsageStatisticLogger(
+            context_data=context_data, repo_path=repo_path,
+        ).pipelines_impression(lambda: total_count)
 
         if not sorts:
             pipeline_uuids = sorted(pipeline_uuids, reverse=reverse_sort)
@@ -205,7 +230,7 @@ class PipelineResource(BaseResource):
             if offset:
                 pipeline_uuids = pipeline_uuids[offset:]
             if limit:
-                pipeline_uuids = pipeline_uuids[:(limit + 1)]
+                pipeline_uuids = pipeline_uuids[: (limit + 1)]
             offset_limit_applied = True
 
         cache = await PipelineCache.initialize_cache()
@@ -215,7 +240,10 @@ class PipelineResource(BaseResource):
                 return await Pipeline.get_async(uuid, repo_path=repo_path)
             except Exception as err:
                 err_message = f'Error loading pipeline {uuid}: {err}.'
-                if err.__class__.__name__ == 'OSError' and 'Too many open files' in err.strerror:
+                if (
+                    err.__class__.__name__ == 'OSError'
+                    and 'Too many open files' in err.strerror
+                ):
                     raise Exception(err_message)
                 else:
                     print(err_message)
@@ -226,7 +254,10 @@ class PipelineResource(BaseResource):
                 return Pipeline(uuid, config=config, repo_path=repo_path)
             except Exception as err:
                 err_message = f'Error loading pipeline sync {uuid}: {err}.'
-                if err.__class__.__name__ == 'OSError' and 'Too many open files' in err.strerror:
+                if (
+                    err.__class__.__name__ == 'OSError'
+                    and 'Too many open files' in err.strerror
+                ):
                     raise Exception(err_message)
                 else:
                     print(err_message)
@@ -239,10 +270,13 @@ class PipelineResource(BaseResource):
             for pipeline_dict in cache.get_models(types=pipeline_types):
                 pipeline_uuid_from_cache = pipeline_dict['pipeline']['uuid']
                 if pipeline_uuid_from_cache in pipeline_uuids:
-                    pipelines.append(Pipeline(
-                        pipeline_uuid_from_cache,
-                        config=pipeline_dict['pipeline'],
-                    ))
+                    pipelines.append(
+                        Pipeline(
+                            pipeline_uuid_from_cache,
+                            config=pipeline_dict['pipeline'],
+                            repo_path=repo_path,
+                        )
+                    )
         else:
             for uuid in pipeline_uuids:
                 pipeline_dict = cache.get_model(dict(uuid=uuid, repo_path=repo_path))
@@ -253,10 +287,13 @@ class PipelineResource(BaseResource):
                     else:
                         # Add pipeline with type "invalid" so pipeline with invalid config
                         # can still be displayed in UI and visible to user
-                        pipelines.append(Pipeline(
-                            uuid,
-                            config=dict(type='invalid'),
-                        ))
+                        pipelines.append(
+                            Pipeline(
+                                uuid,
+                                config=dict(type='invalid'),
+                                repo_path=repo_path,
+                            )
+                        )
                 else:
                     pipeline_uuids_miss.append(uuid)
 
@@ -269,7 +306,11 @@ class PipelineResource(BaseResource):
 
         mapping = {}
         if include_schedules:
-            mapping = query_pipeline_schedules(pipeline_uuids)
+            mapping = query_pipeline_schedules(
+                pipeline_uuids,
+                context_data=context_data,
+                repo_path=repo_path,
+            )
 
         filtered_pipelines = []
         for pipeline in pipelines:
@@ -279,15 +320,23 @@ class PipelineResource(BaseResource):
             pipeline.schedules = schedules
 
             if pipeline_statuses and (
-                (PipelineStatus.ACTIVE in pipeline_statuses and
-                    any(s.status == ScheduleStatus.ACTIVE
-                        for s in pipeline.schedules)) or
-                (PipelineStatus.INACTIVE in pipeline_statuses and
-                    len(pipeline.schedules) > 0 and
-                    all(s.status == ScheduleStatus.INACTIVE
-                        for s in pipeline.schedules)) or
-                (PipelineStatus.NO_SCHEDULES in pipeline_statuses and
-                    len(pipeline.schedules) == 0)
+                (
+                    PipelineStatus.ACTIVE in pipeline_statuses
+                    and any(
+                        s.status == ScheduleStatus.ACTIVE for s in pipeline.schedules
+                    )
+                )
+                or (
+                    PipelineStatus.INACTIVE in pipeline_statuses
+                    and len(pipeline.schedules) > 0
+                    and all(
+                        s.status == ScheduleStatus.INACTIVE for s in pipeline.schedules
+                    )
+                )
+                or (
+                    PipelineStatus.NO_SCHEDULES in pipeline_statuses
+                    and len(pipeline.schedules) == 0
+                )
             ):
                 filtered_pipelines.append(pipeline)
 
@@ -300,6 +349,7 @@ class PipelineResource(BaseResource):
                     pipeline.history = history_by_pipeline_uuid.get(pipeline.uuid)
 
         if sorts:
+
             def _sort_key(p, sorts=sorts, reverse_sort=reverse_sort):
                 bools = []
                 vals = []
@@ -307,24 +357,35 @@ class PipelineResource(BaseResource):
                     if 'blocks' == k.lower():
                         val = len(p.blocks_by_uuid)
                         vals.append(val)
-                        bools.append(val is None if not reverse_sort else val is not None)
+                        bools.append(
+                            val is None if not reverse_sort else val is not None
+                        )
                     elif 'triggers' == k.lower():
                         val = len(p.schedules)
                         vals.append(val)
-                        bools.append(val is None if not reverse_sort else val is not None)
+                        bools.append(
+                            val is None if not reverse_sort else val is not None
+                        )
                     elif 'status' == k.lower():
                         if len(p.schedules) == 0:
                             val = 'no_schedules'
-                        elif find(lambda s: s.status == 'active', p.schedules) is not None:
+                        elif (
+                            find(lambda s: s.status == 'active', p.schedules)
+                            is not None
+                        ):
                             val = 'active'
                         else:
                             val = 'inactive'
                         vals.append(val)
-                        bools.append(val is None if not reverse_sort else val is not None)
+                        bools.append(
+                            val is None if not reverse_sort else val is not None
+                        )
                     elif hasattr(p, k):
                         val = getattr(p, k)
                         vals.append(val)
-                        bools.append(val is None if not reverse_sort else val is not None)
+                        bools.append(
+                            val is None if not reverse_sort else val is not None
+                        )
                     else:
                         bools.append(False)
 
@@ -359,12 +420,15 @@ class PipelineResource(BaseResource):
             'results': len(arr),
             'next': has_next,
         }
-
+        for p in arr:
+            p.context_data = context_data
         return result_set
 
     @classmethod
     @safe_db_query
     async def create(self, payload, user, **kwargs):
+        context_data = kwargs.get('context_data')
+
         clone_pipeline_uuid = payload.get('clone_pipeline_uuid')
         template_uuid = payload.get('custom_template_uuid')
         name = payload.get('name')
@@ -374,18 +438,22 @@ class PipelineResource(BaseResource):
         llm_payload = payload.get('llm')
         pipeline = None
 
+        repo_path = get_repo_path(context_data=context_data, user=user)
         if template_uuid:
-            custom_template = CustomPipelineTemplate.load(template_uuid=template_uuid)
+            custom_template = CustomPipelineTemplate.load(
+                repo_path,
+                template_uuid=template_uuid,
+            )
             pipeline = custom_template.create_pipeline(name)
         elif clone_pipeline_uuid is not None:
-            source = Pipeline.get(clone_pipeline_uuid)
+            source = Pipeline.get(clone_pipeline_uuid, repo_path=repo_path)
             pipeline = await Pipeline.duplicate(source, name)
         else:
             pipeline = Pipeline.create(
                 name,
                 description=description,
                 pipeline_type=pipeline_type,
-                repo_path=get_repo_path(),
+                repo_path=repo_path,
                 tags=tags,
             )
 
@@ -409,25 +477,37 @@ class PipelineResource(BaseResource):
                             block_payload['configuration'] = configuration
 
                         block_uuid = f'{pipeline.uuid}_block_{block_number}'
-                        block_resource = await BlockResource.create(merge_dict(
-                            dict(
-                                name=block_uuid,
-                                type=block_payload.get('block_type'),
+                        block_resource = await BlockResource.create(
+                            merge_dict(
+                                dict(
+                                    name=block_uuid,
+                                    type=block_payload.get('block_type'),
+                                ),
+                                ignore_keys(
+                                    block_payload,
+                                    [
+                                        'block_type',
+                                        'upstream_blocks',
+                                    ],
+                                ),
+                            ),
+                            user,
+                            **merge_dict(
+                                kwargs,
+                                dict(
+                                    parent_model=pipeline,
+                                ),
                             ),
-                            ignore_keys(block_payload, [
-                                'block_type',
-                                'upstream_blocks',
-                            ]),
-                        ), user, **merge_dict(kwargs, dict(
-                            parent_model=pipeline,
-                        )))
+                        )
 
                         upstream_block_uuids = block_payload.get('upstream_blocks')
 
                         pipeline.add_block(
                             block_resource.model,
                             None,
-                            priority=len(upstream_block_uuids) if upstream_block_uuids else 0,
+                            priority=len(upstream_block_uuids)
+                            if upstream_block_uuids
+                            else 0,
                             widget=False,
                         )
 
@@ -441,11 +521,17 @@ class PipelineResource(BaseResource):
 
                         if upstream_block_uuids and len(upstream_block_uuids) >= 1:
                             block = config['block']
-                            arr = [f'{pipeline.uuid}_block_{bn}' for bn in upstream_block_uuids]
+                            arr = [
+                                f'{pipeline.uuid}_block_{bn}'
+                                for bn in upstream_block_uuids
+                            ]
                             block.update(dict(upstream_blocks=arr))
 
         if pipeline:
-            await UsageStatisticLogger().pipeline_create(
+            await UsageStatisticLogger(
+                context_data=context_data,
+                repo_path=repo_path,
+            ).pipeline_create(
                 pipeline,
                 clone_pipeline_uuid=clone_pipeline_uuid,
                 llm_payload=llm_payload,
@@ -472,19 +558,26 @@ class PipelineResource(BaseResource):
 
         self.on_create_callback = _on_create_callback
 
+        pipeline.context_data = context_data
         return self(pipeline, user, **kwargs)
 
     @classmethod
     @safe_db_query
-    async def __fetch_model(self, pipeline_uuid: str, **kqwargs):
+    async def __fetch_model(self, pipeline_uuid: str, repo_path: str, **kwargs):
         all_projects = project_platform_activated()
 
         if all_projects:
             return await get_pipeline_from_platform_async(
                 pipeline_uuid,
+                repo_path=repo_path,
+                context_data=kwargs.get('context_data'),
             )
 
-        return await Pipeline.get_async(pipeline_uuid, all_projects=all_projects)
+        return await Pipeline.get_async(
+            pipeline_uuid,
+            all_projects=all_projects,
+            repo_path=repo_path,
+        )
 
     @classmethod
     @safe_db_query
@@ -494,23 +587,33 @@ class PipelineResource(BaseResource):
         **kwargs,
     ):
         pipeline_uuid = urllib.parse.unquote(pk)
-        return await self.__fetch_model(pipeline_uuid, **kwargs)
+        user = kwargs.get('user')
+        repo_path = get_repo_path(context_data=kwargs.get('context_data'), user=user)
+        return await self.__fetch_model(pipeline_uuid, repo_path, **kwargs)
 
     @classmethod
     @safe_db_query
     async def member(self, pk, user, **kwargs):
-        pipeline = await self.__fetch_model(pk, **kwargs)
+        context_data = kwargs.get('context_data')
+
+        repo_path = get_repo_path(context_data=context_data, user=user)
+        pipeline = await self.__fetch_model(pk, repo_path, **kwargs)
 
         api_operation_action = kwargs.get('api_operation_action', None)
         if api_operation_action != DELETE:
             kernel_name = PIPELINE_TO_KERNEL_NAME[pipeline.type]
             switch_active_kernel(
                 kernel_name,
-                emr_config=pipeline.executor_config if kernel_name == KernelName.PYSPARK else None,
+                emr_config=pipeline.executor_config
+                if kernel_name == KernelName.PYSPARK
+                else None,
             )
 
         if api_operation_action == DETAIL:
-            if Project(pipeline.repo_config).is_feature_enabled(
+            if Project(
+                context_data=context_data,
+                repo_config=pipeline.repo_config,
+            ).is_feature_enabled(
                 FeatureUUID.OPERATION_HISTORY,
             ):
                 record_detail_pipeline(
@@ -537,6 +640,7 @@ class PipelineResource(BaseResource):
             if mapping.get(pipeline.uuid):
                 pipeline.schedules = mapping[pipeline.uuid] or []
 
+        pipeline.context_data = context_data
         return self(pipeline, user, **kwargs)
 
     @safe_db_query
@@ -563,26 +667,30 @@ class PipelineResource(BaseResource):
 
                 block_cache = await BlockCache.initialize_cache()
                 for block in pipeline.blocks_by_uuid.values():
-                    block_cache.remove_pipeline(block, pipeline.uuid, pipeline.repo_path)
+                    block_cache.remove_pipeline(
+                        block, pipeline.uuid, pipeline.repo_path
+                    )
 
         self.on_delete_callback = _on_delete_callback
         return self.model.delete()
 
     @safe_db_query
     async def update(self, payload, **kwargs):
+        context_data = kwargs.get('context_data')
         if 'add_upstream_for_block_uuid' in payload:
             block_uuid = payload['add_upstream_for_block_uuid']
             block = self.model.get_block(block_uuid, widget=False)
             if BlockType.DBT == block.type and block.language == BlockLanguage.SQL:
                 upstream_dbt_blocks_by_uuid = {
-                    block.uuid: block
-                    for block in block.upstream_dbt_blocks()
+                    block.uuid: block for block in block.upstream_dbt_blocks()
                 }
                 self.model.blocks_by_uuid.update(upstream_dbt_blocks_by_uuid)
                 self.model.validate('A cycle was formed while adding a block')
                 self.model.save()
             return self
 
+        repo_path = get_repo_path(context_data=context_data, user=self.current_user)
+
         query = kwargs.get('query', {})
         update_content = query.get('update_content', [False])
         if update_content:
@@ -596,23 +704,34 @@ class PipelineResource(BaseResource):
             if 'pipeline_uuid' not in llm_payload:
                 llm_payload['pipeline_uuid'] = self.model.uuid
 
-            llm_resource = await LlmResource.create(llm_payload, self.current_user, **kwargs)
+            llm_resource = await LlmResource.create(
+                llm_payload, self.current_user, **kwargs
+            )
             llm_response = llm_resource.model.get('response')
 
             pipeline_doc = None
             block_docs = []
             blocks = self.model.blocks_by_uuid.values()
 
-            async def _add_markdown_block(block_doc: str, block_uuid: str, priority: int):
-                return await BlockResource.create(dict(
-                    content=block_doc.strip() if block_doc else block_doc,
-                    language=BlockLanguage.MARKDOWN,
-                    name=f'Documentation for {block_uuid}',
-                    priority=priority,
-                    type=BlockType.MARKDOWN,
-                ), self.current_user, **merge_dict(kwargs, dict(
-                    parent_model=self.model,
-                )))
+            async def _add_markdown_block(
+                block_doc: str, block_uuid: str, priority: int
+            ):
+                return await BlockResource.create(
+                    dict(
+                        content=block_doc.strip() if block_doc else block_doc,
+                        language=BlockLanguage.MARKDOWN,
+                        name=f'Documentation for {block_uuid}',
+                        priority=priority,
+                        type=BlockType.MARKDOWN,
+                    ),
+                    self.current_user,
+                    **merge_dict(
+                        kwargs,
+                        dict(
+                            parent_model=self.model,
+                        ),
+                    ),
+                )
 
             if LLMUseCase.GENERATE_DOC_FOR_BLOCK == llm_use_case:
                 block_doc = llm_response.get('block_doc')
@@ -656,11 +775,15 @@ class PipelineResource(BaseResource):
             switch_active_kernel(
                 kernel_name,
                 emr_config=self.model.executor_config
-                if kernel_name == KernelName.PYSPARK else None,
+                if kernel_name == KernelName.PYSPARK
+                else None,
             )
         except Exception as e:
             pipeline_type_updated = payload.get('type')
-            if pipeline_type_updated is not None and pipeline_type_updated != pipeline_type:
+            if (
+                pipeline_type_updated is not None
+                and pipeline_type_updated != pipeline_type
+            ):
                 await self.model.update(dict(type=pipeline_type))
 
             raise e
@@ -670,9 +793,9 @@ class PipelineResource(BaseResource):
             trigger_configs_by_name = get_trigger_configs_by_name(pipeline_uuid)
             triggers_in_code_to_update = []
             schedules = (
-                PipelineSchedule.
-                query.
-                filter(PipelineSchedule.pipeline_uuid == pipeline_uuid)
+                PipelineSchedule.query.filter(
+                    PipelineSchedule.pipeline_uuid == pipeline_uuid
+                )
             ).all()
             for schedule in schedules:
                 trigger_config = trigger_configs_by_name.get(schedule.name)
@@ -695,29 +818,30 @@ class PipelineResource(BaseResource):
         ):
             if pipeline_runs is not None:
                 pipeline_run_ids = [run.get('id') for run in pipeline_runs]
-                pipeline_runs_to_cancel = (
-                    PipelineRun.
-                    query.
-                    filter(PipelineRun.id.in_(pipeline_run_ids))
+                pipeline_runs_to_cancel = PipelineRun.query.filter(
+                    PipelineRun.id.in_(pipeline_run_ids)
                 )
             elif pipeline_schedule_id is not None:
-                pipeline_runs_to_cancel = PipelineRun.in_progress_runs([pipeline_schedule_id])
+                pipeline_runs_to_cancel = PipelineRun.in_progress_runs(
+                    [pipeline_schedule_id]
+                )
             else:
-                pipeline_runs_to_cancel = (
-                    PipelineRun.
-                    query.
-                    filter(PipelineRun.pipeline_uuid == pipeline_uuid).
-                    filter(PipelineRun.status.in_([
-                        PipelineRun.PipelineRunStatus.INITIAL,
-                        PipelineRun.PipelineRunStatus.RUNNING,
-                    ]))
+                pipeline_runs_to_cancel = PipelineRun.query.filter(
+                    PipelineRun.pipeline_uuid == pipeline_uuid
+                ).filter(
+                    PipelineRun.status.in_(
+                        [
+                            PipelineRun.PipelineRunStatus.INITIAL,
+                            PipelineRun.PipelineRunStatus.RUNNING,
+                        ]
+                    )
                 )
             for pipeline_run in pipeline_runs_to_cancel:
                 PipelineScheduler(pipeline_run).stop()
 
         def retry_pipeline_runs(pipeline_runs):
             for run in pipeline_runs:
-                retry_pipeline_run(run)
+                retry_pipeline_run(run, repo_path)
 
         @safe_db_query
         def query_incomplete_block_runs(pipeline_uuid: str):
@@ -731,15 +855,18 @@ class PipelineResource(BaseResource):
                 a.pipeline_uuid,
             ]
             result = (
-                PipelineRun.
-                select(*columns).
-                join(b, a.id == b.pipeline_run_id).
-                filter(a.pipeline_uuid == pipeline_uuid).
-                filter(a.status == PipelineRun.PipelineRunStatus.FAILED).
-                filter(b.status.not_in([
-                    BlockRun.BlockRunStatus.COMPLETED,
-                    BlockRun.BlockRunStatus.CONDITION_FAILED,
-                ]))
+                PipelineRun.select(*columns)
+                .join(b, a.id == b.pipeline_run_id)
+                .filter(a.pipeline_uuid == pipeline_uuid)
+                .filter(a.status == PipelineRun.PipelineRunStatus.FAILED)
+                .filter(
+                    b.status.not_in(
+                        [
+                            BlockRun.BlockRunStatus.COMPLETED,
+                            BlockRun.BlockRunStatus.CONDITION_FAILED,
+                        ]
+                    )
+                )
             ).all()
 
             return result