From 132ca39607e51fab210c32d4af506b9082b94619 Mon Sep 17 00:00:00 2001 From: rdimaio Date: Mon, 19 Feb 2024 19:35:49 +0100 Subject: [PATCH] Auth: Migrate to SQLAlchemy 2.0 syntax; #6057 --- lib/rucio/core/account_counter.py | 14 ++++++-------- lib/rucio/core/oidc.py | 3 ++- lib/rucio/daemons/abacus/account.py | 14 +++++++------- 3 files changed, 15 insertions(+), 16 deletions(-) diff --git a/lib/rucio/core/account_counter.py b/lib/rucio/core/account_counter.py index bc890e22855..85e7e1e1c0b 100644 --- a/lib/rucio/core/account_counter.py +++ b/lib/rucio/core/account_counter.py @@ -93,17 +93,15 @@ def get_updated_account_counters(total_workers: int, worker_number: int, *, sess :param session: Database session in use. :returns: List of rse_ids whose rse_counters need to be updated. """ - query = session.query(models.UpdatedAccountCounter.account, models.UpdatedAccountCounter.rse_id).\ - distinct(models.UpdatedAccountCounter.account, models.UpdatedAccountCounter.rse_id) - if session.bind.dialect.name == 'oracle': - hash_variable = 'CONCAT(account, rse_id)''' - else: - hash_variable = 'concat(account, rse_id)' + query = (select(models.UpdatedAccountCounter.account, models.UpdatedAccountCounter.rse_id) + .distinct() + ) - query = filter_thread_work(session=session, query=query, total_threads=total_workers, thread_id=worker_number, hash_variable=hash_variable) + query = filter_thread_work(session=session, query=query, total_threads=total_workers, thread_id=worker_number, hash_variable='CONCAT(account, rse_id)') + query_result = session.execute(query) - return query.all() + return [row._asdict() for row in query_result.all()] @transactional_session diff --git a/lib/rucio/core/oidc.py b/lib/rucio/core/oidc.py index d3276483e4d..38b7cf6d941 100644 --- a/lib/rucio/core/oidc.py +++ b/lib/rucio/core/oidc.py @@ -1443,5 +1443,6 @@ def oidc_identity_string(sub: str, iss: str) -> str: """ return 'SUB=' + str(sub) + ', ISS=' + str(iss) + def token_dictionary(token: models.Token) -> TokenDict: - return {'token': token.token, 'expires_at': token.expired_at} \ No newline at end of file + return {'token': token.token, 'expires_at': token.expired_at} diff --git a/lib/rucio/daemons/abacus/account.py b/lib/rucio/daemons/abacus/account.py index 78de5d45c42..273111e6d9d 100644 --- a/lib/rucio/daemons/abacus/account.py +++ b/lib/rucio/daemons/abacus/account.py @@ -55,22 +55,22 @@ def run_once(heartbeat_handler, **_kwargs): worker_number, total_workers, logger = heartbeat_handler.live() start = time.time() # NOQA - account_rse_ids = get_updated_account_counters(total_workers=total_workers, - worker_number=worker_number) - logger(logging.DEBUG, 'Index query time %f size=%d' % (time.time() - start, len(account_rse_ids))) + updated_account_counters = get_updated_account_counters(total_workers=total_workers, + worker_number=worker_number) + logger(logging.DEBUG, 'Index query time %f size=%d' % (time.time() - start, len(updated_account_counters))) # If the list is empty, sent the worker to sleep - if not account_rse_ids: + if not updated_account_counters: logger(logging.INFO, 'did not get any work') return - for account_rse_id in account_rse_ids: + for account_counter in updated_account_counters: worker_number, total_workers, logger = heartbeat_handler.live() if graceful_stop.is_set(): break start_time = time.time() - update_account_counter(account=account_rse_id[0], rse_id=account_rse_id[1]) - logger(logging.DEBUG, 'update of account-rse counter "%s-%s" took %f' % (account_rse_id[0], account_rse_id[1], time.time() - start_time)) + update_account_counter(account=account_counter['account'], rse_id=account_counter['rse_id']) + logger(logging.DEBUG, 'update of account-rse counter "%s-%s" took %f' % (account_counter['account'], account_counter['rse_id'], time.time() - start_time)) def stop(signum: "Optional[int]" = None, frame: "Optional[FrameType]" = None) -> None: