Merge pull request #19323 from davelopez/24.1_fix_quota_usage_with_user_object_stores

[24.1] Fix quota usage with user object stores

lib/galaxy/jobs/__init__.py

@@ -70,6 +70,7 @@ from galaxy.model.base import transaction
 from galaxy.model.store import copy_dataset_instance_metadata_attributes
 from galaxy.model.store.discover import MaxDiscoveredFilesExceededError
 from galaxy.objectstore import (
+    is_user_object_store,
     ObjectStorePopulator,
     serialize_static_object_store_config,
 )
@@ -2328,7 +2329,7 @@ class MinimalJobWrapper(HasResourceParameters):
         required_user_object_store_uris = set()
         for out_dataset_instance in out_data.values():
             object_store_id = out_dataset_instance.dataset.object_store_id
-            if object_store_id and object_store_id.startswith("user_objects://"):
+            if is_user_object_store(object_store_id):
                 required_user_object_store_uris.add(object_store_id)
 
         job_metadata = os.path.join(self.tool_working_directory, self.tool.provided_metadata_file)

lib/galaxy/managers/object_store_instances.py

@@ -31,6 +31,7 @@ from galaxy.objectstore import (
     build_test_object_store_from_user_config,
     ConcreteObjectStoreModel,
     QuotaModel,
+    USER_OBJECTS_SCHEME,
     UserObjectStoresAppConfig,
 )
 from galaxy.objectstore.badges import serialize_badges
@@ -300,7 +301,7 @@ class ObjectStoreInstancesManager:
         )
         secrets = persisted_object_store.template_secrets or []
         uuid = str(persisted_object_store.uuid)
-        object_store_id = f"user_objects://{uuid}"
+        object_store_id = f"{USER_OBJECTS_SCHEME}{uuid}"
 
         return UserConcreteObjectStoreModel(
             uuid=uuid,

lib/galaxy/model/__init__.py

@@ -162,7 +162,10 @@ from galaxy.model.item_attrs import (
 )
 from galaxy.model.orm.now import now
 from galaxy.model.orm.util import add_object_to_object_session
-from galaxy.objectstore import ObjectStorePopulator
+from galaxy.objectstore import (
+    ObjectStorePopulator,
+    USER_OBJECTS_SCHEME,
+)
 from galaxy.objectstore.templates import (
     ObjectStoreConfiguration,
     ObjectStoreTemplate,
@@ -644,6 +647,7 @@ FROM dataset
 LEFT OUTER JOIN library_dataset_dataset_association ON dataset.id = library_dataset_dataset_association.dataset_id
 WHERE dataset.id IN (SELECT dataset_id FROM per_hist_hdas)
     AND library_dataset_dataset_association.id IS NULL
+    AND (dataset.object_store_id NOT LIKE '{user_objects_scheme}%' OR dataset.object_store_id IS NULL)
     {and_dataset_condition}
 """
 
@@ -659,7 +663,9 @@ def calculate_user_disk_usage_statements(user_id, quota_source_map, for_sqlite=F
     default_usage_dataset_condition = f"{default_cond} {use_or} {exclude_cond}"
     if default_usage_dataset_condition.strip():
         default_usage_dataset_condition = f"AND ( {default_usage_dataset_condition} )"
-    default_usage = UNIQUE_DATASET_USER_USAGE.format(and_dataset_condition=default_usage_dataset_condition)
+    default_usage = UNIQUE_DATASET_USER_USAGE.format(
+        and_dataset_condition=default_usage_dataset_condition, user_objects_scheme=USER_OBJECTS_SCHEME
+    )
     default_usage = f"""
 UPDATE galaxy_user SET disk_usage = ({default_usage})
 WHERE id = :id
@@ -673,7 +679,8 @@ WHERE id = :id
     # the object_store_id to quota_source_label into a temp table of values
     for quota_source_label, object_store_ids in source.items():
         label_usage = UNIQUE_DATASET_USER_USAGE.format(
-            and_dataset_condition="AND ( dataset.object_store_id IN :include_object_store_ids )"
+            and_dataset_condition="AND ( dataset.object_store_id IN :include_object_store_ids )",
+            user_objects_scheme=USER_OBJECTS_SCHEME,
         )
         if for_sqlite:
             # hacky alternative for older sqlite
@@ -1126,7 +1133,9 @@ ON CONFLICT
             if exclude_objectstore_ids
             else ""
         )
-        default_usage = UNIQUE_DATASET_USER_USAGE.format(and_dataset_condition=default_usage_dataset_condition)
+        default_usage = UNIQUE_DATASET_USER_USAGE.format(
+            and_dataset_condition=default_usage_dataset_condition, user_objects_scheme=USER_OBJECTS_SCHEME
+        )
         sql_calc = text(default_usage)
         params = {"id": self.id}
         bindparams = [bindparam("id")]

lib/galaxy/objectstore/__init__.py

@@ -75,9 +75,15 @@ DEFAULT_PRIVATE = False
 DEFAULT_QUOTA_SOURCE = None  # Just track quota right on user object in Galaxy.
 DEFAULT_QUOTA_ENABLED = True  # enable quota tracking in object stores by default
 DEFAULT_DEVICE_ID = None
+USER_OBJECTS_SCHEME = "user_objects://"
+
 log = logging.getLogger(__name__)
 
 
+def is_user_object_store(object_store_id: Optional[str]) -> bool:
+    return object_store_id is not None and object_store_id.startswith(USER_OBJECTS_SCHEME)
+
+
 class UserObjectStoreResolver(Protocol):
     def resolve_object_store_uri_config(self, uri: str) -> ObjectStoreConfiguration:
         pass
@@ -1288,7 +1294,7 @@ class DistributedObjectStore(NestedObjectStore):
         try:
             return self.backends[object_store_id]
         except KeyError:
-            if object_store_id.startswith("user_objects://") and self.user_object_store_resolver:
+            if is_user_object_store(object_store_id) and self.user_object_store_resolver:
                 return self.user_object_store_resolver.resolve_object_store_uri(object_store_id)
             raise
 
@@ -1326,7 +1332,7 @@ class DistributedObjectStore(NestedObjectStore):
 
     def __get_store_id_for(self, obj, **kwargs):
         if obj.object_store_id is not None:
-            if obj.object_store_id in self.backends or obj.object_store_id.startswith("user_objects://"):
+            if obj.object_store_id in self.backends or is_user_object_store(obj.object_store_id):
                 return obj.object_store_id
             else:
                 log.warning(
@@ -1365,7 +1371,7 @@ class DistributedObjectStore(NestedObjectStore):
         if parent_check or object_store_id is None:
             return parent_check
         # user selection allowed and object_store_id is not None
-        if object_store_id.startswith("user_objects://"):
+        if is_user_object_store(object_store_id):
             if not user:
                 return "Supplied object store id is not accessible"
             rest_of_uri = object_store_id.split("://", 1)[1]
@@ -1756,12 +1762,16 @@ class QuotaSourceMap:
         self.default_quota_source = source
         self.default_quota_enabled = enabled
         self.info = QuotaSourceInfo(self.default_quota_source, self.default_quota_enabled)
+        # User defined sources are provided by the user and the quota is not tracked
+        self.user_defined_source_info = QuotaSourceInfo(label=None, use=False)
         self.backends = {}
         self._labels = None
 
-    def get_quota_source_info(self, object_store_id):
+    def get_quota_source_info(self, object_store_id: Optional[str]) -> QuotaSourceInfo:
         if object_store_id in self.backends:
             return self.backends[object_store_id].get_quota_source_info(object_store_id)
+        elif is_user_object_store(object_store_id):
+            return self.user_defined_source_info
         else:
             return self.info
 

test/unit/data/test_quota.py

@@ -111,6 +111,21 @@ class TestCalculateUsage(BaseModelTestCase):
 
         assert u.calculate_disk_usage_default_source(object_store) == 10
 
+    def test_calculate_usage_with_user_provided_storage(self):
+        u = self.u
+
+        self._add_dataset(10)
+        # This dataset should not be counted towards the user's disk usage
+        self._add_dataset(30, object_store_id="user_objects://user/provided/storage")
+
+        object_store = MockObjectStore()
+        assert u.calculate_disk_usage_default_source(object_store) == 10
+        assert u.disk_usage is None
+        u.calculate_and_set_disk_usage(object_store)
+        assert u.calculate_disk_usage_default_source(object_store) == 10
+
+        self._refresh_user_and_assert_disk_usage_is(10)
+
     def test_calculate_usage_readjusts_incorrect_quota(self):
         u = self.u