topoteretes--cognee/cognee/modules/data/methods/delete_dataset.py

from cognee.modules.users.models import DatasetDatabase
from sqlalchemy import select, text
from sqlalchemy.orm.attributes import flag_modified

from cognee.modules.data.models import Dataset, DatasetData, Data
from cognee.infrastructure.databases.utils.get_vector_dataset_database_handler import (
    get_vector_dataset_database_handler,
)
from cognee.infrastructure.databases.utils.get_graph_dataset_database_handler import (
    get_graph_dataset_database_handler,
)
from cognee.infrastructure.databases.relational import get_relational_engine


async def delete_dataset(dataset: Dataset):
    db_engine = get_relational_engine()

    async with db_engine.get_async_session() as session:
        if db_engine.engine.dialect.name == "sqlite":
            # Foreign key constraints are disabled by default in SQLite (for backwards compatibility),
            # so must be enabled for each database connection/session separately.
            await session.execute(text("PRAGMA foreign_keys = ON;"))

        stmt = select(DatasetDatabase).where(
            DatasetDatabase.dataset_id == dataset.id,
        )
        dataset_database: DatasetDatabase = await session.scalar(stmt)
        if dataset_database:
            graph_dataset_database_handler = get_graph_dataset_database_handler(dataset_database)
            vector_dataset_database_handler = get_vector_dataset_database_handler(dataset_database)
            await graph_dataset_database_handler["handler_instance"].delete_dataset(
                dataset_database
            )
            await vector_dataset_database_handler["handler_instance"].delete_dataset(
                dataset_database
            )

        # Clear pipeline_status entries for this dataset from related Data objects
        # so re-adding the same data isn't blocked by stale "completed" status.
        data_ids_query = select(DatasetData.data_id).where(DatasetData.dataset_id == dataset.id)
        data_records = (
            (await session.execute(select(Data).where(Data.id.in_(data_ids_query)))).scalars().all()
        )

        dataset_id_str = str(dataset.id)
        for data_record in data_records:
            if not data_record.pipeline_status:
                continue
            updated = False
            for pipeline_name in list(data_record.pipeline_status.keys()):
                if dataset_id_str in data_record.pipeline_status[pipeline_name]:
                    del data_record.pipeline_status[pipeline_name][dataset_id_str]
                    updated = True
            if updated:
                # MutableDict only tracks top-level dict changes. Nested dict
                # modifications (del status[pipeline][dataset_id]) are invisible
                # to SQLAlchemy's change tracking, so we must explicitly mark
                # the column as dirty to ensure the update is persisted.
                flag_modified(data_record, "pipeline_status")

        await session.commit()

    # Use ORM session.delete() instead of raw table reflection with hardcoded
    # schema.  The previous ``delete_entity_by_id`` call reflected the table
    # using ``schema_name="public"`` by default, which fails when PostgreSQL
    # tables live in a non-public schema (e.g. ``cognee``).  session.delete()
    # resolves the table through SQLAlchemy metadata and triggers ORM cascades
    # (e.g. deleting related ACL rows via cascade="all, delete-orphan").
    async with db_engine.get_async_session() as session:
        if db_engine.engine.dialect.name == "sqlite":
            # Foreign key constraints are disabled by default in SQLite (for backwards compatibility),
            # so must be enabled for each database connection/session separately.
            await session.execute(text("PRAGMA foreign_keys = ON;"))

        dataset = await session.get(Dataset, dataset.id)
        if dataset:
            await session.delete(dataset)
            await session.commit()