cognee 1.2.2.dev2__py3-none-any.whl → 1.2.2.dev4__py3-none-any.whl

cognee/api/v1/export/export.py

@@ -18,6 +18,7 @@ async def export(
     destination: Optional[Union[str, Path]] = None,
     user=None,
     link_relations: bool = False,
+    include_permissions: bool = False,
 ) -> GraphSnapshot: ...
 
 
@@ -28,6 +29,7 @@ async def export(
     destination: Optional[Union[str, Path]] = None,
     user=None,
     link_relations: bool = False,
+    include_permissions: bool = False,
 ) -> ExportResult: ...
 
 
@@ -39,6 +41,7 @@ async def export(
     destination: Optional[Union[str, Path]] = None,
     user=None,
     link_relations: bool = False,
+    include_permissions: bool = False,
 ) -> ExportResult: ...
 
 
@@ -48,6 +51,7 @@ async def export(
     destination: Optional[Union[str, Path]] = None,
     user=None,
     link_relations: bool = False,
+    include_permissions: bool = False,
 ) -> Union[ExportResult, GraphSnapshot]:
     """Export a dataset's knowledge graph.
 
@@ -97,4 +101,5 @@ async def export(
             destination=destination,
             user=user,
             link_relations=link_relations,
+            include_permissions=include_permissions,
         )

cognee/modules/migration/cogx.py

@@ -169,6 +169,7 @@ RECORD_FILES: Dict[str, str] = {
 
 MANIFEST_FILE = "manifest.json"
 RAW_NODES_FILE = "nodes.jsonl"
+PERMISSIONS_FILE = "permissions.json"
 
 
 class COGXManifest(BaseModel):
@@ -177,6 +178,11 @@ class COGXManifest(BaseModel):
     exported_at: Optional[datetime] = None
     counts: Dict[str, int] = Field(default_factory=dict)
     embedding_model: Optional[str] = None
+    # Cognee-origin archives only: the source store's stamped data-migration
+    # revision at export time. Raw nodes keep their source-store ids, so the
+    # importing store must not claim a newer revision than the exported data
+    # actually has (None = unknown; the import keeps its own stamp).
+    migration_revision: Optional[str] = None
     notes: List[str] = Field(default_factory=list)
 
 
@@ -209,12 +215,13 @@ class COGXArchiveWriter:
     def __init__(self, directory: Union[str, Path], source_system: str = "cognee"):
         self.directory = Path(directory)
         self.directory.mkdir(parents=True, exist_ok=True)
-        for file_name in (*RECORD_FILES.values(), RAW_NODES_FILE, MANIFEST_FILE):
+        for file_name in (*RECORD_FILES.values(), RAW_NODES_FILE, MANIFEST_FILE, PERMISSIONS_FILE):
             (self.directory / file_name).unlink(missing_ok=True)
         self.source_system = source_system
         self.counts: Dict[str, int] = {}
         self.notes: List[str] = []
         self.embedding_model: Optional[str] = None
+        self.migration_revision: Optional[str] = None
         self._handles: Dict[str, Any] = {}
 
     def __enter__(self) -> "COGXArchiveWriter":
@@ -254,6 +261,7 @@ class COGXArchiveWriter:
                 exported_at=datetime.now(timezone.utc),
                 counts=self.counts,
                 embedding_model=self.embedding_model,
+                migration_revision=self.migration_revision,
                 notes=self.notes,
             )
             manifest_path = self.directory / MANIFEST_FILE
@@ -269,6 +277,25 @@ def read_manifest(directory: Union[str, Path]) -> Optional[COGXManifest]:
     return manifest
 
 
+def write_social_layer(directory: Union[str, Path], payload: Dict[str, Any]) -> None:
+    """Persist the dataset's social layer (owner + grants, WITH credentials).
+
+    Written only on explicit request (``export(include_permissions=True)``):
+    the payload carries user emails, password HASHES, and account flags, so an
+    archive containing this file must be handled as a secret, not just data.
+    """
+    (Path(directory) / PERMISSIONS_FILE).write_text(
+        json.dumps(payload, indent=2, default=str), encoding="utf-8"
+    )
+
+
+def read_social_layer(directory: Union[str, Path]) -> Optional[Dict[str, Any]]:
+    permissions_path = Path(directory) / PERMISSIONS_FILE
+    if not permissions_path.exists():
+        return None
+    return json.loads(permissions_path.read_text(encoding="utf-8"))
+
+
 def _archive_file(base: Path, file_name: str) -> Optional[Path]:
     """Resolve ``base / file_name`` and confirm it stays inside ``base``.
 

cognee/modules/migration/export.py

@@ -14,7 +14,7 @@ Formats:
 
 from dataclasses import dataclass
 from pathlib import Path
-from typing import Optional, Union
+from typing import Any, Dict, Optional, Union
 from uuid import UUID
 
 from cognee.modules.migration.cogx import (
@@ -23,6 +23,7 @@ from cognee.modules.migration.cogx import (
     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
@@ -51,13 +52,98 @@ class ExportResult:
         )
 
 
+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
+    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")
@@ -112,9 +198,16 @@ async def export_dataset(
     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`.
@@ -179,7 +272,16 @@ async def export_dataset(
             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)
+        _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":

cognee/modules/migration/import_source.py

@@ -42,6 +42,210 @@ DATA_ITEMS_PER_ADD = 200
 _GRAPH_RECORD_KINDS = ("entity", "fact", "raw_node")
 
 
+async def _ensure_user(user_payload: Dict[str, Any]):
+    """Create-or-match a user by email, transferring credentials on creation.
+
+    An existing target user is returned untouched — their credentials are
+    never clobbered by an import. A missing user is created and then given
+    the archived credentials (hashed password + account flags) directly, so
+    restored accounts authenticate with their original passwords.
+    """
+    import secrets
+
+    from cognee.infrastructure.databases.relational import get_relational_engine
+    from cognee.modules.users.methods import create_user, get_user_by_email
+    from cognee.modules.users.models import User
+
+    existing = await get_user_by_email(user_payload["email"])
+    if existing is not None:
+        return existing
+
+    created = await create_user(user_payload["email"], secrets.token_urlsafe(32))
+    db_engine = get_relational_engine()
+    async with db_engine.get_async_session() as session:
+        record = await session.get(User, created.id)
+        record.hashed_password = user_payload["hashed_password"]
+        record.is_active = user_payload.get("is_active", True)
+        record.is_superuser = user_payload.get("is_superuser", False)
+        record.is_verified = user_payload.get("is_verified", False)
+        await session.commit()
+    logger.info("Restored user %s from archive social layer.", user_payload["email"])
+    return created
+
+
+async def _resolve_import_user(source: MemorySource, user):
+    """The identity the import runs as.
+
+    Archives carrying a social layer import AS the archived dataset OWNER
+    (created/matched by email first): per-dataset databases derive their
+    physical location from the owner id, so ownership must be right BEFORE
+    the rows land — it cannot be reassigned afterwards. All other imports
+    run as the caller's user, exactly as before.
+
+    Processing a social layer requires a SUPERUSER importer: the archive
+    supplies emails, password hashes, and account flags verbatim, so an
+    unprivileged importer could otherwise mint arbitrary accounts (including
+    superusers) with credentials of their choosing — both via the SDK and via
+    the /v1/remember archive-upload endpoint.
+    """
+    social_layer = getattr(source, "social_layer", None)
+    owner_payload = (social_layer or {}).get("owner")
+    if owner_payload is None:
+        return user
+
+    importer = user
+    if importer is None:
+        from cognee.modules.users.methods import get_default_user
+
+        importer = await get_default_user()
+    if not importer.is_superuser:
+        from cognee.modules.users.exceptions.exceptions import PermissionDeniedError
+
+        raise PermissionDeniedError(
+            message="Importing an archive that carries a social layer (permissions.json) "
+            "requires a superuser: it restores user accounts and credentials."
+        )
+    return await _ensure_user(owner_payload)
+
+
+async def _apply_social_grants(source: MemorySource, dataset_name: str, owner, importer) -> None:
+    """Re-apply the archive's ACL grants onto the freshly imported dataset.
+
+    Users are created/matched by email (credentials transfer on creation);
+    ``give_permission_on_dataset`` deduplicates existing ACL rows, so
+    re-importing is idempotent. The importing user additionally keeps read
+    access — they held the archive bytes, and this prevents silent lockout
+    when restoring someone else's dataset.
+    """
+    social_layer = getattr(source, "social_layer", None)
+    if not social_layer:
+        return
+
+    from cognee.modules.data.methods import get_authorized_existing_datasets
+    from cognee.modules.users.permissions.methods import give_permission_on_dataset
+
+    datasets = await get_authorized_existing_datasets([dataset_name], "read", owner)
+    if not datasets:
+        logger.warning(
+            "No dataset %r found after import; cannot restore its social layer.", dataset_name
+        )
+        return
+    dataset_id = datasets[0].id
+
+    for grant in social_layer.get("grants", []):
+        principal = await _ensure_user(grant["user"])
+        for permission_name in grant.get("permissions", []):
+            await give_permission_on_dataset(principal, dataset_id, permission_name)
+
+    if importer is None:
+        from cognee.modules.users.methods import get_default_user
+
+        importer = await get_default_user()
+    if importer.id != owner.id:
+        await give_permission_on_dataset(importer, dataset_id, "read")
+
+
+def _revision_to_stamp(
+    archive_revision: Optional[str],
+    stored_revision: Optional[str],
+    ordered_revisions: List[str],
+) -> Optional[str]:
+    """The revision the imported store should be re-stamped at, or None.
+
+    Stamps only BACKWARD — when the archive's revision is strictly behind the
+    store's current stamp — so the next migration run replays exactly
+    archive -> head over the imported rows (idempotent over already-current
+    data). Never forward: stamping ahead would mark the store's own data as
+    migrated when it is not. Unknown revisions (either side written by newer
+    code) and an unstamped store (None = base, already minimal) leave the
+    stamp untouched.
+    """
+    if archive_revision is None or stored_revision is None:
+        return None
+    if archive_revision not in ordered_revisions or stored_revision not in ordered_revisions:
+        return None
+    if ordered_revisions.index(archive_revision) < ordered_revisions.index(stored_revision):
+        return archive_revision
+    return None
+
+
+async def _restamp_to_source_revision(source: MemorySource, dataset_name: str, user) -> None:
+    """Align the target's migration stamp with a cognee-origin archive.
+
+    Preserve/hybrid imports write the archive's raw nodes with their
+    source-store ids, so the target must not claim a newer data-migration
+    revision than the exported data actually has. When the archive carries a
+    revision behind the target's stamp, re-stamp backward; the next migration
+    gate then replays revision -> head over the imported data. External
+    sources carry no revision (their records are written entirely by
+    current-code pipelines) and are skipped, as are re-derive imports (no raw
+    nodes land).
+    """
+    archive_revision = getattr(source, "migration_revision", None)
+    if archive_revision is None or source.mode == "re-derive":
+        return
+
+    from cognee.context_global_variables import backend_access_control_enabled
+    from cognee.infrastructure.databases.relational import get_relational_engine
+    from cognee.modules.migrations.migration import order_migrations
+    from cognee.modules.migrations.registry import MIGRATIONS
+    from cognee.modules.migrations.runner import stamp_revisions
+
+    ordered_revisions = [migration.revision for migration in order_migrations(MIGRATIONS)]
+    if archive_revision not in ordered_revisions:
+        logger.warning(
+            "Archive migration revision %r is unknown to this chain — the archive was "
+            "exported by newer code; leaving the store's migration stamp unchanged.",
+            archive_revision,
+        )
+        return
+
+    db_engine = get_relational_engine()
+    if backend_access_control_enabled():
+        from cognee.modules.data.methods import get_authorized_existing_datasets
+        from cognee.modules.users.methods import get_default_user
+        from cognee.modules.users.models import DatasetDatabase
+
+        if user is None:
+            user = await get_default_user()
+        datasets = await get_authorized_existing_datasets([dataset_name], "read", user)
+        if not datasets:
+            logger.warning(
+                "No dataset %r found after import; cannot align its migration stamp.",
+                dataset_name,
+            )
+            return
+        dataset_id = datasets[0].id
+        async with db_engine.get_async_session() as session:
+            record = await session.get(DatasetDatabase, dataset_id)
+        stored_revision = record.migration_revision if record else None
+        target = _revision_to_stamp(archive_revision, stored_revision, ordered_revisions)
+        if target is None:
+            return
+        await stamp_revisions(target=target, dataset_ids=[dataset_id])
+    else:
+        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)
+        stored_revision = record.global_migration_revision if record else None
+        target = _revision_to_stamp(archive_revision, stored_revision, ordered_revisions)
+        if target is None:
+            return
+        await stamp_revisions(target=target)
+
+    logger.info(
+        "Stamped store back to archive migration revision %r (was %r); the next "
+        "migration run replays %r -> head over the imported data.",
+        target,
+        stored_revision,
+        target,
+    )
+
+
 def _pipeline_run_id(pipeline_result: Any) -> Optional[str]:
     """Extract the pipeline run id from a run_custom_pipeline return value.
 
@@ -84,14 +288,34 @@ async def import_memory_source(
     # before its own gate). This also records the data-migration revision —
     # stamping a fresh store at head — BEFORE the imported rows arrive;
     # without it the populated store has no recorded revision and the first
-    # migration-aware startup replays the entire data chain over it.
+    # migration-aware startup replays the entire data chain over it. It must
+    # run before user resolution: on a fresh store it also creates the
+    # relational schema that user lookup needs.
     await run_migrations_and_block(dataset_name, user)
 
+    # Archives carrying a social layer import AS the archived owner (see
+    # _resolve_import_user); everything else runs as the caller's user.
+    importer = user
+    user = await _resolve_import_user(source, user)
+
     node_set = node_set or [f"import:{source.source_system}"]
 
     if source.mode == "preserve" and getattr(source, "replayable", False):
-        return await _import_streaming(source, dataset_name, user, run_in_background, node_set)
-    return await _import_buffered(source, dataset_name, user, run_in_background, node_set, **kwargs)
+        result = await _import_streaming(source, dataset_name, user, run_in_background, node_set)
+    else:
+        result = await _import_buffered(
+            source, dataset_name, user, run_in_background, node_set, **kwargs
+        )
+
+    # After the rows land: cognee-origin archives may need the migration
+    # stamp aligned backward to the SOURCE store's revision (see
+    # _restamp_to_source_revision) and their social layer restored (grants
+    # re-applied for recreated users).
+    await _restamp_to_source_revision(source, dataset_name, user)
+    if getattr(source, "social_layer", None):
+        await _apply_social_grants(source, dataset_name, owner=user, importer=importer)
+
+    return result
 
 
 async def _import_streaming(
@@ -130,7 +354,12 @@ async def _import_streaming(
 
     logger.info("Importing from %s (mode=preserve, streaming): %s", source.source_system, counts)
 
-    stats: Dict[str, int] = {"graph_nodes": 0, "graph_edges": 0, "skipped_facts": 0}
+    stats: Dict[str, int] = {
+        "graph_nodes": 0,
+        "graph_edges": 0,
+        "skipped_facts": 0,
+        "deduped_edges": 0,
+    }
     pipeline_result = None
     has_graph_records = any(counts.get(kind) for kind in _GRAPH_RECORD_KINDS)
     if has_graph_records:
@@ -175,6 +404,7 @@ async def _import_streaming(
         "graph_nodes": stats["graph_nodes"],
         "graph_edges": stats["graph_edges"],
         "skipped_facts": stats["skipped_facts"],
+        "deduped_edges": stats["deduped_edges"],
         "pipeline_run_id": run_id,
     }
     if backgrounded:
@@ -208,7 +438,13 @@ async def _import_buffered(
 
     # Translate the record stream directly: no full raw-record list is kept,
     # so peak memory is bounded by the translation output alone.
-    translation = await translate_record_stream(source.records(), source.mode)
+    translation = await translate_record_stream(
+        source.records(),
+        source.mode,
+        # Cognee-origin archives keep source node UUIDs verbatim; other
+        # systems get class-namespaced ids (see _register_entity).
+        preserve_source_ids=source.source_system == "cognee",
+    )
 
     logger.info(
         "Importing %d records from %s (mode=%s): %s",