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']}'"