rucio 32.8.6__py3-none-any.whl → 35.8.0__py3-none-any.whl

rucio/core/__init__.py

@@ -1,4 +1,3 @@
-# -*- coding: utf-8 -*-
 # Copyright European Organization for Nuclear Research (CERN) since 2012
 #
 # Licensed under the Apache License, Version 2.0 (the "License");

rucio/core/account.py

@@ -1,4 +1,3 @@
-# -*- coding: utf-8 -*-
 # Copyright European Organization for Nuclear Research (CERN) since 2012
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
@@ -17,9 +16,9 @@ from datetime import datetime
 from enum import Enum
 from re import match
 from traceback import format_exc
-from typing import TYPE_CHECKING
+from typing import TYPE_CHECKING, Any, Optional
 
-from sqlalchemy import select, and_
+from sqlalchemy import and_, select
 from sqlalchemy.exc import IntegrityError
 from sqlalchemy.orm import exc
 
@@ -30,14 +29,24 @@ from rucio.common.config import config_get_bool
 from rucio.core.vo import vo_exists
 from rucio.db.sqla import models
 from rucio.db.sqla.constants import AccountStatus, AccountType
-from rucio.db.sqla.session import read_session, transactional_session, stream_session
+from rucio.db.sqla.session import read_session, stream_session, transactional_session
 
 if TYPE_CHECKING:
+    from collections.abc import Iterator, Mapping
+
     from sqlalchemy.orm import Session
 
+    from rucio.common.types import AccountAttributesDict, AccountDict, AccountUsageModelDict, IdentityDict, InternalAccount, UsageDict
+
 
 @transactional_session
-def add_account(account, type_, email, *, session: "Session"):
+def add_account(
+    account: "InternalAccount",
+    type_: AccountType,
+    email: str,
+    *,
+    session: "Session"
+) -> None:
     """ Add an account with the given account name and type.
 
     :param account: the name of the new account.
@@ -63,7 +72,11 @@ def add_account(account, type_, email, *, session: "Session"):
 
 
 @read_session
-def account_exists(account, *, session: "Session"):
+def account_exists(
+    account: "InternalAccount",
+    *,
+    session: "Session"
+) -> bool:
     """ Checks to see if account exists and is active.
 
     :param account: Name of the account.
@@ -81,7 +94,11 @@ def account_exists(account, *, session: "Session"):
 
 
 @read_session
-def get_account(account, *, session: "Session"):
+def get_account(
+    account: "InternalAccount",
+    *,
+    session: "Session"
+) -> models.Account:
     """ Returns an account for the given account name.
 
     :param account: the name of the account.
@@ -102,7 +119,11 @@ def get_account(account, *, session: "Session"):
 
 
 @transactional_session
-def del_account(account, *, session: "Session"):
+def del_account(
+    account: "InternalAccount",
+    *,
+    session: "Session"
+) -> None:
     """ Disable an account with the given account name.
 
     :param account: the account name.
@@ -115,15 +136,21 @@ def del_account(account, *, session: "Session"):
         models.Account.status == AccountStatus.ACTIVE
     )
     try:
-        account = session.execute(query).scalar_one()
+        query_result = session.execute(query).scalar_one()
     except exc.NoResultFound:
         raise exception.AccountNotFound('Account with ID \'%s\' cannot be found' % account)
 
-    account.update({'status': AccountStatus.DELETED, 'deleted_at': datetime.utcnow()})
+    query_result.update({'status': AccountStatus.DELETED, 'deleted_at': datetime.utcnow()})
 
 
 @transactional_session
-def update_account(account, key, value, *, session: "Session"):
+def update_account(
+    account: "InternalAccount",
+    key: str,
+    value: Any,
+    *,
+    session: "Session"
+) -> None:
     """ Update a property of an account.
 
     :param account: Name of the account.
@@ -137,22 +164,26 @@ def update_account(account, key, value, *, session: "Session"):
         models.Account.account == account
     )
     try:
-        account = session.execute(query).scalar_one()
+        query_result = session.execute(query).scalar_one()
     except exc.NoResultFound:
         raise exception.AccountNotFound('Account with ID \'%s\' cannot be found' % account)
     if key == 'status':
         if isinstance(value, str):
             value = AccountStatus[value]
         if value == AccountStatus.SUSPENDED:
-            account.update({'status': value, 'suspended_at': datetime.utcnow()})
+            query_result.update({'status': value, 'suspended_at': datetime.utcnow()})
         elif value == AccountStatus.ACTIVE:
-            account.update({'status': value, 'suspended_at': None})
+            query_result.update({'status': value, 'suspended_at': None})
     else:
-        account.update({key: value})
+        query_result.update({key: value})
 
 
 @stream_session
-def list_accounts(filter_=None, *, session: "Session"):
+def list_accounts(
+    filter_: Optional["Mapping[str, Any]"] = None,
+    *,
+    session: "Session"
+) -> "Iterator[AccountDict]":
     """ Returns a list of all account names.
 
     :param filter_: Dictionary of attributes by which the input data should be filtered
@@ -213,7 +244,11 @@ def list_accounts(filter_=None, *, session: "Session"):
 
 
 @read_session
-def list_identities(account, *, session: "Session"):
+def list_identities(
+    account: "InternalAccount",
+    *,
+    session: "Session"
+) -> list["IdentityDict"]:
     """
     List all identities on an account.
 
@@ -244,11 +279,16 @@ def list_identities(account, *, session: "Session"):
     ).where(
         models.IdentityAccountAssociation.account == account
     )
-    return [row._asdict() for row in session.execute(query)]
+
+    return [row._asdict() for row in session.execute(query)]  # type: ignore (pending SQLA2.1: https://github.com/rucio/rucio/discussions/6615)
 
 
 @read_session
-def list_account_attributes(account, *, session: "Session"):
+def list_account_attributes(
+    account: "InternalAccount",
+    *,
+    session: "Session"
+) -> list["AccountAttributesDict"]:
     """
     Get all attributes defined for an account.
 
@@ -274,11 +314,16 @@ def list_account_attributes(account, *, session: "Session"):
     ).where(
         models.AccountAttrAssociation.account == account
     )
-    return [row._asdict() for row in session.execute(query)]
+    return [row._asdict() for row in session.execute(query)]  # type: ignore (pending SQLA2.1: https://github.com/rucio/rucio/discussions/6615)
 
 
 @read_session
-def has_account_attribute(account, key, *, session: "Session"):
+def has_account_attribute(
+    account: "InternalAccount",
+    key: str,
+    *,
+    session: "Session"
+) -> bool:
     """
     Indicates whether the named key is present for the account.
 
@@ -298,7 +343,13 @@ def has_account_attribute(account, key, *, session: "Session"):
 
 
 @transactional_session
-def add_account_attribute(account, key, value, *, session: "Session"):
+def add_account_attribute(
+    account: "InternalAccount",
+    key: str,
+    value: Any,
+    *,
+    session: "Session"
+) -> None:
     """
     Add an attribute for the given account name.
 
@@ -334,7 +385,12 @@ def add_account_attribute(account, key, value, *, session: "Session"):
 
 
 @transactional_session
-def del_account_attribute(account, key, *, session: "Session"):
+def del_account_attribute(
+    account: "InternalAccount",
+    key: str,
+    *,
+    session: "Session"
+) -> None:
     """
     Add an attribute for the given account name.
 
@@ -355,7 +411,12 @@ def del_account_attribute(account, key, *, session: "Session"):
 
 
 @read_session
-def get_usage(rse_id, account, *, session: "Session"):
+def get_usage(
+    rse_id: str,
+    account: "InternalAccount",
+    *,
+    session: "Session"
+) -> "UsageDict":
     """
     Returns current values of the specified counter, or raises CounterNotFound if the counter does not exist.
 
@@ -373,13 +434,17 @@ def get_usage(rse_id, account, *, session: "Session"):
         models.AccountUsage.account == account
     )
     try:
-        return session.execute(query).one()._asdict()
+        return session.execute(query).one()._asdict()  # type: ignore (pending SQLA2.1: https://github.com/rucio/rucio/discussions/6615)
     except exc.NoResultFound:
         return {'bytes': 0, 'files': 0, 'updated_at': None}
 
 
 @read_session
-def get_all_rse_usages_per_account(account, *, session: "Session"):
+def get_all_rse_usages_per_account(
+    account: "InternalAccount",
+    *,
+    session: "Session"
+) -> list["AccountUsageModelDict"]:
     """
     Returns current values of the specified counter, or raises CounterNotFound if the counter does not exist.
 
@@ -394,13 +459,18 @@ def get_all_rse_usages_per_account(account, *, session: "Session"):
         models.AccountUsage.account == account
     )
     try:
-        return [result.to_dict() for result in session.execute(query).scalars()]
+        return [result.to_dict() for result in session.execute(query).scalars()]  # type: ignore (pending SQLA2.1: https://github.com/rucio/rucio/discussions/6615)
     except exc.NoResultFound:
         return []
 
 
 @read_session
-def get_usage_history(rse_id, account, *, session: "Session"):
+def get_usage_history(
+    rse_id: str,
+    account: "InternalAccount",
+    *,
+    session: "Session"
+) -> list["UsageDict"]:
     """
     Returns historical values of the specified counter, or raises CounterNotFound if the counter does not exist.
 
@@ -420,7 +490,7 @@ def get_usage_history(rse_id, account, *, session: "Session"):
         models.AccountUsageHistory.updated_at
     )
     try:
-        return [row._asdict() for row in session.execute(query)]
+        return [row._asdict() for row in session.execute(query)]  # type: ignore (pending SQLA2.1: https://github.com/rucio/rucio/discussions/6615)
     except exc.NoResultFound:
         raise exception.CounterNotFound('No usage can be found for account %s on RSE %s' % (account, rucio.core.rse.get_rse_name(rse_id=rse_id, session=session)))
     return []

rucio/core/account_counter.py

@@ -1,4 +1,3 @@
-# -*- coding: utf-8 -*-
 # Copyright European Organization for Nuclear Research (CERN) since 2012
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
@@ -15,20 +14,27 @@
 import datetime
 from typing import TYPE_CHECKING
 
-from sqlalchemy import literal, insert, select
-from sqlalchemy.orm.exc import NoResultFound
+from sqlalchemy import and_, delete, insert, literal, select
+from sqlalchemy.exc import NoResultFound
 
-from rucio.db.sqla import models, filter_thread_work
+from rucio.db.sqla import filter_thread_work, models
 from rucio.db.sqla.session import read_session, transactional_session
 
 if TYPE_CHECKING:
     from sqlalchemy.orm import Session
 
+    from rucio.common.types import InternalAccount, RSEAccountCounterDict
+
 MAX_COUNTERS = 10
 
 
 @transactional_session
-def add_counter(rse_id, account, *, session: "Session"):
+def add_counter(
+    rse_id: str,
+    account: "InternalAccount",
+    *,
+    session: "Session"
+) -> None:
     """
     Creates the specified counter for a rse_id and account.
 
@@ -41,7 +47,14 @@ def add_counter(rse_id, account, *, session: "Session"):
 
 
 @transactional_session
-def increase(rse_id, account, files, bytes_, *, session: "Session"):
+def increase(
+    rse_id: str,
+    account: "InternalAccount",
+    files: int,
+    bytes_: int,
+    *,
+    session: "Session"
+) -> None:
     """
     Increments the specified counter by the specified amount.
 
@@ -55,7 +68,14 @@ def increase(rse_id, account, files, bytes_, *, session: "Session"):
 
 
 @transactional_session
-def decrease(rse_id, account, files, bytes_, *, session: "Session"):
+def decrease(
+    rse_id: str,
+    account: "InternalAccount",
+    files: int,
+    bytes_: int,
+    *,
+    session: "Session"
+) -> None:
     """
     Decreases the specified counter by the specified amount.
 
@@ -69,7 +89,12 @@ def decrease(rse_id, account, files, bytes_, *, session: "Session"):
 
 
 @transactional_session
-def del_counter(rse_id, account, *, session: "Session"):
+def del_counter(
+    rse_id: str,
+    account: "InternalAccount",
+    *,
+    session: "Session"
+) -> None:
     """
     Resets the specified counter and initializes it by the specified amounts.
 
@@ -78,11 +103,24 @@ def del_counter(rse_id, account, *, session: "Session"):
     :param session: The database session in use.
     """
 
-    session.query(models.AccountUsage).filter_by(rse_id=rse_id, account=account).delete(synchronize_session=False)
+    stmt = delete(
+        models.AccountUsage
+    ).where(
+        and_(models.AccountUsage.rse_id == rse_id,
+             models.AccountUsage.account == account)
+    ).execution_options(
+        synchronize_session=False
+    )
+    session.execute(stmt)
 
 
 @read_session
-def get_updated_account_counters(total_workers, worker_number, *, session: "Session"):
+def get_updated_account_counters(
+    total_workers: int,
+    worker_number: int,
+    *,
+    session: "Session"
+) -> list["RSEAccountCounterDict"]:
     """
     Get updated rse_counters.
 
@@ -91,21 +129,25 @@ def get_updated_account_counters(total_workers, worker_number, *, session: "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)')
 
-    return query.all()
+    return [row._asdict() for row in session.execute(query).all()]  # type: ignore (pending SQLA2.1: https://github.com/rucio/rucio/discussions/6615)
 
 
 @transactional_session
-def update_account_counter(account, rse_id, *, session: "Session"):
+def update_account_counter(
+    account: "InternalAccount",
+    rse_id: str,
+    *,
+    session: "Session"
+) -> None:
     """
     Read the updated_account_counters and update the account_counter.
 
@@ -114,10 +156,22 @@ def update_account_counter(account, rse_id, *, session: "Session"):
     :param session:  Database session in use.
     """
 
-    updated_account_counters = session.query(models.UpdatedAccountCounter).filter_by(account=account, rse_id=rse_id).all()
+    stmt = select(
+        models.UpdatedAccountCounter
+    ).where(
+        and_(models.UpdatedAccountCounter.account == account,
+             models.UpdatedAccountCounter.rse_id == rse_id)
+    )
+    updated_account_counters = session.execute(stmt).scalars().all()
 
     try:
-        account_counter = session.query(models.AccountUsage).filter_by(account=account, rse_id=rse_id).one()
+        stmt = select(
+            models.AccountUsage
+        ).where(
+            and_(models.AccountUsage.account == account,
+                 models.AccountUsage.rse_id == rse_id)
+        )
+        account_counter = session.execute(stmt).scalar_one()
         account_counter.bytes += sum([updated_account_counter.bytes for updated_account_counter in updated_account_counters])
         account_counter.files += sum([updated_account_counter.files for updated_account_counter in updated_account_counters])
     except NoResultFound:
@@ -131,7 +185,12 @@ def update_account_counter(account, rse_id, *, session: "Session"):
 
 
 @transactional_session
-def update_account_counter_history(account, rse_id, *, session: "Session"):
+def update_account_counter_history(
+    account: "InternalAccount",
+    rse_id: str,
+    *,
+    session: "Session"
+) -> None:
     """
     Read the AccountUsage and update the AccountUsageHistory.
 
@@ -139,7 +198,13 @@ def update_account_counter_history(account, rse_id, *, session: "Session"):
     :param rse_id:   The rse_id to update.
     :param session:  Database session in use.
     """
-    counter = session.query(models.AccountUsage).filter_by(rse_id=rse_id, account=account).one_or_none()
+    stmt = select(
+        models.AccountUsage
+    ).where(
+        and_(models.AccountUsage.account == account,
+             models.AccountUsage.rse_id == rse_id)
+    )
+    counter = session.execute(stmt).scalar_one_or_none()
     if counter:
         models.AccountUsageHistory(rse_id=rse_id, account=account, files=counter.files, bytes=counter.bytes).save(session=session)
     else:
@@ -147,7 +212,7 @@ def update_account_counter_history(account, rse_id, *, session: "Session"):
 
 
 @transactional_session
-def fill_account_counter_history_table(*, session: "Session"):
+def fill_account_counter_history_table(*, session: "Session") -> None:
     """
     Make a snapshot of current counters