Repository: topoteretes/cognee
File: cognee/modules/pipelines/operations/pipeline.py (Python)

Code before:

    user, authorized_datasets = await resolve_authorized_user_datasets(datasets, user)

    for dataset in authorized_datasets:
        async for run_info in run_pipeline_per_dataset(
            dataset=dataset,
            user=user,
            tasks=tasks,
            data=data,
            pipeline_name=pipeline_name,
            context={"dataset": dataset},
            pipeline_cache=pipeline_cache,
            incremental_loading=incremental_loading,
            data_per_batch=data_per_batch,
        ):
            yield run_info


async def run_pipeline_per_dataset(
    dataset: Dataset,
    user: User,
    tasks: list[Task],
    data=None,
    pipeline_name: str = "custom_pipeline",
    context: dict = None,
    pipeline_cache=False,
    incremental_loading=False,
    data_per_batch: int = 20,
):
    # Will only be used if ENABLE_BACKEND_ACCESS_CONTROL is set to True
    await set_database_global_context_variables(dataset.id, dataset.owner_id)

    if not data:
        data: list[Data] = await get_dataset_data(dataset_id=dataset.id)

    process_pipeline_status = await check_pipeline_run_qualification(dataset, data, pipeline_name)
    if process_pipeline_status:
        # If pipeline was already processed or is currently being processed
        # return status information to async generator and finish execution
        if pipeline_cache:
            # If pipeline caching is enabled we do not proceed with re-processing
            yield process_pipeline_status
            return
        else:
            # If pipeline caching is disabled we always return pipeline started information and proceed with re-processing
            yield PipelineRunStarted(
                pipeline_run_id=process_pipeline_status.pipeline_run_id,
                dataset_id=dataset.id,
                dataset_name=dataset.name,
                payload=data,
            )

Review comment (performance):
`pipeline_cache` is a bit confusing, can we use the name that is more descriptive? Also, something like `use_pipeline_cache` associates more to a boolean. Since this is high-level API, we will have to live with this for quite some time.

Code after:

    user, authorized_datasets = await resolve_authorized_user_datasets(datasets, user)

    for dataset in authorized_datasets:
        async for run_info in run_pipeline_per_dataset(
            dataset=dataset,
            user=user,
            tasks=tasks,
            data=data,
            pipeline_name=pipeline_name,
            context={"dataset": dataset},
            use_pipeline_cache=use_pipeline_cache,
            incremental_loading=incremental_loading,
            data_per_batch=data_per_batch,
        ):
            yield run_info


async def run_pipeline_per_dataset(
    dataset: Dataset,
    user: User,
    tasks: list[Task],
    data=None,
    pipeline_name: str = "custom_pipeline",
    context: dict = None,
    use_pipeline_cache=False,
    incremental_loading=False,
    data_per_batch: int = 20,
):
    # Will only be used if ENABLE_BACKEND_ACCESS_CONTROL is set to True
    await set_database_global_context_variables(dataset.id, dataset.owner_id)

    if not data:
        data: list[Data] = await get_dataset_data(dataset_id=dataset.id)

    process_pipeline_status = await check_pipeline_run_qualification(dataset, data, pipeline_name)
    if process_pipeline_status:
        # If pipeline was already processed or is currently being processed
        # return status information to async generator and finish execution
        if use_pipeline_cache:
            # If pipeline caching is enabled we do not proceed with re-processing
            yield process_pipeline_status
            return
        else:
            # If pipeline caching is disabled we always return pipeline started information and proceed with re-processing
            yield PipelineRunStarted(
                pipeline_run_id=process_pipeline_status.pipeline_run_id,
                dataset_id=dataset.id,
                dataset_name=dataset.name,
                payload=data,
            )