feat: new Dataset permissions (#869)

<!-- .github/pull_request_template.md -->

## Description
<!-- Provide a clear description of the changes in this PR -->

## DCO Affirmation
I affirm that all code in every commit of this pull request conforms to
the terms of the Topoteretes Developer Certificate of Origin.

---------

Co-authored-by: Boris Arzentar <borisarzentar@gmail.com>
Co-authored-by: Boris <boris@topoteretes.com>

Author: dexters1

.env.template

@@ -69,3 +69,11 @@ LITELLM_LOG="ERROR"
 # Set this environment variable to disable sending telemetry data
 # TELEMETRY_DISABLED=1
 
+# Set this variable to True to enforce usage of backend access control for Cognee
+# Note: This is only currently supported by the following databases:
+#       Relational: SQLite, Postgres
+#       Vector: LanceDB
+#       Graph: KuzuDB
+#
+# It enforces LanceDB and KuzuDB use and uses them to create databases per Cognee user + dataset
+ENABLE_BACKEND_ACCESS_CONTROL=False

.github/workflows/e2e_tests.yml

@@ -215,3 +215,34 @@ jobs:
             AWS_ACCESS_KEY_ID: ${{ secrets.AWS_ACCESS_KEY_ID }}
             AWS_SECRET_ACCESS_KEY: ${{ secrets.AWS_SECRET_ACCESS_KEY }}
           run: poetry run python ./cognee/tests/test_s3.py
+
+  test-parallel-databases:
+    name: Test using different async databases in parallel in Cognee
+    runs-on: ubuntu-22.04
+    steps:
+      - name: Check out repository
+        uses: actions/checkout@v4
+
+      - name: Cognee Setup
+        uses: ./.github/actions/cognee_setup
+        with:
+          python-version: '3.11.x'
+
+      - name: Install specific graph db dependency
+        run: |
+          poetry install -E kuzu
+
+      - name: Run parallel databases test
+        env:
+          ENV: 'dev'
+          LLM_MODEL: ${{ secrets.LLM_MODEL }}
+          LLM_ENDPOINT: ${{ secrets.LLM_ENDPOINT }}
+          LLM_API_KEY: ${{ secrets.LLM_API_KEY }}
+          LLM_API_VERSION: ${{ secrets.LLM_API_VERSION }}
+          EMBEDDING_MODEL: ${{ secrets.EMBEDDING_MODEL }}
+          EMBEDDING_ENDPOINT: ${{ secrets.EMBEDDING_ENDPOINT }}
+          EMBEDDING_API_KEY: ${{ secrets.EMBEDDING_API_KEY }}
+          EMBEDDING_API_VERSION: ${{ secrets.EMBEDDING_API_VERSION }}
+          AWS_ACCESS_KEY_ID: ${{ secrets.AWS_ACCESS_KEY_ID }}
+          AWS_SECRET_ACCESS_KEY: ${{ secrets.AWS_SECRET_ACCESS_KEY }}
+        run: poetry run python ./cognee/tests/test_parallel_databases.py

.github/workflows/test_suites.yml

@@ -47,7 +47,7 @@ jobs:
 
   relational-db-migration-tests:
     name: Relational DB Migration Tests
-    needs: [ basic-tests, e2e-tests ]
+    needs: [ basic-tests, e2e-tests, graph-db-tests]
     uses: ./.github/workflows/relational_db_migration_tests.yml
     secrets: inherit
 
@@ -79,7 +79,7 @@ jobs:
 
   db-examples-tests:
     name: DB Examples Tests
-    needs: [vector-db-tests]
+    needs: [vector-db-tests, graph-db-tests, relational-db-migration-tests]
     uses: ./.github/workflows/db_examples_tests.yml
     secrets: inherit
 

.github/workflows/vector_db_tests.yml

@@ -135,6 +135,16 @@ jobs:
       run:
         shell: bash
 
+    services:
+      qdrant:
+        image: qdrant/qdrant:v1.14.1
+        env:
+          QDRANT__LOG_LEVEL: ERROR
+          QDRANT__SERVICE__API_KEY: qdrant_api_key
+          QDRANT__SERVICE__ENABLE_TLS: 0
+        ports:
+          - 6333:6333
+
     steps:
       - name: Check out
         uses: actions/checkout@master
@@ -148,6 +158,19 @@ jobs:
         run: |
           poetry install -E qdrant
 
+      - name: Wait for Qdrant to be healthy
+        run: |
+          for i in {1..10}; do
+            if curl -f http://127.0.0.1:6333/healthz; then
+              echo "Qdrant is healthy!"
+              exit 0
+            fi
+            echo "Waiting for Qdrant to be healthy..."
+            sleep 3
+          done
+          echo "Qdrant failed to become healthy in time"
+          exit 1
+
       - name: Run default Qdrant
         env:
           ENV: 'dev'
@@ -159,8 +182,8 @@ jobs:
           EMBEDDING_ENDPOINT: ${{ secrets.EMBEDDING_ENDPOINT }}
           EMBEDDING_API_KEY: ${{ secrets.EMBEDDING_API_KEY }}
           EMBEDDING_API_VERSION: ${{ secrets.EMBEDDING_API_VERSION }}
-          VECTOR_DB_URL: ${{ secrets.QDRANT_API_URL }}
-          VECTOR_DB_KEY: ${{ secrets.QDRANT_API_KEY }}
+          VECTOR_DB_URL: 127.0.0.1
+          VECTOR_DB_KEY: qdrant_api_key
         run: poetry run python ./cognee/tests/test_qdrant.py
 
   run-postgres-tests:

cognee/api/client.py

@@ -1,6 +1,7 @@
 """FastAPI server for the Cognee API."""
 
 import os
+
 import uvicorn
 from cognee.shared.logging_utils import get_logger
 import sentry_sdk
@@ -63,6 +64,7 @@ async def lifespan(app: FastAPI):
 
 app = FastAPI(debug=app_environment != "prod", lifespan=lifespan)
 
+
 app.add_middleware(
     CORSMiddleware,
     allow_origins=["*"],

cognee/api/v1/add/add.py

@@ -1,3 +1,4 @@
+from uuid import UUID
 from typing import Union, BinaryIO, List, Optional
 
 from cognee.modules.pipelines import Task
@@ -11,9 +12,21 @@ async def add(
     dataset_name: str = "main_dataset",
     user: User = None,
     node_set: Optional[List[str]] = None,
+    vector_db_config: dict = None,
+    graph_db_config: dict = None,
+    dataset_id: UUID = None,
 ):
-    tasks = [Task(resolve_data_directories), Task(ingest_data, dataset_name, user, node_set)]
+    tasks = [
+        Task(resolve_data_directories),
+        Task(ingest_data, dataset_name, user, node_set, dataset_id),
+    ]
 
     await cognee_pipeline(
-        tasks=tasks, datasets=dataset_name, data=data, user=user, pipeline_name="add_pipeline"
+        tasks=tasks,
+        datasets=dataset_id if dataset_id else dataset_name,
+        data=data,
+        user=user,
+        pipeline_name="add_pipeline",
+        vector_db_config=vector_db_config,
+        graph_db_config=graph_db_config,
     )

cognee/api/v1/add/routers/get_add_router.py

@@ -1,4 +1,5 @@
 from uuid import UUID
+
 from fastapi import Form, UploadFile, Depends
 from fastapi.responses import JSONResponse
 from fastapi import APIRouter
@@ -20,8 +21,8 @@ def get_add_router() -> APIRouter:
     @router.post("/", response_model=None)
     async def add(
         data: List[UploadFile],
+        datasetName: str,
         datasetId: Optional[UUID] = Form(default=None),
-        datasetName: Optional[str] = Form(default=None),
         user: User = Depends(get_authenticated_user),
     ):
         """This endpoint is responsible for adding data to the graph."""
@@ -30,19 +31,13 @@ async def add(
         if not datasetId and not datasetName:
             raise ValueError("Either datasetId or datasetName must be provided.")
 
-        if datasetId and not datasetName:
-            dataset = await get_dataset(user_id=user.id, dataset_id=datasetId)
-            try:
-                datasetName = dataset.name
-            except IndexError:
-                raise ValueError("No dataset found with the provided datasetName.")
-
         try:
             if isinstance(data, str) and data.startswith("http"):
                 if "github" in data:
                     # Perform git clone if the URL is from GitHub
                     repo_name = data.split("/")[-1].replace(".git", "")
                     subprocess.run(["git", "clone", data, f".data/{repo_name}"], check=True)
+                    # TODO: Update add call with dataset info
                     await cognee_add(
                         "data://.data/",
                         f"{repo_name}",
@@ -53,10 +48,10 @@ async def add(
                     response.raise_for_status()
 
                     file_data = await response.content()
-
+                    # TODO: Update add call with dataset info
                     return await cognee_add(file_data)
             else:
-                await cognee_add(data, datasetName, user=user)
+                await cognee_add(data, dataset_name=datasetName, user=user, dataset_id=datasetId)
         except Exception as error:
             return JSONResponse(status_code=409, content={"error": str(error)})
 

cognee/api/v1/cognify/cognify.py

@@ -9,7 +9,7 @@
 from cognee.modules.users.models import User
 from cognee.shared.data_models import KnowledgeGraph
 from cognee.tasks.documents import (
-    check_permissions_on_documents,
+    check_permissions_on_dataset,
     classify_documents,
     extract_chunks_from_documents,
 )
@@ -31,11 +31,18 @@ async def cognify(
     chunker=TextChunker,
     chunk_size: int = None,
     ontology_file_path: Optional[str] = None,
+    vector_db_config: dict = None,
+    graph_db_config: dict = None,
 ):
     tasks = await get_default_tasks(user, graph_model, chunker, chunk_size, ontology_file_path)
 
     return await cognee_pipeline(
-        tasks=tasks, datasets=datasets, user=user, pipeline_name="cognify_pipeline"
+        tasks=tasks,
+        datasets=datasets,
+        user=user,
+        pipeline_name="cognify_pipeline",
+        vector_db_config=vector_db_config,
+        graph_db_config=graph_db_config,
     )
 
 
@@ -48,7 +55,7 @@ async def get_default_tasks(  # TODO: Find out a better way to do this (Boris's
 ) -> list[Task]:
     default_tasks = [
         Task(classify_documents),
-        Task(check_permissions_on_documents, user=user, permissions=["write"]),
+        Task(check_permissions_on_dataset, user=user, permissions=["write"]),
         Task(
             extract_chunks_from_documents,
             max_chunk_size=chunk_size or get_max_chunk_tokens(),

cognee/api/v1/cognify/routers/get_cognify_router.py

@@ -1,3 +1,4 @@
+from uuid import UUID
 from typing import List, Optional
 from pydantic import BaseModel
 from fastapi import Depends
@@ -10,6 +11,7 @@
 
 class CognifyPayloadDTO(BaseModel):
     datasets: List[str]
+    dataset_ids: Optional[List[UUID]]
     graph_model: Optional[BaseModel] = KnowledgeGraph
 
 
@@ -22,7 +24,9 @@ async def cognify(payload: CognifyPayloadDTO, user: User = Depends(get_authentic
         from cognee.api.v1.cognify import cognify as cognee_cognify
 
         try:
-            await cognee_cognify(payload.datasets, user, payload.graph_model)
+            # Send dataset UUIDs if they are given, if not send dataset names
+            datasets = payload.dataset_ids if payload.dataset_ids else payload.datasets
+            await cognee_cognify(datasets, user, payload.graph_model)
         except Exception as error:
             return JSONResponse(status_code=409, content={"error": str(error)})
 

cognee/api/v1/permissions/routers/get_permissions_router.py

@@ -1,66 +1,69 @@
 from uuid import UUID
+from typing import List
 
-from fastapi import APIRouter
+from fastapi import APIRouter, Depends
 from fastapi.responses import JSONResponse
 
+from cognee.modules.users.models import User
+from cognee.modules.users.methods import get_authenticated_user
+
 
 def get_permissions_router() -> APIRouter:
     permissions_router = APIRouter()
 
-    @permissions_router.post("/roles/{role_id}/permissions")
-    async def give_default_permission_to_role(role_id: UUID, permission_name: str):
-        from cognee.modules.users.permissions.methods import (
-            give_default_permission_to_role as set_default_permission_to_role,
-        )
-
-        await set_default_permission_to_role(role_id, permission_name)
-
-        return JSONResponse(status_code=200, content={"message": "Permission assigned to role"})
+    @permissions_router.post("/datasets/{principal_id}/")
+    async def give_datasets_permission_to_principal(
+        permission_name: str,
+        dataset_ids: List[UUID],
+        principal_id: UUID,
+        user: User = Depends(get_authenticated_user),
+    ):
+        from cognee.modules.users.permissions.methods import authorized_give_permission_on_datasets
 
-    @permissions_router.post("/tenants/{tenant_id}/permissions")
-    async def give_default_permission_to_tenant(tenant_id: UUID, permission_name: str):
-        from cognee.modules.users.permissions.methods import (
-            give_default_permission_to_tenant as set_tenant_default_permissions,
+        await authorized_give_permission_on_datasets(
+            principal_id,
+            [dataset_id for dataset_id in dataset_ids],
+            permission_name,
+            user.id,
         )
 
-        await set_tenant_default_permissions(tenant_id, permission_name)
-
-        return JSONResponse(status_code=200, content={"message": "Permission assigned to tenant"})
-
-    @permissions_router.post("/users/{user_id}/permissions")
-    async def give_default_permission_to_user(user_id: UUID, permission_name: str):
-        from cognee.modules.users.permissions.methods import (
-            give_default_permission_to_user as set_default_permission_to_user,
+        return JSONResponse(
+            status_code=200, content={"message": "Permission assigned to principal"}
         )
 
-        await set_default_permission_to_user(user_id, permission_name)
-
-        return JSONResponse(status_code=200, content={"message": "Permission assigned to user"})
-
     @permissions_router.post("/roles")
-    async def create_role(
-        role_name: str,
-        tenant_id: UUID,
-    ):
+    async def create_role(role_name: str, user: User = Depends(get_authenticated_user)):
         from cognee.modules.users.roles.methods import create_role as create_role_method
 
-        await create_role_method(role_name=role_name, tenant_id=tenant_id)
+        await create_role_method(role_name=role_name, owner_id=user.id)
 
         return JSONResponse(status_code=200, content={"message": "Role created for tenant"})
 
     @permissions_router.post("/users/{user_id}/roles")
-    async def add_user_to_role(user_id: UUID, role_id: UUID):
+    async def add_user_to_role(
+        user_id: UUID, role_id: UUID, user: User = Depends(get_authenticated_user)
+    ):
         from cognee.modules.users.roles.methods import add_user_to_role as add_user_to_role_method
 
-        await add_user_to_role_method(user_id=user_id, role_id=role_id)
+        await add_user_to_role_method(user_id=user_id, role_id=role_id, owner_id=user.id)
 
         return JSONResponse(status_code=200, content={"message": "User added to role"})
 
+    @permissions_router.post("/users/{user_id}/tenants")
+    async def add_user_to_tenant(
+        user_id: UUID, tenant_id: UUID, user: User = Depends(get_authenticated_user)
+    ):
+        from cognee.modules.users.tenants.methods import add_user_to_tenant
+
+        await add_user_to_tenant(user_id=user_id, tenant_id=tenant_id, owner_id=user.id)
+
+        return JSONResponse(status_code=200, content={"message": "User added to tenant"})
+
     @permissions_router.post("/tenants")
-    async def create_tenant(tenant_name: str):
+    async def create_tenant(tenant_name: str, user: User = Depends(get_authenticated_user)):
         from cognee.modules.users.tenants.methods import create_tenant as create_tenant_method
 
-        await create_tenant_method(tenant_name=tenant_name)
+        await create_tenant_method(tenant_name=tenant_name, user_id=user.id)
 
         return JSONResponse(status_code=200, content={"message": "Tenant created."})