get_quota in SQL

galaxyproject · Jul 26, 2022 · fbb9bbf · fbb9bbf
1 parent 0f77b3e
commit fbb9bbf
Show file tree

Hide file tree

Showing 2 changed files with 58 additions and 42 deletions.
diff --git a/lib/galaxy/quota/__init__.py b/lib/galaxy/quota/__init__.py
@@ -1,6 +1,8 @@
 """Galaxy Quotas"""
 import logging
 
+from sqlalchemy.sql import text
+
 import galaxy.util
 
 log = logging.getLogger(__name__)
@@ -98,51 +100,64 @@ def get_quota(self, user):
                quotas.
         """
         if not user:
-            return self.default_unregistered_quota
-        quotas = []
-        for group in [uga.group for uga in user.groups]:
-            for quota in [gqa.quota for gqa in group.quotas]:
-                if quota not in quotas:
-                    quotas.append(quota)
-        for quota in [uqa.quota for uqa in user.quotas]:
-            if quota not in quotas:
-                quotas.append(quota)
-        use_default = True
-        max = 0
-        adjustment = 0
-        rval = 0
-        for quota in quotas:
-            if quota.deleted:
-                continue
-            if quota.operation == "=" and quota.bytes == -1:
-                rval = None
-                break
-            elif quota.operation == "=":
-                use_default = False
-                if quota.bytes > max:
-                    max = quota.bytes
-            elif quota.operation == "+":
-                adjustment += quota.bytes
-            elif quota.operation == "-":
-                adjustment -= quota.bytes
-        if use_default:
-            max = self.default_registered_quota
-            if max is None:
-                rval = None
-        if rval is not None:
-            rval = max + adjustment
-            if rval <= 0:
-                rval = 0
-        return rval
+            return self._default_unregistered_quota
+        query = text(
+            """
+SELECT (
+        COALESCE(MAX(CASE WHEN union_quota.operation = '='
+                          THEN union_quota.bytes
+                          ELSE NULL
+                          END),
+                 (SELECT default_quota.bytes
+                  FROM quota as default_quota
+                      LEFT JOIN default_quota_association on default_quota.id = default_quota_association.quota_id
+                      WHERE default_quota_association.type == 'registered'
+                          AND default_quota.deleted != :is_true))
+        +
+        (CASE WHEN SUM(CASE WHEN union_quota.operation = '=' AND union_quota.bytes = -1
+                            THEN 1 ELSE 0
+                            END) > 0
+              THEN NULL
+              ELSE 0 END)
+        +
+        (COALESCE(SUM(
+                CASE WHEN union_quota.operation = '+' THEN union_quota.bytes
+                     WHEN union_quota.operation = '-' THEN -1 * union_quota.bytes
+                     ELSE 0
+                     END
+              ), 0))
+       )
+FROM (
+    SELECT user_quota.operation as operation, user_quota.bytes as bytes
+    FROM galaxy_user as user
+        LEFT JOIN user_quota_association as uqa on user.id = uqa.user_id
+        LEFT JOIN quota as user_quota on user_quota.id = uqa.quota_id
+    WHERE user_quota.deleted != :is_true
+        AND user.id = :user_id
+    UNION ALL
+    SELECT group_quota.operation as operation, group_quota.bytes as bytes
+    FROM galaxy_user as user
+        LEFT JOIN user_group_association as uga on user.id = uga.user_id
+        LEFT JOIN galaxy_group on galaxy_group.id = uga.group_id
+        LEFT JOIN group_quota_association as gqa on galaxy_group.id = gqa.group_id
+        LEFT JOIN quota as group_quota on group_quota.id = gqa.quota_id
+    WHERE group_quota.deleted != :is_true
+        AND user.id = :user_id
+) as union_quota
+"""
+        )
+        conn = self.sa_session.connection()
+        with conn.begin():
+            res = conn.execute(query, is_true=True, user_id=user.id).fetchone()
+            if res:
+                return res[0]
+            else:
+                return None
 
     @property
-    def default_unregistered_quota(self):
+    def _default_unregistered_quota(self):
         return self._default_quota(self.model.DefaultQuotaAssociation.types.UNREGISTERED)
 
-    @property
-    def default_registered_quota(self):
-        return self._default_quota(self.model.DefaultQuotaAssociation.types.REGISTERED)
-
     def _default_quota(self, default_type):
         dqa = (
             self.sa_session.query(self.model.DefaultQuotaAssociation)

diff --git a/test/unit/data/test_quota.py b/test/unit/data/test_quota.py
@@ -98,7 +98,8 @@ def _add_user_quota(self, user, quota):
         self.persist(quota, uqa, user)
 
     def _assert_user_quota_is(self, user, amount):
-        assert amount == self.quota_agent.get_quota(user)
+        actual_quota = self.quota_agent.get_quota(user)
+        assert amount == actual_quota, "Expected quota [%s], got [%s]" % (amount, actual_quota)
         if amount is None:
             user.total_disk_usage = 1000
             job = model.Job()