Skip to content

Commit

Permalink
Auth: Migrate to SQLAlchemy 2.0 syntax; #6057
Browse files Browse the repository at this point in the history
  • Loading branch information
rdimaio committed Feb 19, 2024
1 parent 52de58a commit 0e4e829
Show file tree
Hide file tree
Showing 3 changed files with 16 additions and 12 deletions.
11 changes: 7 additions & 4 deletions lib/rucio/core/account_counter.py
Expand Up @@ -84,7 +84,7 @@ def del_counter(rse_id: str, account: InternalAccount, *, session: "Session") ->


@read_session
def get_updated_account_counters(total_workers: int, worker_number: int, *, session: "Session") -> list[tuple[InternalAccount, str]]:
def get_updated_account_counters(total_workers: int, worker_number: int, *, session: "Session") -> list[dict[str, any]]:
"""
Get updated rse_counters.
Expand All @@ -93,17 +93,20 @@ 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)

query = (select(models.UpdatedAccountCounter.account, models.UpdatedAccountCounter.rse_id)
.distinct()
)

if session.bind.dialect.name == 'oracle':
hash_variable = 'CONCAT(account, rse_id)'''
else:
hash_variable = 'concat(account, rse_id)'

query = filter_thread_work(session=session, query=query, total_threads=total_workers, thread_id=worker_number, hash_variable=hash_variable)
query_result = session.execute(query)

return query.all()
return [row._asdict() for row in query_result.all()]


@transactional_session
Expand Down
3 changes: 2 additions & 1 deletion lib/rucio/core/oidc.py
Expand Up @@ -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}
return {'token': token.token, 'expires_at': token.expired_at}
14 changes: 7 additions & 7 deletions lib/rucio/daemons/abacus/account.py
Expand Up @@ -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:
Expand Down

0 comments on commit 0e4e829

Please sign in to comment.