Optimize email_in_org_cached

trypromptly · Sep 27, 2024 · 245b057 · 245b057
1 parent 3c89bb5
commit 245b057
Show file tree

Hide file tree

Showing 3 changed files with 11 additions and 11 deletions.
diff --git a/llmstack/apps/models.py b/llmstack/apps/models.py
@@ -387,7 +387,7 @@ def has_write_permission(self, user):
             self.owner == user
             or user.email in self.write_accessible_by
             or (
-                Profile.email_in_org_cached(self.owner_profile.organization, user.email)
+                Profile.user_in_org_cached(self.owner_profile.organization, user)
                 and self.visibility >= AppVisibility.ORGANIZATION
                 and self.owner_profile.organization.settings.default_app_access_permission == AppAccessPermission.WRITE
             )
@@ -402,7 +402,7 @@ def has_read_permission(self, user):
             or self.has_write_permission(user)
             or user.email in self.read_accessible_by
             or (
-                Profile.email_in_org_cached(self.owner_profile.organization, user.email)
+                Profile.user_in_org_cached(self.owner_profile.organization, user)
                 and self.visibility >= AppVisibility.ORGANIZATION
             )
         )

diff --git a/llmstack/base/models.py b/llmstack/base/models.py
@@ -732,13 +732,13 @@ def get_by_encrypted_uuid(cls, encrypted_uuid):
 
     @staticmethod
     @cache
-    def email_in_org_cached(organization, email):
-        profiles = Profile.objects.filter(user__email=email)
-        return (
-            profiles.exists()
-            and profiles.values("organization").distinct().count() == 1
-            and profiles.first().organization == organization
-        )
+    def user_in_org_cached(organization, user):
+        profile = Profile.objects.get(user=user)
+
+        if not profile:
+            return False
+
+        return profile.organization == organization
 
 
 class DefaultProfile(AbstractProfile):

diff --git a/llmstack/data/models.py b/llmstack/data/models.py
@@ -173,7 +173,7 @@ def has_read_permission(self, user):
             self.owner == user
             or user.email in self.read_accessible_by
             or (
-                Profile.email_in_org_cached(self.profile.organization, user.email)
+                Profile.user_in_org_cached(self.profile.organization, user)
                 and self.visibility == DataSourceVisibility.ORGANIZATION
             )
         )
@@ -183,7 +183,7 @@ def has_write_permission(self, user):
             self.owner == user
             or user.email in self.write_accessible_by
             or (
-                Profile.email_in_org_cached(self.profile.organization, user.email)
+                Profile.user_in_org_cached(self.profile.organization, user)
                 and self.visibility == DataSourceVisibility.ORGANIZATION
                 and self.profile.organization.settings.default_datasource_access_permission
                 == DataSourceAccessPermission.WRITE