topoteretes--cognee/cognee/modules/migration/export.py

"""Export a Cognee dataset's knowledge graph to portable formats.

Formats:

- ``pydantic`` — a :class:`GraphSnapshot` of typed DataPoint instances (the
  Pydantic-native export: real ``Entity``/``DocumentChunk``/custom-model
  objects, losslessly serializable via ``model_dump_json``)
- ``cogx``    — a COGX archive directory (the canonical portable dump;
  re-importable via :class:`COGXArchiveSource`)
- ``json``    — full-fidelity nodes/edges JSON
- ``graphml`` — Gephi/yEd/NetworkX interop
- ``cypher`` — MERGE script loadable into any Neo4j-compatible database
"""

from dataclasses import dataclass
from pathlib import Path
from typing import Any, Dict, Optional, Union
from uuid import UUID

from cognee.modules.migration.cogx import (
    COGXArchiveWriter,
    COGXDocument,
    COGXEntity,
    COGXFact,
    parse_timestamp,
    write_social_layer,
)
from cognee.modules.migration.formats import write_cypher, write_graphml, write_json
from cognee.modules.migration.snapshot import GraphSnapshot, build_snapshot
from cognee.shared.logging_utils import get_logger

logger = get_logger("migration.export")

EXPORT_FORMATS = ("pydantic", "cogx", "json", "graphml", "cypher")

_FORMAT_SUFFIX = {"json": ".json", "graphml": ".graphml", "cypher": ".cypher"}


@dataclass
class ExportResult:
    format: str
    destination: str
    dataset_name: str
    dataset_id: str
    num_nodes: int
    num_edges: int

    def __repr__(self):
        return (
            f"ExportResult(format={self.format!r}, destination={self.destination!r}, "
            f"dataset={self.dataset_name!r}, nodes={self.num_nodes}, edges={self.num_edges})"
        )


def _user_payload(user_row) -> Dict[str, Any]:
    return {
        "email": user_row.email,
        "hashed_password": user_row.hashed_password,
        "is_active": user_row.is_active,
        "is_superuser": user_row.is_superuser,
        "is_verified": user_row.is_verified,
    }


async def _gather_social_layer(dataset_obj) -> Dict[str, Any]:
    """The dataset's owner and ACL grants, including user credentials.

    Emails, password hashes, and account flags are included so a
    cognee-to-cognee restore can recreate functional accounts — which makes
    an archive carrying this payload a SECRET. Only gathered on explicit
    request (``include_permissions=True``).
    """
    from sqlalchemy import select

    from cognee.infrastructure.databases.relational import get_relational_engine
    from cognee.modules.users.models import ACL, Permission, User

    db_engine = get_relational_engine()
    async with db_engine.get_async_session() as session:
        owner = await session.get(User, dataset_obj.owner_id)
        rows = (
            await session.execute(
                select(User, Permission.name)
                .join(ACL, ACL.principal_id == User.id)
                .join(Permission, Permission.id == ACL.permission_id)
                .where(ACL.dataset_id == dataset_obj.id)
            )
        ).all()

    grants: Dict[str, Dict[str, Any]] = {}
    for user_row, permission_name in rows:
        grant = grants.setdefault(
            user_row.email, {"user": _user_payload(user_row), "permissions": []}
        )
        grant["permissions"].append(permission_name)
    return {
        "owner": _user_payload(owner) if owner else None,
        "grants": sorted(grants.values(), key=lambda g: g["user"]["email"]),
    }


async def _stored_migration_revision(dataset_id) -> Optional[str]:
    """The source store's stamped data-migration revision, for the manifest.

    Per-dataset bookkeeping with backend access control on, the global row
    with it off. None (no row / never migrated / bookkeeping unreadable)
    means "unknown" — the importer then keeps its own stamp.
    """
    from cognee.context_global_variables import backend_access_control_enabled
    from cognee.infrastructure.databases.relational import get_relational_engine

    try:
        db_engine = get_relational_engine()
        if backend_access_control_enabled():
            from cognee.modules.users.models import DatasetDatabase

            async with db_engine.get_async_session() as session:
                record = await session.get(DatasetDatabase, dataset_id)
            return record.migration_revision if record else None

        from cognee.modules.migrations.models import (
            GLOBAL_DATABASE_VERSION_ROW_ID,
            GlobalDatabaseVersion,
        )

        async with db_engine.get_async_session() as session:
            record = await session.get(GlobalDatabaseVersion, GLOBAL_DATABASE_VERSION_ROW_ID)
        return record.global_migration_revision if record else None
    except Exception as error:  # noqa: BLE001 — manifest metadata is best effort
        logger.debug("Could not determine migration revision for manifest: %s", error)
        return None


def _write_cogx(
    nodes,
    edges,
    destination: Path,
    dataset_name: str,
    embedding_model: Optional[str] = None,
    migration_revision: Optional[str] = None,
) -> None:
    """Map graph nodes/edges onto typed COGX records; keep the rest raw."""
    with COGXArchiveWriter(destination, source_system="cognee") as writer:
        writer.add_note(f"Exported from Cognee dataset {dataset_name!r}.")
        writer.embedding_model = embedding_model
        writer.migration_revision = migration_revision
        for node_id, properties in nodes:
            properties = properties or {}
            node_type = properties.get("type")
            if node_type == "Entity" and properties.get("name"):
                writer.write(
                    COGXEntity(
                        external_system="cognee",
                        external_id=str(node_id),
                        name=properties["name"],
                        description=properties.get("description"),
                        created_at=parse_timestamp(properties.get("created_at")),
                        updated_at=parse_timestamp(properties.get("updated_at")),
                    )
                )
            elif node_type == "DocumentChunk" and properties.get("text"):
                writer.write(
                    COGXDocument(
                        external_system="cognee",
                        external_id=str(node_id),
                        content=properties["text"],
                        created_at=parse_timestamp(properties.get("created_at")),
                    )
                )
                # Also persist the chunk as a raw node: preserve-mode restore
                # rehydrates it as a graph node so facts referencing the chunk
                # (e.g. DocumentChunk -contains-> Entity) keep their topology
                # instead of dangling. The COGXDocument record above carries
                # the chunk's content for re-derive/cross-provider imports.
                writer.write_raw_node({"id": str(node_id), **properties})
            else:
                writer.write_raw_node({"id": str(node_id), **properties})

        for source, target, relationship, properties in edges:
            properties = properties or {}
            writer.write(
                COGXFact(
                    external_system="cognee",
                    external_id=f"{source}:{relationship}:{target}",
                    subject_ref=str(source),
                    predicate=str(relationship),
                    object_ref=str(target),
                    fact_text=properties.get("edge_text"),
                    valid_at=parse_timestamp(properties.get("valid_at")),
                    invalid_at=parse_timestamp(properties.get("invalid_at")),
                )
            )


async def export_dataset(
    dataset: Union[str, UUID] = "main_dataset",
    format: str = "pydantic",
    destination: Optional[Union[str, Path]] = None,
    user=None,
    link_relations: bool = False,
    include_permissions: bool = False,
) -> Union[ExportResult, GraphSnapshot]:
    """Export an authorized dataset's graph. Requires read permission.

    ``include_permissions`` (cogx only): additionally write the dataset's
    social layer — owner and ACL grants WITH user credentials — into
    ``permissions.json`` so a cognee-to-cognee restore recreates functional
    accounts and grants. The archive then contains password hashes: treat it
    as a secret.

    ``format="pydantic"`` returns a :class:`GraphSnapshot` (typed DataPoint
    instances, in memory; pass ``destination`` to also persist it as JSON).
    All other formats write a file and return an :class:`ExportResult`.
    ``link_relations`` (pydantic only) re-attaches edges as object references
    on declared relation fields, e.g. ``entity.is_a``.
    """
    if format not in EXPORT_FORMATS:
        raise ValueError(f"Unknown export format {format!r}. Expected one of {EXPORT_FORMATS}.")

    from cognee.context_global_variables import set_database_global_context_variables
    from cognee.infrastructure.databases.graph import get_graph_engine
    from cognee.modules.data.exceptions.exceptions import DatasetNotFoundError
    from cognee.modules.data.methods import get_authorized_existing_datasets
    from cognee.modules.users.methods import get_default_user

    if user is None:
        user = await get_default_user()

    datasets = await get_authorized_existing_datasets([dataset], "read", user)
    if not datasets:
        raise DatasetNotFoundError(message=f"Dataset not found or not readable: {dataset}")
    dataset_obj = datasets[0]

    async with set_database_global_context_variables(dataset_obj.id, dataset_obj.owner_id):
        graph_engine = await get_graph_engine()
        nodes, edges = await graph_engine.get_graph_data()

    # Ladybug's get_graph_data() synthesizes "SELF" self-loops when a graph
    # has no edges; they are not real relationships, so no emitter sees them.
    edges = [edge for edge in edges if not (edge[2] == "SELF" and edge[0] == edge[1])]

    if format == "pydantic":
        snapshot = build_snapshot(
            nodes,
            edges,
            dataset_name=dataset_obj.name,
            dataset_id=str(dataset_obj.id),
            link_relations=link_relations,
        )
        if destination is not None:
            snapshot.save(destination)
        logger.info(
            "Exported dataset %s as GraphSnapshot: %d nodes, %d edges",
            dataset_obj.name,
            len(snapshot.nodes),
            len(snapshot.edges),
        )
        return snapshot

    if destination is None:
        suffix = _FORMAT_SUFFIX.get(format, "")
        destination = f"{dataset_obj.name}_export{suffix}" if suffix else f"{dataset_obj.name}_cogx"
    destination = Path(destination)

    if format == "cogx":
        embedding_model = None
        try:
            from cognee.infrastructure.databases.vector.embeddings.config import (
                get_embedding_context_config,
            )

            embedding_model = get_embedding_context_config().embedding_model
        except Exception as error:  # noqa: BLE001 — manifest metadata is best effort
            logger.debug("Could not determine embedding model for manifest: %s", error)
        _write_cogx(
            nodes,
            edges,
            destination,
            dataset_obj.name,
            embedding_model=embedding_model,
            migration_revision=await _stored_migration_revision(dataset_obj.id),
        )
        if include_permissions:
            write_social_layer(destination, await _gather_social_layer(dataset_obj))
    elif format == "json":
        write_json(nodes, edges, destination)
    elif format == "graphml":
        write_graphml(nodes, edges, destination)
    elif format == "cypher":
        write_cypher(nodes, edges, destination)

    logger.info(
        "Exported dataset %s: %d nodes, %d edges -> %s (%s)",
        dataset_obj.name,
        len(nodes),
        len(edges),
        destination,
        format,
    )

    return ExportResult(
        format=format,
        destination=str(destination),
        dataset_name=dataset_obj.name,
        dataset_id=str(dataset_obj.id),
        num_nodes=len(nodes),
        num_edges=len(edges),
    )