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

rucio/core/heartbeat.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,26 +14,45 @@
 
 import datetime
 import hashlib
-from typing import TYPE_CHECKING
+from typing import TYPE_CHECKING, Optional
 
-from sqlalchemy import func
-from sqlalchemy.sql import distinct
+from sqlalchemy import and_, delete, func, select, update
 
 from rucio.common.exception import DatabaseException
 from rucio.common.utils import pid_exists
-from rucio.db.sqla.models import Heartbeats
+from rucio.db.sqla.models import Heartbeat
 from rucio.db.sqla.session import read_session, transactional_session
 
 if TYPE_CHECKING:
+    from threading import Thread
+    from typing import TypedDict
+
     from sqlalchemy.orm import Session
 
+    class HeartbeatDict(TypedDict):
+        readable: Optional[str]
+        hostname: str
+        pid: int
+        thread_name: Optional[str]
+        updated_at: datetime.datetime
+        created_at: datetime.datetime
+        payload: Optional[str]
+        test: int
 
 DEFAULT_EXPIRATION_DELAY = datetime.timedelta(days=1).total_seconds()
 
 
 @transactional_session
-def sanity_check(executable, hostname, hash_executable=None, pid=None, thread=None,
-                 expiration_delay=DEFAULT_EXPIRATION_DELAY, *, session: "Session"):
+def sanity_check(
+    executable: str,
+    hostname: str,
+    hash_executable: Optional[str] = None,
+    pid: Optional[int] = None,
+    thread: Optional["Thread"] = None,
+    expiration_delay: float = DEFAULT_EXPIRATION_DELAY,
+    *,
+    session: "Session"
+) -> None:
     """
     sanity_check wrapper to ignore DatabaseException errors.
 
@@ -57,7 +75,14 @@ def sanity_check(executable, hostname, hash_executable=None, pid=None, thread=No
 
 
 @transactional_session
-def _sanity_check(executable, hostname, hash_executable=None, expiration_delay=DEFAULT_EXPIRATION_DELAY, *, session: "Session"):
+def _sanity_check(
+    executable: str,
+    hostname: str,
+    hash_executable: Optional[str] = None,
+    expiration_delay: float = DEFAULT_EXPIRATION_DELAY,
+    *,
+    session: "Session"
+) -> None:
     """
     Check if processes on the host are still running.
 
@@ -67,24 +92,56 @@ def _sanity_check(executable, hostname, hash_executable=None, expiration_delay=D
     :param expiration_delay: time (in seconds) after which any inactive health check will be removed
     :param session: The database session in use.
     """
+    base_stmt = select(
+        Heartbeat.pid
+    ).distinct(
+    ).where(
+        Heartbeat.hostname == hostname
+    )
     if executable:
         if not hash_executable:
             hash_executable = calc_hash(executable)
 
-        for pid, in session.query(distinct(Heartbeats.pid)).filter_by(executable=hash_executable, hostname=hostname):
+        stmt = base_stmt.where(
+            Heartbeat.executable == hash_executable
+        )
+        for pid in session.execute(stmt).scalars().all():
             if not pid_exists(pid):
-                session.query(Heartbeats).filter_by(executable=hash_executable, hostname=hostname, pid=pid).delete()
+                stmt = delete(
+                    Heartbeat
+                ).where(
+                    and_(Heartbeat.executable == hash_executable,
+                         Heartbeat.hostname == hostname,
+                         Heartbeat.pid == pid)
+                )
+                session.execute(stmt)
     else:
-        for pid, in session.query(distinct(Heartbeats.pid)).filter_by(hostname=hostname):
+        for pid in session.execute(base_stmt).scalars().all():
             if not pid_exists(pid):
-                session.query(Heartbeats).filter_by(hostname=hostname, pid=pid).delete()
+                stmt = delete(
+                    Heartbeat
+                ).where(
+                    and_(Heartbeat.hostname == hostname,
+                         Heartbeat.pid == pid)
+                )
+                session.execute(stmt)
 
     if expiration_delay:
         cardiac_arrest(older_than=expiration_delay, session=session)
 
 
 @transactional_session
-def live(executable, hostname, pid, thread=None, older_than=600, hash_executable=None, payload=None, *, session: "Session"):
+def live(
+    executable: str,
+    hostname: str,
+    pid: int,
+    thread: Optional["Thread"] = None,
+    older_than: int = 600,
+    hash_executable: Optional[str] = None,
+    payload: Optional[str] = None,
+    *,
+    session: "Session"
+) -> dict[str, int]:
     """
     Register a heartbeat for a process/thread on a given node.
     The executable name is used for the calculation of thread assignments.
@@ -114,35 +171,48 @@ def live(executable, hostname, pid, thread=None, older_than=600, hash_executable
         thread_name = "thread"
 
     # upsert the heartbeat
-    rowcount = session.query(Heartbeats)\
-        .filter_by(executable=hash_executable,
-                   hostname=hostname,
-                   pid=pid,
-                   thread_id=thread_id)\
-        .update({'updated_at': datetime.datetime.utcnow(), 'payload': payload})
-    if not rowcount:
-        Heartbeats(executable=hash_executable,
-                   readable=executable[:Heartbeats.readable.property.columns[0].type.length],
-                   hostname=hostname,
-                   pid=pid,
-                   thread_id=thread_id,
-                   thread_name=thread_name,
-                   payload=payload).save(session=session)
+    stmt = update(
+        Heartbeat
+    ).where(
+        and_(Heartbeat.executable == hash_executable,
+             Heartbeat.hostname == hostname,
+             Heartbeat.pid == pid,
+             Heartbeat.thread_id == thread_id)
+    ).values({
+        Heartbeat.updated_at: datetime.datetime.utcnow(),
+        Heartbeat.payload: payload
+    })
+    if not session.execute(stmt).rowcount:
+        Heartbeat(executable=hash_executable,
+                  readable=executable[:Heartbeat.readable.property.columns[0].type.length],
+                  hostname=hostname,
+                  pid=pid,
+                  thread_id=thread_id,
+                  thread_name=thread_name,
+                  payload=payload).save(session=session)
 
     # assign thread identifier
-    query = session.query(Heartbeats.hostname,
-                          Heartbeats.pid,
-                          Heartbeats.thread_id)\
-                   .with_hint(Heartbeats, "index(HEARTBEATS HEARTBEATS_PK)", 'oracle')\
-                   .filter(Heartbeats.executable == hash_executable)\
-                   .filter(Heartbeats.updated_at >= datetime.datetime.utcnow() - datetime.timedelta(seconds=older_than))\
-                   .group_by(Heartbeats.hostname,
-                             Heartbeats.pid,
-                             Heartbeats.thread_id)\
-                   .order_by(Heartbeats.hostname,
-                             Heartbeats.pid,
-                             Heartbeats.thread_id)
-    result = query.all()
+    stmt = select(
+        Heartbeat.hostname,
+        Heartbeat.pid,
+        Heartbeat.thread_id
+    ).with_hint(
+        Heartbeat,
+        'INDEX(HEARTBEATS HEARTBEATS_PK)',
+        'oracle'
+    ).where(
+        and_(Heartbeat.executable == hash_executable,
+             Heartbeat.updated_at >= datetime.datetime.utcnow() - datetime.timedelta(seconds=older_than))
+    ).group_by(
+        Heartbeat.hostname,
+        Heartbeat.pid,
+        Heartbeat.thread_id
+    ).order_by(
+        Heartbeat.hostname,
+        Heartbeat.pid,
+        Heartbeat.thread_id
+    )
+    result = session.execute(stmt).all()
 
     # there is no universally applicable rownumber in SQLAlchemy
     # so we have to do it in Python
@@ -157,7 +227,16 @@ def live(executable, hostname, pid, thread=None, older_than=600, hash_executable
 
 
 @transactional_session
-def die(executable, hostname, pid, thread, older_than=None, hash_executable=None, *, session: "Session"):
+def die(
+    executable: str,
+    hostname: str,
+    pid: int,
+    thread: "Thread",
+    older_than: Optional[int] = None,
+    hash_executable: Optional[str] = None,
+    *,
+    session: "Session"
+) -> None:
     """
     Remove a single heartbeat older than specified.
 
@@ -168,25 +247,28 @@ def die(executable, hostname, pid, thread, older_than=None, hash_executable=None
     :param older_than: Removes specified heartbeats older than specified nr of seconds
     :param hash_executable: Hash of the executable.
     :param session: The database session in use.
-
-    :returns heartbeats: Dictionary {assign_thread, nr_threads}
     """
     if not hash_executable:
         hash_executable = calc_hash(executable)
 
-    query = session.query(Heartbeats).filter_by(executable=hash_executable,
-                                                hostname=hostname,
-                                                pid=pid,
-                                                thread_id=thread.ident)
+    stmt = delete(
+        Heartbeat
+    ).where(
+        and_(Heartbeat.executable == hash_executable,
+             Heartbeat.hostname == hostname,
+             Heartbeat.pid == pid,
+             Heartbeat.thread_id == thread.ident)
+    )
 
     if older_than:
-        query = query.filter(Heartbeats.updated_at < datetime.datetime.utcnow() - datetime.timedelta(seconds=older_than))
-
-    query.delete()
+        stmt = stmt.where(
+            Heartbeat.updated_at < datetime.datetime.utcnow() - datetime.timedelta(seconds=older_than)
+        )
+    session.execute(stmt)
 
 
 @transactional_session
-def cardiac_arrest(older_than=None, *, session: "Session"):
+def cardiac_arrest(older_than: Optional[int] = None, *, session: "Session") -> None:
     """
     Removes all heartbeats older than specified.
 
@@ -194,35 +276,42 @@ def cardiac_arrest(older_than=None, *, session: "Session"):
     :param session: The database session in use.
     """
 
-    query = session.query(Heartbeats)
+    stmt = delete(
+        Heartbeat
+    )
 
     if older_than:
-        query = query.filter(Heartbeats.updated_at < datetime.datetime.utcnow() - datetime.timedelta(seconds=older_than))
-
-    query.delete()
+        stmt = stmt.where(
+            Heartbeat.updated_at < datetime.datetime.utcnow() - datetime.timedelta(seconds=older_than)
+        )
+    session.execute(stmt)
 
 
 @read_session
-def list_heartbeats(*, session: "Session"):
+def list_heartbeats(*, session: "Session") -> list["HeartbeatDict"]:
     """
     List all heartbeats.
 
     :param session: The database session in use.
 
-    :returns: List of tuples
+    :returns: List of dicts
     """
 
-    query = session.query(Heartbeats.readable,
-                          Heartbeats.hostname,
-                          Heartbeats.pid,
-                          Heartbeats.thread_name,
-                          Heartbeats.updated_at,
-                          Heartbeats.created_at,
-                          Heartbeats.payload).order_by(Heartbeats.readable,
-                                                       Heartbeats.hostname,
-                                                       Heartbeats.thread_name)
-
-    result = query.all()
+    stmt = select(
+        Heartbeat.readable,
+        Heartbeat.hostname,
+        Heartbeat.pid,
+        Heartbeat.thread_name,
+        Heartbeat.updated_at,
+        Heartbeat.created_at,
+        Heartbeat.payload
+    ).order_by(
+        Heartbeat.readable,
+        Heartbeat.hostname,
+        Heartbeat.thread_name
+    )
+
+    result = session.execute(stmt).all()
     json_result = []
     for element in result:
         json_result.append(element._asdict())
@@ -230,7 +319,13 @@ def list_heartbeats(*, session: "Session"):
 
 
 @read_session
-def list_payload_counts(executable, older_than=600, hash_executable=None, *, session: "Session"):
+def list_payload_counts(
+    executable: str,
+    older_than: int = 600,
+    hash_executable: Optional[str] = None,
+    *,
+    session: "Session"
+) -> dict[str, int]:
     """
     Give the counts of number of threads per payload for a certain executable.
 
@@ -239,22 +334,27 @@ def list_payload_counts(executable, older_than=600, hash_executable=None, *, ses
     :param hash_executable: Hash of the executable.
     :param session: The database session in use.
 
-    :returns: List of tuples
+    :returns: Dict
     """
 
     if not hash_executable:
         hash_executable = calc_hash(executable)
-    query = session.query(Heartbeats.payload,
-                          func.count(Heartbeats.payload))\
-                   .filter(Heartbeats.executable == hash_executable)\
-                   .filter(Heartbeats.updated_at >= datetime.datetime.utcnow() - datetime.timedelta(seconds=older_than))\
-                   .group_by(Heartbeats.payload)\
-                   .order_by(Heartbeats.payload)
-
-    return dict((payload, count) for payload, count in query.all() if payload)
-
-
-def calc_hash(executable):
+    stmt = select(
+        Heartbeat.payload,
+        func.count(Heartbeat.payload)
+    ).where(
+        and_(Heartbeat.executable == hash_executable,
+             Heartbeat.updated_at >= datetime.datetime.utcnow() - datetime.timedelta(seconds=older_than))
+    ).group_by(
+        Heartbeat.payload
+    ).order_by(
+        Heartbeat.payload
+    )
+
+    return dict((payload, count) for payload, count in session.execute(stmt).all() if payload)
+
+
+def calc_hash(executable: str) -> str:
     """
     Calculates a SHA256 hash.
 

rucio/core/identity.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");
@@ -16,8 +15,7 @@
 import hashlib
 import os
 from re import match
-from typing import Optional
-from typing import TYPE_CHECKING
+from typing import TYPE_CHECKING, Optional
 
 from sqlalchemy import select, true
 from sqlalchemy.exc import IntegrityError
@@ -30,11 +28,14 @@ from rucio.db.sqla.constants import IdentityType
 from rucio.db.sqla.session import read_session, transactional_session
 
 if TYPE_CHECKING:
+    from collections.abc import Sequence
+
+    from sqlalchemy import Row
     from sqlalchemy.orm import Session
 
 
 @transactional_session
-def add_identity(identity: str, type_: IdentityType, email: str, password: Optional[str] = None, *, session: "Session"):
+def add_identity(identity: str, type_: IdentityType, email: str, password: Optional[str] = None, *, session: "Session") -> None:
     """
     Creates a user identity.
 
@@ -103,7 +104,7 @@ def verify_identity(identity: str, type_: IdentityType, password: Optional[str]
 
 
 @transactional_session
-def del_identity(identity: str, type_: IdentityType, *, session: "Session"):
+def del_identity(identity: str, type_: IdentityType, *, session: "Session") -> None:
     """
     Deletes a user identity.
 
@@ -125,7 +126,16 @@ def del_identity(identity: str, type_: IdentityType, *, session: "Session"):
 
 
 @transactional_session
-def add_account_identity(identity: str, type_: IdentityType, account: InternalAccount, email: str, default: bool = False, password: str = None, *, session: "Session"):
+def add_account_identity(
+    identity: str,
+    type_: IdentityType,
+    account: InternalAccount,
+    email: str,
+    default: bool = False,
+    password: Optional[str] = None,
+    *,
+    session: "Session"
+) -> None:
     """
     Adds a membership association between identity and account.
 
@@ -167,7 +177,7 @@ def add_account_identity(identity: str, type_: IdentityType, account: InternalAc
 
 
 @read_session
-def exist_identity_account(identity: str, type_: IdentityType, account: InternalAccount, *, session: "Session"):
+def exist_identity_account(identity: str, type_: IdentityType, account: InternalAccount, *, session: "Session") -> bool:
     """
     Check if an identity is mapped to an account.
 
@@ -189,13 +199,13 @@ def exist_identity_account(identity: str, type_: IdentityType, account: Internal
 
 
 @read_session
-def get_default_account(identity: str, type_: IdentityType, oldest_if_none: bool = False, *, session: "Session"):
+def get_default_account(identity: str, type_: IdentityType, oldest_if_none: bool = False, *, session: "Session") -> Optional[InternalAccount]:
     """
     Retrieves the default account mapped to an identity.
 
     :param identity: The identity key name. For example, x509DN, or a username.
     :param type_: The type of the authentication (x509, gss, userpass, saml, oidc).
-    :param oldest_if_none: If True and no default account it found the oldes known
+    :param oldest_if_none: If True and no default account it found the oldest known
                            account of that identity will be chosen, if False and
                            no default account is found, exception will be raised.
     :param session: The database session to use.
@@ -230,7 +240,7 @@ def get_default_account(identity: str, type_: IdentityType, oldest_if_none: bool
 
 
 @transactional_session
-def del_account_identity(identity: str, type_: IdentityType, account: InternalAccount, *, session: "Session"):
+def del_account_identity(identity: str, type_: IdentityType, account: InternalAccount, *, session: "Session") -> None:
     """
     Removes a membership association between identity and account.
 
@@ -253,7 +263,7 @@ def del_account_identity(identity: str, type_: IdentityType, account: InternalAc
 
 
 @read_session
-def list_identities(*, session: "Session", **kwargs):
+def list_identities(*, session: "Session", **kwargs) -> "Sequence[Row[tuple[str, IdentityType]]]":
     """
     Returns a list of all identities.
 
@@ -271,7 +281,7 @@ def list_identities(*, session: "Session", **kwargs):
 
 
 @read_session
-def list_accounts_for_identity(identity: str, type_: IdentityType, *, session: "Session"):
+def list_accounts_for_identity(identity: str, type_: IdentityType, *, session: "Session") -> "Sequence[InternalAccount]":
     """
     Returns a list of all accounts for an identity.
 

rucio/core/importer.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");
@@ -13,14 +12,21 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-from typing import TYPE_CHECKING
+from collections.abc import Iterable
+from typing import TYPE_CHECKING, Any
+
+from sqlalchemy import select
 
 from rucio.common.config import config_get
+from rucio.common.constants import RseAttr
 from rucio.common.exception import RSEOperationNotSupported
 from rucio.common.types import InternalAccount
-from rucio.core import rse as rse_module, distance as distance_module, account as account_module, identity as identity_module
+from rucio.core import account as account_module
+from rucio.core import distance as distance_module
+from rucio.core import identity as identity_module
+from rucio.core import rse as rse_module
 from rucio.db.sqla import models
-from rucio.db.sqla.constants import RSEType, AccountType, IdentityType
+from rucio.db.sqla.constants import AccountType, IdentityType, RSEType
 from rucio.db.sqla.session import transactional_session
 
 if TYPE_CHECKING:
@@ -28,7 +34,7 @@ if TYPE_CHECKING:
 
 
 @transactional_session
-def import_rses(rses, rse_sync_method='edit', attr_sync_method='edit', protocol_sync_method='edit', vo='def', *, session: "Session"):
+def import_rses(rses: dict[str, dict[str, Any]], rse_sync_method: str = 'edit', attr_sync_method: str = 'edit', protocol_sync_method: str = 'edit', vo: str = 'def', *, session: "Session") -> None:
     new_rses = []
     for rse_name in rses:
         rse = rses[rse_name]
@@ -104,8 +110,8 @@ def import_rses(rses, rse_sync_method='edit', attr_sync_method='edit', protocol_
 
         # Attributes
         attributes = rse.get('attributes', {})
-        attributes['lfn2pfn_algorithm'] = rse.get('lfn2pfn_algorithm')
-        attributes['verify_checksum'] = rse.get('verify_checksum')
+        attributes[RseAttr.LFN2PFN_ALGORITHM] = rse.get('lfn2pfn_algorithm')
+        attributes[RseAttr.VERIFY_CHECKSUM] = rse.get('verify_checksum')
 
         old_attributes = rse_module.list_rse_attributes(rse_id=rse_id, session=session)
         missing_attributes = [attribute for attribute in old_attributes if attribute not in attributes]
@@ -136,7 +142,7 @@ def import_rses(rses, rse_sync_method='edit', attr_sync_method='edit', protocol_
 
 
 @transactional_session
-def import_distances(distances, vo='def', *, session: "Session"):
+def import_distances(distances, vo: str = 'def', *, session: "Session") -> None:
     for src_rse_name in distances:
         src = rse_module.get_rse_id(rse=src_rse_name, vo=vo, session=session)
         for dest_rse_name in distances[src_rse_name]:
@@ -156,7 +162,7 @@ def import_distances(distances, vo='def', *, session: "Session"):
 
 
 @transactional_session
-def import_identities(identities, account_name, old_identities, old_identity_account, account_email, *, session: "Session"):
+def import_identities(identities: Iterable[dict[str, Any]], account_name: str, old_identities: Iterable[tuple], old_identity_account: tuple[str, str, str], account_email: str, *, session: "Session") -> None:
     for identity in identities:
         identity['type'] = IdentityType[identity['type'].upper()]
 
@@ -185,14 +191,19 @@ def import_identities(identities, account_name, old_identities, old_identity_acc
 
 
 @transactional_session
-def import_accounts(accounts, vo='def', *, session: "Session"):
+def import_accounts(accounts: Iterable[dict[str, Any]], vo: str = 'def', *, session: "Session") -> None:
     vo_filter = {'account': InternalAccount(account='*', vo=vo)}
     old_accounts = {account['account']: account for account in account_module.list_accounts(filter_=vo_filter, session=session)}
     missing_accounts = [account for account in accounts if account['account'] not in old_accounts]
     outdated_accounts = [account for account in accounts if account['account'] in old_accounts]
     to_be_removed_accounts = [old_account for old_account in old_accounts if old_account not in [account['account'] for account in accounts]]
     old_identities = identity_module.list_identities(session=session)
-    old_identity_account = session.query(models.IdentityAccountAssociation.identity, models.IdentityAccountAssociation.identity_type, models.IdentityAccountAssociation.account).all()
+    stmt = select(
+        models.IdentityAccountAssociation.identity,
+        models.IdentityAccountAssociation.identity_type,
+        models.IdentityAccountAssociation.account
+    )
+    old_identity_account = session.execute(stmt).all()
 
     # add missing accounts
     for account_dict in missing_accounts:
@@ -222,7 +233,7 @@ def import_accounts(accounts, vo='def', *, session: "Session"):
 
 
 @transactional_session
-def import_data(data, vo='def', *, session: "Session"):
+def import_data(data: dict[str, Any], vo: str = 'def', *, session: "Session") -> None:
     """
     Import data to add and update records in Rucio.