Map DefaultQuotaAssociation declaratively

assuntad23 · Sep 7, 2021 · 37edac4 · 37edac4
1 parent 06df91c
commit 37edac4
Show file tree

Hide file tree

Showing 4 changed files with 27 additions and 15 deletions.
diff --git a/lib/galaxy/model/__init__.py b/lib/galaxy/model/__init__.py
@@ -2460,6 +2460,7 @@ class Quota(Base, Dictifiable, RepresentById):
     bytes = Column('bytes', BigInteger)
     operation = Column('operation', String(8))
     deleted = Column('deleted', Boolean, index=True, default=False)
+    default = relationship('DefaultQuotaAssociation', back_populates='quota')
     groups = relationship('GroupQuotaAssociation', back_populates='quota')
     users = relationship('UserQuotaAssociation', back_populates='quota')
 
@@ -2496,7 +2497,16 @@ def display_amount(self):
             return galaxy.util.nice_size(self.bytes)
 
 
-class DefaultQuotaAssociation(Dictifiable, RepresentById):
+class DefaultQuotaAssociation(Base, Dictifiable, RepresentById):
+    __tablename__ = 'default_quota_association'
+
+    id = Column('id', Integer, primary_key=True)
+    create_time = Column('create_time', DateTime, default=now)
+    update_time = Column('update_time', DateTime, default=now, onupdate=now)
+    type = Column('type', String(32), index=True, unique=True)
+    quota_id = Column('quota_id', Integer, ForeignKey('quota.id'), index=True)
+    quota = relationship('Quota', back_populates='default')
+
     dict_element_visible_keys = ['type']
 
     class types(str, Enum):

diff --git a/lib/galaxy/model/mapping.py b/lib/galaxy/model/mapping.py
@@ -240,14 +240,6 @@
     Column("create_time", DateTime, default=now),
     Column("update_time", DateTime, default=now, onupdate=now))
 
-model.DefaultQuotaAssociation.table = Table(
-    "default_quota_association", metadata,
-    Column("id", Integer, primary_key=True),
-    Column("create_time", DateTime, default=now),
-    Column("update_time", DateTime, default=now, onupdate=now),
-    Column("type", String(32), index=True, unique=True),
-    Column("quota_id", Integer, ForeignKey("quota.id"), index=True))
-
 model.LibraryPermissions.table = Table(
     "library_permissions", metadata,
     Column("id", Integer, primary_key=True),
@@ -1689,10 +1681,6 @@ def simple_mapping(model, **kwds):
     )
 ))
 
-mapper_registry.map_imperatively(model.DefaultQuotaAssociation, model.DefaultQuotaAssociation.table, properties=dict(
-    quota=relation(model.Quota, backref="default")
-))
-
 mapper_registry.map_imperatively(model.LibraryPermissions, model.LibraryPermissions.table, properties=dict(
     library=relation(model.Library, backref="actions"),
     role=relation(model.Role, backref="library_actions")

diff --git a/lib/galaxy/quota/__init__.py b/lib/galaxy/quota/__init__.py
@@ -144,7 +144,7 @@ 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).filter(self.model.DefaultQuotaAssociation.table.c.type == default_type).first()
+        dqa = self.sa_session.query(self.model.DefaultQuotaAssociation).filter(self.model.DefaultQuotaAssociation.type == default_type).first()
         if not dqa:
             return None
         if dqa.quota.bytes < 0:
@@ -161,7 +161,7 @@ def set_default_quota(self, default_type, quota):
         for gqa in quota.groups:
             self.sa_session.delete(gqa)
         # Find the old default, assign the new quota if it exists
-        dqa = self.sa_session.query(self.model.DefaultQuotaAssociation).filter(self.model.DefaultQuotaAssociation.table.c.type == default_type).first()
+        dqa = self.sa_session.query(self.model.DefaultQuotaAssociation).filter(self.model.DefaultQuotaAssociation.type == default_type).first()
         if dqa:
             dqa.quota = quota
         # Or create if necessary

diff --git a/test/unit/model/test_model.py b/test/unit/model/test_model.py
@@ -94,6 +94,20 @@ def test_DatasetPermissions(model, session, dataset, role):
         assert stored_obj.role == role
 
 
+def test_DefaultQuotaAssociation(model, session, quota):
+    cls = model.DefaultQuotaAssociation
+    assert cls.__tablename__ == 'default_quota_association'
+    with dbcleanup(session, cls):
+        type = cls.types.REGISTERED
+        obj = cls(type, quota)
+        obj_id = persist(session, obj)
+
+        stored_obj = get_stored_obj(session, cls, obj_id)
+        assert stored_obj.id == obj_id
+        assert stored_obj.type == type
+        assert stored_obj.quota == quota
+
+
 def test_DefaultUserPermissions(model, session, user, role):
     cls = model.DefaultUserPermissions
     assert cls.__tablename__ == 'default_user_permissions'