gmakstutis/cognee
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

fix: revert dataset fetching in run_tasks

Browse source
This commit is contained in:
Boris Arzentar 2025-10-31 10:48:30 +01:00
parent 1864fc8f2b
commit 737b6dc634
No known key found for this signature in database
GPG key ID: D5CC274C784807B7
2 changed files with 11 additions and 4 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

7
cognee/modules/pipelines/operations/run_tasks.py
View file

@ -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

8
cognee/modules/pipelines/operations/run_tasks_distributed.py
View file

@ -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