fix: Resolve dataset storage under its owner

cognee/context_global_variables.py

@@ -110,6 +110,35 @@ def backend_access_control_enabled():
 GRAPH_DBS_WITH_MULTI_USER_SUPPORT = ["ladybug", "kuzu", "falkor", "postgres"]
 
 
+async def _get_dataset_owner_id(dataset_id: UUID) -> UUID:
+    """Return the owner id of an existing dataset; raise if it does not exist.
+
+    No authorization here on purpose — this resolves where the dataset's
+    databases live (storage identity), not who may touch them.
+    """
+    # Imported lazily to avoid circular imports at module load.
+    from cognee.infrastructure.databases.relational import get_relational_engine
+    from cognee.modules.data.models import Dataset
+    from cognee.modules.data.exceptions import DatasetNotFoundError
+
+    db_engine = get_relational_engine()
+    async with db_engine.get_async_session() as session:
+        dataset = await session.get(Dataset, dataset_id)
+
+    if dataset is None:
+        raise DatasetNotFoundError(f"Dataset {dataset_id} does not exist.")
+    return dataset.owner_id
+
+
+async def _verify_dataset_access(user_id: UUID, dataset_id: UUID, permission_type: str) -> None:
+    """Raise PermissionDeniedError unless the user holds the given permission
+    on the dataset."""
+    # Imported lazily to avoid circular imports at module load.
+    from cognee.modules.users.permissions.methods import get_specific_user_permission_datasets
+
+    await get_specific_user_permission_datasets(user_id, permission_type, [dataset_id])
+
+
 class DatabaseContextManager:
     """Dual-mode helper returned by :func:`set_database_global_context_variables`.
 
@@ -123,6 +152,7 @@ class DatabaseContextManager:
         "_user_id",
         "_llm_config",
         "_embedding_config",
+        "_permission_type",
         "_applied",
         "_dataset_token",
         "_llm_token",
@@ -132,21 +162,26 @@ class DatabaseContextManager:
     def __init__(
         self,
         dataset: Optional[UUID],
-        user_id: UUID,
+        user_id: Optional[UUID] = None,
         llm_config: Optional[LLMConfig] = None,
         embedding_config: Optional[EmbeddingConfig] = None,
+        permission_type: Optional[str] = None,
     ) -> None:
         self._dataset = dataset
         self._user_id = user_id
         self._llm_config = llm_config
         self._embedding_config = embedding_config
+        self._permission_type = permission_type
         self._applied = False
         self._dataset_token = None
         self._llm_token = None
         self._embedding_token = None
 
     async def apply_database_context_variables(
-        self, dataset: Optional[UUID], user_id: UUID
+        self,
+        dataset: Optional[UUID],
+        user_id: Optional[UUID] = None,
+        permission_type: Optional[str] = None,
     ) -> None:
         # current_dataset_id always carries a dataset *id* (a UUID object) or
         # None. Exactly one input type: callers resolve names/strings to a UUID
@@ -184,7 +219,22 @@ async def apply_database_context_variables(
 
         await dataset_queue().ensure_slot(dataset)
 
-        user = await get_user(user_id)
+        # Optional permission gate: checked only when the caller asked for it
+        # by passing a permission_type — callers that already authorized at the
+        # API layer pass nothing and no check is performed here.
+        if permission_type is not None:
+            if user_id is None:
+                raise CogneeValidationError(
+                    "permission_type requires a user_id to check the permission for."
+                )
+            await _verify_dataset_access(user_id, dataset, permission_type)
+
+        # Storage is namespaced by the dataset's OWNER, never by the caller:
+        # the dataset_database row, the physical database paths, and the data
+        # root all live under the owner (#4829). Resolve the owner from the
+        # dataset itself so a caller identity can never redirect them.
+        owner_id = await _get_dataset_owner_id(dataset)
+        user = await get_user(owner_id)
 
         # To ensure permissions are enforced properly all datasets will have their own databases
         dataset_database = await get_or_create_dataset_database(dataset, user)
@@ -277,7 +327,9 @@ async def apply_database_context_variables(
     async def _apply(self) -> None:
         if self._applied:
             return
-        await self.apply_database_context_variables(self._dataset, self._user_id)
+        await self.apply_database_context_variables(
+            self._dataset, self._user_id, self._permission_type
+        )
         self._applied = True
 
     def __await__(self):
@@ -323,9 +375,10 @@ async def __aexit__(self, exc_type, exc, tb) -> None:
 
 def set_database_global_context_variables(
     dataset: Optional[UUID],
-    user_id: UUID,
+    user_id: Optional[UUID] = None,
     llm_config: Optional[LLMConfig] = None,
     embedding_config: Optional[EmbeddingConfig] = None,
+    permission_type: Optional[str] = None,
 ) -> "DatabaseContextManager":
     """Returns a dual-mode helper that is both awaitable and an async context manager.
 
@@ -358,14 +411,22 @@ def set_database_global_context_variables(
     callers can keep reading the dataset databases.
 
     Args:
-        dataset: Cognee dataset name or id
-        user_id: UUID of the owner of the dataset
+        dataset: Cognee dataset id
+        user_id: Optional UUID of the user entering the context. Storage
+            (registry row, database paths, data root) always resolves under
+            the dataset's owner, never this value — it is consulted only by
+            the optional permission check below, and therefore required when
+            ``permission_type`` is set.
         llm_config: Optional ``LLMConfig`` to use for LLM calls in this context.
         embedding_config: Optional ``EmbeddingConfig`` to use for embedding calls
             in this context.
+        permission_type: Optional permission ("read"/"write"/"delete"/"share")
+            to verify for ``user_id`` on the dataset before entering. When
+            omitted, no permission check is performed here — callers are
+            expected to have authorized at the API layer.
 
     Returns:
         A :class:`DatabaseContextManager` that can be awaited or used as an
         async context manager.
     """
-    return DatabaseContextManager(dataset, user_id, llm_config, embedding_config)
+    return DatabaseContextManager(dataset, user_id, llm_config, embedding_config, permission_type)

cognee/infrastructure/databases/utils/get_or_create_dataset_database.py

@@ -41,25 +41,26 @@ async def _get_graph_db_info(dataset_id: UUID, user: User) -> dict:
 
 async def _existing_dataset_database(
     dataset_id: UUID,
-    user: User,
 ) -> Optional[DatasetDatabase]:
     """
-    Check if a DatasetDatabase row already exists for the given owner + dataset.
+    Check if a DatasetDatabase row already exists for the given dataset.
     Return None if it doesn't exist, return the row if it does.
+
+    dataset_id is the table's primary key — one row per dataset, shared by the
+    owner and every ACL-granted user — so the lookup must not filter by owner:
+    for a non-owner caller that turns a hit into a miss and the fall-through
+    INSERT crashes on the primary key (#4829).
+
     Args:
         dataset_id:
-        user:
 
     Returns:
         DatasetDatabase or None
     """
     db_engine = get_relational_engine()
 
     async with db_engine.get_async_session() as session:
-        stmt = select(DatasetDatabase).where(
-            DatasetDatabase.owner_id == user.id,
-            DatasetDatabase.dataset_id == dataset_id,
-        )
+        stmt = select(DatasetDatabase).where(DatasetDatabase.dataset_id == dataset_id)
         existing: DatasetDatabase = await session.scalar(stmt)
         return existing
 
@@ -79,7 +80,10 @@ async def get_or_create_dataset_database(
     Parameters
     ----------
     user : User
-        Principal that owns this dataset.
+        Principal that owns this dataset. An existing row is returned regardless
+        of who calls (the row is keyed by dataset alone), but when the row does
+        not exist yet this user is recorded as its owner — so creation must
+        happen as the owner, never as an ACL-granted caller.
     dataset : Union[str, UUID]
         Dataset being linked.
     """
@@ -99,7 +103,7 @@ async def get_or_create_dataset_database(
         dataset = await create_authorized_dataset(dataset, user)
 
     # If dataset database already exists return it
-    existing_dataset_database = await _existing_dataset_database(dataset_id, user)
+    existing_dataset_database = await _existing_dataset_database(dataset_id)
     if existing_dataset_database:
         return existing_dataset_database
 

cognee/tests/test_delete_permission.py

@@ -120,7 +120,19 @@ class CustomData(BaseModel):
     nodes, edges = await graph_engine.get_graph_data()
     assert len(nodes) == 2 and len(edges) == 1, "Nodes and edges are not deleted properly."
 
-    await datasets.delete_data(dataset.id, data2.id, user2)
+    # Regression for #4829: the dataset_database row is keyed by dataset_id alone,
+    # so an ACL-granted non-owner must get the owner's existing row back instead
+    # of a false negative that crashes on a duplicate INSERT.
+    from cognee.infrastructure.databases.utils import get_or_create_dataset_database
+
+    dataset_database = await get_or_create_dataset_database(dataset.id, user2)
+    assert dataset_database.owner_id == user1.id, (
+        "get_or_create_dataset_database must return the owner's existing row for a grantee."
+    )
+
+    # Regression for #4829: forget() as an ACL-granted non-owner must work end-to-end
+    # (it used to crash provisioning a duplicate dataset_database row on context entry).
+    await cognee.forget(data_id=data2.id, dataset_id=dataset.id, user=user2)
 
     nodes, edges = await graph_engine.get_graph_data()
     assert len(nodes) == 0 and len(edges) == 0, "Nodes and edges are not deleted."