From 737b6dc63422d62d544d1b9478c34397f9b19f58 Mon Sep 17 00:00:00 2001 From: Boris Arzentar Date: Fri, 31 Oct 2025 10:48:30 +0100 Subject: [PATCH] fix: revert dataset fetching in run_tasks --- cognee/modules/pipelines/operations/run_tasks.py | 7 +++++-- .../modules/pipelines/operations/run_tasks_distributed.py | 8 ++++++-- 2 files changed, 11 insertions(+), 4 deletions(-) diff --git a/cognee/modules/pipelines/operations/run_tasks.py b/cognee/modules/pipelines/operations/run_tasks.py index 4414b73bb..c9ae1da1a 100644 --- a/cognee/modules/pipelines/operations/run_tasks.py +++ b/cognee/modules/pipelines/operations/run_tasks.py @@ -7,7 +7,6 @@ from uuid import UUID from cognee.infrastructure.databases.graph import get_graph_engine from cognee.infrastructure.databases.relational import get_relational_engine -from cognee.modules.data.methods import get_authorized_dataset from cognee.modules.pipelines.operations.run_tasks_distributed import run_tasks_distributed from cognee.modules.users.models import User from cognee.shared.logging_utils import get_logger @@ -65,7 +64,11 @@ async def run_tasks( if not user: user = await get_default_user() - dataset = await get_authorized_dataset(user, dataset_id, "write") + async with get_relational_engine().get_async_session() as session: + from cognee.modules.data.models import Dataset + + dataset = await session.get(Dataset, dataset_id) + pipeline_id = generate_pipeline_id(user.id, dataset.id, pipeline_name) pipeline_run = await log_pipeline_run_start(pipeline_id, pipeline_name, dataset.id, data) pipeline_run_id = pipeline_run.pipeline_run_id diff --git a/cognee/modules/pipelines/operations/run_tasks_distributed.py b/cognee/modules/pipelines/operations/run_tasks_distributed.py index 997ffd571..b997e92a4 100644 --- a/cognee/modules/pipelines/operations/run_tasks_distributed.py +++ b/cognee/modules/pipelines/operations/run_tasks_distributed.py @@ -6,7 +6,7 @@ except ModuleNotFoundError: from typing import Any, List, Optional from uuid import UUID -from cognee.modules.data.methods import get_authorized_dataset +from cognee.infrastructure.databases.relational import get_relational_engine from cognee.modules.data.models import Dataset from cognee.modules.pipelines.tasks.task import Task from cognee.modules.pipelines.models import ( @@ -93,7 +93,11 @@ async def run_tasks_distributed( if not user: user = await get_default_user() - dataset = await get_authorized_dataset(user, dataset_id, "write") + async with get_relational_engine().get_async_session() as session: + from cognee.modules.data.models import Dataset + + dataset = await session.get(Dataset, dataset_id) + pipeline_id: UUID = generate_pipeline_id(user.id, dataset.id, pipeline_name) pipeline_run = await log_pipeline_run_start(pipeline_id, pipeline_name, dataset.id, data) pipeline_run_id: UUID = pipeline_run.pipeline_run_id