Fixes to the sqlalchemy adapter

cognee/api/v1/cognify/cognify_v2.py

@@ -40,8 +40,12 @@ async def cognify(datasets: Union[str, list[str]] = None, user: User = None):
     if datasets is None or len(datasets) == 0:
         return await cognify(await db_engine.get_datasets())
 
+
+    db_engine = get_relational_engine()
+    async with db_engine.get_async_session() as session:
+
         if user is None:
-        user = await get_default_user()
+            user = await get_default_user(session=  session)
 
         async def run_cognify_pipeline(dataset_name: str, files: list[dict]):
             documents = [
@@ -52,7 +56,7 @@ async def cognify(datasets: Union[str, list[str]] = None, user: User = None):
                 for file in files
             ]
 
-        await check_permissions_on_documents(user, "read", [document.id for document in documents])
+            await check_permissions_on_documents(user, "read", [document.id for document in documents], session=session)
 
             async with update_status_lock:
                 task_status = get_task_status([dataset_name])

cognee/modules/users/methods/get_default_user.py

@@ -3,9 +3,8 @@ from cognee.infrastructure.databases.relational import get_relational_engine
 
 from sqlalchemy.future import select
 
-async def get_default_user():
+async def get_default_user(session):
-    db_engine = get_relational_engine()
+
-    async with db_engine.get_async_session() as session:
     stmt = select(User).where(User.email == "default_user@example.com")
     result = await session.execute(stmt)
     user = result.scalars().first()

cognee/modules/users/permissions/methods/check_permissions_on_documents.py

@@ -1,32 +1,37 @@
 import  logging
+
+from sqlalchemy import select
+
 from cognee.infrastructure.databases.relational import get_relational_engine
 from ...models.User import User
 from ...models.ACL import ACL
 
 logger = logging.getLogger(__name__)
 
-async def check_permissions_on_documents(
+class PermissionDeniedException(Exception):
-    user: User,
+    def __init__(self, message: str):
-    permission_type: str,
+        self.message = message
-    document_ids: list[str],
+        super().__init__(self.message)
-):
-    try:
-        relational_engine = get_relational_engine()
 
-        async with relational_engine.get_async_session() as session:
+
+async def check_permissions_on_documents(user: User, permission_type: str, document_ids: list[str], session):
+    try:
         user_group_ids = [group.id for group in user.groups]
 
-            acls = session.query(ACL) \
+        result = await session.execute(
-                .filter(ACL.principal_id.in_([user.id, *user_group_ids])) \
+            select(ACL).filter(
-                .filter(ACL.permission.name == permission_type) \
+                ACL.principal_id.in_([user.id, *user_group_ids]),
-                .all()
+                ACL.permission.name == permission_type
+            )
+        )
+        acls = result.scalars().all()
 
-            resource_ids = [resource.resource_id for resource in acl.resources for acl in acls]
+        resource_ids = [resource.resource_id for acl in acls for resource in acl.resources]
-
+        has_permissions = all(document_id in resource_ids for document_id in document_ids)
-            has_permissions = all([document_id in resource_ids for document_id in document_ids])
 
         if not has_permissions:
-                raise Exception(f"User {user.username} does not have {permission_type} permission on documents")
+            raise PermissionDeniedException(f"User {user.username} does not have {permission_type} permission on documents")
     except Exception as error:
         logger.error("Error checking permissions on documents: %s", str(error))
-        raise error
+        raise
+