diff --git a/lib/galaxy/celery/base_task.py b/lib/galaxy/celery/base_task.py index 72447276539..6737aba0231 100644 --- a/lib/galaxy/celery/base_task.py +++ b/lib/galaxy/celery/base_task.py @@ -80,37 +80,28 @@ class GalaxyTaskBeforeStartUserRateLimitPostgres(GalaxyTaskBeforeStartUserRateLi We take advantage of efficiencies in its dialect. """ - _update_stmt = ( - update(CeleryUserRateLimit) - .where(CeleryUserRateLimit.user_id == bindparam("userid")) - .values(last_scheduled_time=text("greatest(last_scheduled_time + ':interval second', " ":now) ")) - .returning(CeleryUserRateLimit.last_scheduled_time) - ) - _insert_stmt = ( - ps_insert(CeleryUserRateLimit) - .values(user_id=bindparam("userid"), last_scheduled_time=bindparam("now")) - .returning(CeleryUserRateLimit.last_scheduled_time) - ) - - _upsert_stmt = _insert_stmt.on_conflict_do_update( # type:ignore[attr-defined] - index_elements=["user_id"], set_=dict(last_scheduled_time=bindparam("sched_time")) - ) - def calculate_task_start_time( # type: ignore self, user_id: int, sa_session: galaxy_scoped_session, task_interval_secs: float, now: datetime.datetime ) -> datetime.datetime: with transaction(sa_session): - result = sa_session.execute( - self._update_stmt, {"userid": user_id, "interval": task_interval_secs, "now": now} + update_stmt = ( + update(CeleryUserRateLimit) + .where(CeleryUserRateLimit.user_id == user_id) + .values(last_scheduled_time=text("greatest(last_scheduled_time + ':interval second', " ":now) ")) + .returning(CeleryUserRateLimit.last_scheduled_time) ) - if result.rowcount == 0: + result = sa_session.execute(update_stmt, {"interval": task_interval_secs, "now": now}).all() + if not result: sched_time = now + datetime.timedelta(seconds=task_interval_secs) - result = sa_session.execute( - self._upsert_stmt, {"userid": user_id, "now": now, "sched_time": sched_time} + upsert_stmt = ( + ps_insert(CeleryUserRateLimit) # type:ignore[attr-defined] + .values(user_id=user_id, last_scheduled_time=now) + .returning(CeleryUserRateLimit.last_scheduled_time) + .on_conflict_do_update(index_elements=["user_id"], set_=dict(last_scheduled_time=sched_time)) ) - for row in result: - return row[0] + result = sa_session.execute(upsert_stmt).all() sa_session.commit() + return result[0][0] class GalaxyTaskBeforeStartUserRateLimitStandard(GalaxyTaskBeforeStartUserRateLimit):