From 88770b1b111fe43d72f50e8c6507f65820b145f6 Mon Sep 17 00:00:00 2001
From: Igor Ilic <30923996+dexters1@users.noreply.github.com>
Date: Wed, 17 Sep 2025 11:16:20 +0200
Subject: [PATCH] refactor: Use awaitable attrs for getting roles (#1405)

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

## Description
Resolve issue with getting role objects for user

## Type of Change
<!-- Please check the relevant option -->
- [ ] Bug fix (non-breaking change that fixes an issue)
- [ ] New feature (non-breaking change that adds functionality)
- [ ] Breaking change (fix or feature that would cause existing
functionality to change)
- [ ] Documentation update
- [ ] Code refactoring
- [ ] Performance improvement
- [ ] Other (please specify):

## Changes Made
<!-- List the specific changes made in this PR -->
-
-
-

## Testing
<!-- Describe how you tested your changes -->

## Screenshots/Videos (if applicable)
<!-- Add screenshots or videos to help explain your changes -->

## Pre-submission Checklist
<!-- Please check all boxes that apply before submitting your PR -->
- [ ] **I have tested my changes thoroughly before submitting this PR**
- [ ] **This PR contains minimal changes necessary to address the
issue/feature**
- [ ] My code follows the project's coding standards and style
guidelines
- [ ] I have added tests that prove my fix is effective or that my
feature works
- [ ] I have added necessary documentation (if applicable)
- [ ] All new and existing tests pass
- [ ] I have searched existing PRs to ensure this change hasn't been
submitted already
- [ ] I have linked any relevant issues in the description
- [ ] My commits have clear and descriptive messages

## Related Issues
<!-- Link any related issues using "Fixes #issue_number" or "Relates to
#issue_number" -->

## Additional Notes
<!-- Add any additional notes, concerns, or context for reviewers -->

## 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 <boris@topoteretes.com>
---
 .../infrastructure/databases/relational/ModelBase.py  |  3 ++-
 .../methods/get_all_user_permission_datasets.py       | 11 +++++++++--
 2 files changed, 11 insertions(+), 3 deletions(-)

diff --git a/cognee/infrastructure/databases/relational/ModelBase.py b/cognee/infrastructure/databases/relational/ModelBase.py
index a4d3a1a19..3a2054207 100644
--- a/cognee/infrastructure/databases/relational/ModelBase.py
+++ b/cognee/infrastructure/databases/relational/ModelBase.py
@@ -1,7 +1,8 @@
 from sqlalchemy.orm import DeclarativeBase
+from sqlalchemy.ext.asyncio import AsyncAttrs
 
 
-class Base(DeclarativeBase):
+class Base(AsyncAttrs, DeclarativeBase):
     """
     Represents a base class for declarative models using SQLAlchemy.
 
diff --git a/cognee/modules/users/permissions/methods/get_all_user_permission_datasets.py b/cognee/modules/users/permissions/methods/get_all_user_permission_datasets.py
index 5b242baa4..a8731a773 100644
--- a/cognee/modules/users/permissions/methods/get_all_user_permission_datasets.py
+++ b/cognee/modules/users/permissions/methods/get_all_user_permission_datasets.py
@@ -1,3 +1,5 @@
+from types import SimpleNamespace
+
 from cognee.shared.logging_utils import get_logger
 
 from ...models.User import User
@@ -17,9 +19,14 @@ async def get_all_user_permission_datasets(user: User, permission_type: str) ->
         # Get all datasets all tenants have access to
         tenant = await get_tenant(user.tenant_id)
         datasets.extend(await get_principal_datasets(tenant, permission_type))
+
         # Get all datasets Users roles have access to
-        for role_name in user.roles:
-            role = await get_role(user.tenant_id, role_name)
+        if isinstance(user, SimpleNamespace):
+            # If simple namespace use roles defined in user
+            roles = user.roles
+        else:
+            roles = await user.awaitable_attrs.roles
+        for role in roles:
             datasets.extend(await get_principal_datasets(role, permission_type))
 
     # Deduplicate datasets with same ID