From dfc4ca1cd01429ff5c1e597442d4113ac283a8b8 Mon Sep 17 00:00:00 2001
From: mvdbeek <m.vandenbeek@gmail.com>
Date: Thu, 26 Feb 2026 15:25:38 +0100
Subject: [PATCH] Enforce storage quota for Celery-based data fetch jobs

The Celery data fetch path (setup_fetch_data -> fetch_data ->
set_job_metadata -> finish_job) bypassed quota checks that the
traditional job runner path enforced via JobWrapper.enqueue().
This caused users over their disk quota to still be able to
fetch remote data when Celery task execution was enabled.

Add _pause_job_if_over_quota() call in setup_fetch_data after
object store IDs are assigned, matching the traditional path.
Propagate the paused state through the task chain so downstream
tasks (fetch_data, set_job_metadata, finish_job) gracefully
short-circuit when the job has been paused.

Fixes #21642
---
 lib/galaxy/celery/tasks.py                  | 12 +++++-
 test/integration/test_celery_fetch_quota.py | 47 +++++++++++++++++++++
 2 files changed, 57 insertions(+), 2 deletions(-)
 create mode 100644 test/integration/test_celery_fetch_quota.py

diff --git a/lib/galaxy/celery/tasks.py b/lib/galaxy/celery/tasks.py
index bfbd8b3155e..7a3ab934366 100644
--- a/lib/galaxy/celery/tasks.py
+++ b/lib/galaxy/celery/tasks.py
@@ -243,6 +243,11 @@ def setup_fetch_data(
     mini_job_wrapper.change_state(model.Job.states.QUEUED, flush=False, job=job)
     # Set object store after job destination so can leverage parameters...
     mini_job_wrapper._set_object_store_ids(job)
+    # Now that we have the object store id, check if we are over the limit
+    mini_job_wrapper._pause_job_if_over_quota(job)
+    if job.state == model.Job.states.PAUSED:
+        sa_session.commit()
+        return None
     request_json = Path(mini_job_wrapper.working_directory) / "request.json"
     request_json_value = next(iter(p.value for p in job.parameters if p.name == "request_json"))
     request_json.write_text(json.loads(request_json_value))
@@ -268,7 +273,8 @@ def finish_job(
     tool = cached_create_tool_from_representation(app=app, raw_tool_source=raw_tool_source)
     job = sa_session.get(Job, job_id)
     assert job
-    # TODO: assert state ?
+    if job.state == model.Job.states.PAUSED:
+        return
     mini_job_wrapper = MinimalJobWrapper(job=job, app=app, tool=tool)
     mini_job_wrapper.finish("", "")
 
@@ -327,7 +333,9 @@ def fetch_data(
     app: MinimalManagerApp,
     sa_session: galaxy_scoped_session,
     task_user_id: Optional[int] = None,
-) -> str:
+) -> Optional[str]:
+    if setup_return is None:
+        return None
     job = sa_session.get(Job, job_id)
     assert job
     mini_job_wrapper = MinimalJobWrapper(job=job, app=app)
diff --git a/test/integration/test_celery_fetch_quota.py b/test/integration/test_celery_fetch_quota.py
new file mode 100644
index 00000000000..ae2d53af453
--- /dev/null
+++ b/test/integration/test_celery_fetch_quota.py
@@ -0,0 +1,47 @@
+"""Integration test for quota enforcement during Celery-based data fetch.
+
+Verifies that when a user is over their disk quota, data fetch jobs
+submitted via the Celery task chain are correctly paused rather than
+being allowed to proceed with the download.
+"""
+
+from galaxy_test.base.populators import DatasetPopulator
+from galaxy_test.driver import integration_util
+
+
+class TestCeleryFetchQuotaEnforcement(integration_util.IntegrationTestCase):
+    dataset_populator: DatasetPopulator
+    require_admin_user = True
+
+    @classmethod
+    def handle_galaxy_config_kwds(cls, config):
+        super().handle_galaxy_config_kwds(config)
+        config["enable_quotas"] = True
+
+    def setUp(self):
+        super().setUp()
+        self.dataset_populator = DatasetPopulator(self.galaxy_interactor)
+
+    def test_fetch_paused_when_over_quota(self):
+        with self.dataset_populator.test_history() as history_id:
+            # Upload an initial dataset so the user has some disk usage
+            self.dataset_populator.new_dataset(history_id, content="initial content", wait=True)
+
+            # Set a very low quota (1 byte) so the user is over quota
+            payload = {
+                "name": "default-celery-fetch-quota",
+                "description": "very low default quota for testing",
+                "amount": "1 bytes",
+                "operation": "=",
+                "default": "registered",
+            }
+            self.dataset_populator.create_quota(payload)
+
+            # Now try to upload another dataset - should be paused due to quota
+            hda = self.dataset_populator.new_dataset(history_id, content="more data", wait=False)
+            self.dataset_populator.wait_for_history(history_id, assert_ok=False)
+
+            details = self.dataset_populator.get_history_dataset_details(
+                history_id, dataset=hda, wait=False, assert_ok=False
+            )
+            assert details["state"] == "paused", f"Expected dataset state 'paused', got '{details['state']}'"