Quota recalculation logic consolidation

lib/galaxy/model/__init__.py

@@ -28,6 +28,7 @@
 import galaxy.security.passwords
 import galaxy.util
 from galaxy.model.item_attrs import UsesAnnotations
+from galaxy.model.util import pgcalc
 from galaxy.security import get_permitted_actions
 from galaxy.util import (directory_hash_id, Params, ready_name_for_url,
                          restore_text, send_mail, unicodify, unique_id)
@@ -275,6 +276,31 @@ def calculate_disk_usage( self ):
                     total += hda.dataset.get_total_size()
         return total
 
+    def calculate_and_set_disk_usage( self ):
+        """
+        Calculates and sets user disk usage.
+        """
+        new = None
+        db_session = object_session(self)
+        current = self.get_disk_usage()
+        if db_session.get_bind().dialect.name not in ( 'postgres', 'postgresql' ):
+            done = False
+            while not done:
+                new = self.calculate_disk_usage()
+                db_session.refresh( self )
+                # make sure usage didn't change while calculating
+                # set done if it has not, otherwise reset current and iterate again.
+                if self.get_disk_usage() == current:
+                    done = True
+                else:
+                    current = self.get_disk_usage()
+        else:
+            new = pgcalc(db_session, self.id)
+        if new not in (current, None):
+            self.set_disk_usage( new )
+            db_session.add( self )
+            db_session.flush()
+
     @staticmethod
     def user_template_environment( user ):
         """

lib/galaxy/queue_worker.py

@@ -9,7 +9,6 @@
 
 import galaxy.queues
 from galaxy import util
-from galaxy.model.util import pgcalc
 
 from kombu import Connection
 from kombu.mixins import ConsumerMixin
@@ -143,13 +142,11 @@ def recalculate_user_disk_usage(app, **kwargs):
     if user_id:
         user = sa_session.query( app.model.User ).get( app.security.decode_id( user_id ) )
         if user:
-            if sa_session.get_bind().dialect.name not in ( 'postgres', 'postgresql' ):
-                new = user.calculate_disk_usage()
-            else:
-                new = pgcalc(sa_session, user.id)
-            user.set_disk_usage(new)
-            sa_session.add(user)
-            sa_session.flush()
+            user.calculate_and_set_disk_usage()
+        else:
+            log.error("Recalculate user disk usage task failed, user %s not found" % user_id)
+    else:
+        log.error("Recalculate user disk usage task received without user_id.")
 
 
 def reload_tool_data_tables(app, **kwargs):

lib/galaxy/webapps/galaxy/controllers/admin.py

@@ -10,7 +10,6 @@
 from galaxy.actions.admin import AdminActions
 from galaxy.exceptions import MessageException
 from galaxy.model import tool_shed_install as install_model
-from galaxy.model.util import pgcalc
 from galaxy.util import nice_size, sanitize_text, url_get
 from galaxy.util.odict import odict
 from galaxy.web import url_for
@@ -874,27 +873,9 @@ def recalculate_user_disk_usage( self, trans, **kwd ):
         user = trans.sa_session.query( trans.model.User ).get( trans.security.decode_id( user_id ) )
         if not user:
             return trans.show_error_message( "User not found for id (%s)" % sanitize_text( str( user_id ) ) )
-        engine = None
-        if trans.app.config.database_connection:
-            engine = trans.app.config.database_connection.split(':')[0]
-        if engine not in ( 'postgres', 'postgresql' ):
-            done = False
-            while not done:
-                current = user.get_disk_usage()
-                new = user.calculate_disk_usage()
-                trans.sa_session.refresh( user )
-                # make sure usage didn't change while calculating, set done
-                if user.get_disk_usage() == current:
-                    done = True
-                if new not in (current, None):
-                    user.set_disk_usage( new )
-                    trans.sa_session.add( user )
-                    trans.sa_session.flush()
-        else:
-            # We can use the lightning fast pgcalc!
-            current = user.get_disk_usage()
-            new = pgcalc( self.sa_session, user.id )
-        # yes, still a small race condition between here and the flush
+        current = user.get_disk_usage()
+        user.calculate_and_set_disk_usage()
+        new = user.get_disk_usage()
         if new in ( current, None ):
             message = 'Usage is unchanged at %s.' % nice_size( current )
         else: