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

rucio/core/message.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,11 +15,11 @@
 import json
 from typing import TYPE_CHECKING
 
-from sqlalchemy import or_, delete, update, insert
+from sqlalchemy import delete, insert, or_, select, update
 from sqlalchemy.exc import IntegrityError
 
 from rucio.common.config import config_get_list
-from rucio.common.constants import HermesService, MAX_MESSAGE_LENGTH
+from rucio.common.constants import MAX_MESSAGE_LENGTH, HermesService
 from rucio.common.exception import InvalidObject, RucioException
 from rucio.common.utils import APIEncoder, chunks
 from rucio.db.sqla import filter_thread_work
@@ -29,6 +28,7 @@ from rucio.db.sqla.session import transactional_session
 
 if TYPE_CHECKING:
     from typing import Any, Optional
+
     from sqlalchemy.orm import Session
 
     MessageType = dict[str, Any]
@@ -74,7 +74,10 @@ def add_messages(messages: "MessagesListType", *, session: "Session") -> None:
             except TypeError as err:  # noqa: F841
                 raise InvalidObject(f'Invalid JSON for payload: {err}')
     for messages_chunk in chunks(msgs, 1000):
-        session.execute(insert(Message), messages_chunk)
+        stmt = insert(
+            Message
+        )
+        session.execute(stmt, messages_chunk)
 
 
 @transactional_session
@@ -114,45 +117,58 @@ def retrieve_messages(bulk: int = 1000,
     """
     messages = []
     try:
-        subquery = session.query(Message.id)
-        subquery = filter_thread_work(session=session, query=subquery, total_threads=total_threads, thread_id=thread)
+        stmt_subquery = select(
+            Message.id
+        ).order_by(
+            Message.created_at
+        )
+        stmt_subquery = filter_thread_work(session=session, query=stmt_subquery, total_threads=total_threads, thread_id=thread)
         if event_type:
-            subquery = subquery.filter_by(event_type=event_type)
+            stmt_subquery = stmt_subquery.where(
+                Message.event_type == event_type
+            )
         elif old_mode:
-            subquery = subquery.filter(Message.event_type != 'email')
+            stmt_subquery = stmt_subquery.where(
+                Message.event_type != 'email'
+            )
 
         # Step 1:
         # MySQL does not support limits in nested queries, limit on the outer query instead.
         # This is not as performant, but the best we can get from MySQL.
         # FIXME: SQLAlchemy generates wrong nowait MySQL8 statement for MySQL5
         #        Remove once this is resolved in SQLAlchemy
+        stmt = select(
+            Message.id,
+            Message.created_at,
+            Message.event_type,
+            Message.payload,
+            Message.services
+        )
         if session.bind.dialect.name == 'mysql':
-            subquery = subquery.order_by(Message.created_at)
-            query = session.query(Message.id,
-                                  Message.created_at,
-                                  Message.event_type,
-                                  Message.payload,
-                                  Message.services)\
-                           .filter(Message.id.in_(subquery))
+            stmt = stmt.where(
+                Message.id.in_(stmt_subquery)
+            )
         else:
-            subquery = subquery.order_by(Message.created_at).limit(bulk)
-            query = session.query(Message.id,
-                                  Message.created_at,
-                                  Message.event_type,
-                                  Message.payload,
-                                  Message.services)\
-                           .filter(Message.id.in_(subquery))\
-                           .with_for_update(nowait=True)
+            stmt_subquery = stmt_subquery.limit(
+                bulk
+            )
+            stmt = stmt.where(
+                Message.id.in_(stmt_subquery)
+            ).with_for_update(
+                nowait=True
+            )
 
         # Step 2:
         # MySQL does not support limits in nested queries, limit on the outer query instead.
         # This is not as performant, but the best we can get from MySQL.
         if session.bind.dialect.name == 'mysql':
-            query = query.limit(bulk)
+            stmt = stmt.limit(
+                bulk
+            )
 
         # Step 3:
         # Assemble message object
-        for id_, created_at, event_type, payload, services in query:
+        for id_, created_at, event_type, payload, services in session.execute(stmt).all():
             message = {'id': id_,
                        'created_at': created_at,
                        'event_type': event_type,
@@ -160,8 +176,12 @@ def retrieve_messages(bulk: int = 1000,
 
             # Only switch SQL context when necessary
             if payload == 'nolimit':
-                nolimit_query = session.query(Message.payload_nolimit).filter(Message.id == id_).one()[0]
-                message['payload'] = json.loads(str(nolimit_query))
+                nolimit_stmt = select(
+                    Message.payload_nolimit
+                ).where(
+                    Message.id == id_
+                )
+                message['payload'] = json.loads(str(session.execute(nolimit_stmt).scalar_one()))
             else:
                 message['payload'] = json.loads(str(payload))
 
@@ -188,13 +208,22 @@ def delete_messages(messages: "MessagesListType", *, session: "Session") -> None
 
     try:
         if message_condition:
-            stmt = delete(Message).\
-                prefix_with("/*+ index(messages MESSAGES_ID_PK) */", dialect='oracle').\
-                where(or_(*message_condition)).\
-                execution_options(synchronize_session=False)
+            stmt = delete(
+                Message
+            ).prefix_with(
+                '/*+ INDEX(messages MESSAGES_ID_PK) */',
+                dialect='oracle'
+            ).where(
+                or_(*message_condition)
+            ).execution_options(
+                synchronize_session=False
+            )
             session.execute(stmt)
 
-            session.execute(insert(MessageHistory), messages)
+            stmt = insert(
+                MessageHistory
+            )
+            session.execute(stmt, messages)
     except IntegrityError as e:
         raise RucioException(e.args)
 
@@ -208,7 +237,12 @@ def truncate_messages(*, session: "Session") -> None:
     """
 
     try:
-        session.query(Message).delete(synchronize_session=False)
+        stmt = delete(
+            Message
+        ).execution_options(
+            synchronize_session=False
+        )
+        session.execute(stmt)
     except IntegrityError as e:
         raise RucioException(e.args)
 
@@ -230,11 +264,18 @@ def update_messages_services(messages: "MessagesListType", services: str, *, ses
 
     try:
         if message_condition:
-            stmt = update(Message).\
-                prefix_with("/*+ index(messages MESSAGES_ID_PK) */", dialect='oracle').\
-                where(or_(*message_condition)).\
-                execution_options(synchronize_session=False).\
-                values(services=services)
+            stmt = update(
+                Message
+            ).prefix_with(
+                '/*+ INDEX(messages MESSAGES_ID_PK) */',
+                dialect='oracle'
+            ).where(
+                or_(*message_condition)
+            ).execution_options(
+                synchronize_session=False
+            ).values({
+                Message.services: services
+            })
             session.execute(stmt)
 
     except IntegrityError as err:

rucio/core/{meta.py → meta_conventions.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");
@@ -14,15 +13,13 @@
 # limitations under the License.
 
 from re import match
-from typing import TYPE_CHECKING
+from typing import TYPE_CHECKING, Optional, Union
 
-from sqlalchemy.exc import IntegrityError
-from sqlalchemy.orm.exc import NoResultFound
+from sqlalchemy import and_, delete, select
+from sqlalchemy.exc import IntegrityError, NoResultFound
 
 from rucio.common.constraints import AUTHORIZED_VALUE_TYPES
-from rucio.common.exception import (Duplicate, RucioException,
-                                    KeyNotFound, InvalidValueForKey, UnsupportedValueType,
-                                    InvalidObject, UnsupportedKeyType)
+from rucio.common.exception import Duplicate, InvalidObject, InvalidValueForKey, KeyNotFound, RucioException, UnsupportedKeyType, UnsupportedValueType
 from rucio.db.sqla import models
 from rucio.db.sqla.constants import DIDType, KeyType
 from rucio.db.sqla.session import read_session, transactional_session
@@ -32,9 +29,9 @@ if TYPE_CHECKING:
 
 
 @transactional_session
-def add_key(key, key_type, value_type=None, value_regexp=None, *, session: "Session"):
+def add_key(key: str, key_type: Union[KeyType, str], value_type: Optional[str] = None, value_regexp: Optional[str] = None, *, session: "Session") -> None:
     """
-    Adds a new allowed key.
+    Add an allowed key for DID metadata (update the DID Metadata Conventions table with a new key).
 
     :param key: the name for the new key.
     :param key_type: the type of the key: all(container, dataset, file), collection(dataset or container), file, derived(compute from file for collection).
@@ -65,7 +62,7 @@ def add_key(key, key_type, value_type=None, value_regexp=None, *, session: "Sess
     except ValueError:
         raise UnsupportedKeyType('The type \'%s\' is not supported for keys!' % str(key_type))
 
-    new_key = models.DIDKey(key=key, value_type=value_type and str(value_type), value_regexp=value_regexp, key_type=key_type)
+    new_key = models.DIDMetaConventionsKey(key=key, value_type=value_type and str(value_type), value_regexp=value_regexp, key_type=key_type)
     try:
         new_key.save(session=session)
     except IntegrityError as error:
@@ -77,46 +74,54 @@ def add_key(key, key_type, value_type=None, value_regexp=None, *, session: "Sess
            or match('.*UniqueViolation.*duplicate key value violates unique constraint.*', error.args[0]) \
            or match('.*IntegrityError.*columns? key.*not unique.*', error.args[0]):
             raise Duplicate(f"key '{key}' already exists!")
-        raise
+        raise RucioException(error.args)
 
 
 @transactional_session
-def del_key(key, *, session: "Session"):
+def del_key(key: str, *, session: "Session") -> None:
     """
-    Deletes a key.
+    Delete a key in the DID Metadata Conventions table.
 
     :param key: the name for the key.
     :param session: The database session in use.
     """
-    session.query(models.DIDKey).filter(key == key).delete()
+    stmt = delete(
+        models.DIDMetaConventionsKey
+    ).where(
+        models.DIDMetaConventionsKey.key == key
+    )
+    session.execute(stmt)
 
 
 @read_session
-def list_keys(*, session: "Session"):
+def list_keys(*, session: "Session") -> list[str]:
     """
-    Lists all keys.
+    Lists all keys for DID Metadata Conventions.
 
     :param session: The database session in use.
 
     :returns: A list containing all keys.
     """
-    key_list = []
-    query = session.query(models.DIDKey)
-    for row in query:
-        key_list.append(row.key)
-    return key_list
+    stmt = select(
+        models.DIDMetaConventionsKey.key
+    )
+    return list(session.execute(stmt).scalars().all())
 
 
 @transactional_session
-def add_value(key, value, *, session: "Session"):
+def add_value(key: str, value: str, *, session: "Session") -> None:
     """
-    Adds a new value to a key.
+    Adds a new value for a key in DID Metadata Convention.
 
     :param key: the name for the key.
     :param value: the value.
     :param session: The database session in use.
+
+    :raises Duplicate: Key-Value pair exists
+    :raises KeyNotFound: Key not in metadata conventions table
+    :raises InvalidValueForKey: Value conflicts with rse expression for key values or does not have the correct type
     """
-    new_value = models.DIDKeyValueAssociation(key=key, value=value)
+    new_value = models.DIDMetaConventionsConstraint(key=key, value=value)
     try:
         new_value.save(session=session)
     except IntegrityError as error:
@@ -137,37 +142,43 @@ def add_value(key, value, *, session: "Session"):
 
         raise RucioException(error.args)
 
-    k = session.query(models.DIDKey).filter_by(key=key).one()
+    statement = select(
+        models.DIDMetaConventionsKey,
+    ).where(
+        models.DIDMetaConventionsKey.key == key
+    )
+    query = session.execute(statement).scalar_one()
 
     # Check value against regexp, if defined
-    if k.value_regexp and not match(k.value_regexp, value):
-        raise InvalidValueForKey("The value '%s' for the key '%s' does not match the regular expression '%s'" % (value, key, k.value_regexp))
+    if query.value_regexp and not match(query.value_regexp, value):
+        raise InvalidValueForKey(f"The value {value} for the key {key} does not match the regular expression {query.value_regexp}")
 
     # Check value type, if defined
     type_map = dict([(str(t), t) for t in AUTHORIZED_VALUE_TYPES])
-    if k.value_type and not isinstance(value, type_map.get(k.value_type)):
-        raise InvalidValueForKey("The value '%s' for the key '%s' does not match the required type '%s'" % (value, key, k.value_type))
+    if query.value_type and not isinstance(value, type_map.get(query.value_type)):  # type: ignore ; Typing error caused by 'isinstaince' not thinking types count as classes
+        raise InvalidValueForKey(f"The value {value} for the key {key} does not match the required type {query.value_type}")
 
 
 @read_session
-def list_values(key, *, session: "Session"):
+def list_values(key: str, *, session: "Session") -> list[str]:
     """
-    Lists all values for a key.
+    Lists all allowed values for a DID key (all values for a key in DID Metadata Conventions).
 
     :param key: the name for the key.
     :param session: The database session in use.
 
     :returns: A list containing all values.
     """
-    value_list = []
-    query = session.query(models.DIDKeyValueAssociation).filter_by(key=key)
-    for row in query:
-        value_list.append(row.value)
-    return value_list
+    statement = select(
+        models.DIDMetaConventionsConstraint.value
+    ).where(
+        models.DIDMetaConventionsConstraint.key == key
+    )
+    return list(session.execute(statement).scalars().all())
 
 
 @read_session
-def validate_meta(meta, did_type, *, session: "Session"):
+def validate_meta(meta: dict, did_type: DIDType, *, session: "Session") -> None:
     """
     Validates metadata for a did.
 
@@ -175,16 +186,18 @@ def validate_meta(meta, did_type, *, session: "Session"):
     :param meta: the type of the did, e.g, DATASET, CONTAINER, FILE.
     :param session: The database session in use.
 
-    :returns: True
+    :raises InvalidObject:
     """
     # For now only validate the datatype for datasets
     key = 'datatype'
     if did_type == DIDType.DATASET and key in meta:
         try:
-            session.query(models.DIDKeyValueAssociation.value).\
-                filter_by(key=key).\
-                filter_by(value=meta[key]).\
-                one()
+            statement = select(
+                models.DIDMetaConventionsConstraint.value
+            ).where(
+                and_(models.DIDMetaConventionsConstraint.value == meta[key],
+                     models.DIDMetaConventionsConstraint.key == key)
+            )
+            session.execute(statement).one()
         except NoResultFound:
-            print("The value '%s' for the key '%s' is not valid" % (meta[key], key))
-            raise InvalidObject("The value '%s' for the key '%s' is not valid" % (meta[key], key))
+            raise InvalidObject(f"The value {meta[key]}' for the key {key} is not valid")

rucio/core/monitor.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,7 +16,6 @@
 Graphite and prometheus metrics
 """
 
-import __main__ as main
 import atexit
 import logging
 import os
@@ -28,23 +26,26 @@ from datetime import datetime, timedelta
 from functools import wraps
 from pathlib import Path
 from threading import Lock
-from typing import Any, Optional, TypeVar, Union
+from typing import TYPE_CHECKING, Any, Optional, TypeVar, Union
 
-from prometheus_client import (Counter, Gauge, Histogram, REGISTRY, CollectorRegistry, generate_latest, multiprocess,
-                               push_to_gateway, start_http_server, values)
+from prometheus_client import REGISTRY, CollectorRegistry, Counter, Gauge, Histogram, generate_latest, multiprocess, push_to_gateway, start_http_server, values
 from statsd import StatsClient
 
+import __main__ as main
 from rucio.common.config import config_get, config_get_bool, config_get_int
 from rucio.common.stopwatch import Stopwatch
 from rucio.common.utils import retrying
 
+if TYPE_CHECKING:
+    from rucio.common.types import LoggerFunction
+
 _T = TypeVar('_T')
 _M = TypeVar('_M', bound="_MultiMetric")
 
 PROMETHEUS_MULTIPROC_DIR = os.environ.get('PROMETHEUS_MULTIPROC_DIR', os.environ.get('prometheus_multiproc_dir', None))
 
 
-def cleanup_prometheus_files_at_exit():
+def cleanup_prometheus_files_at_exit() -> None:
     if PROMETHEUS_MULTIPROC_DIR:
         multiprocess.mark_process_dead(os.getpid())
 
@@ -101,7 +102,12 @@ METRICS_LOCK = Lock()
 _HISTOGRAM_DEFAULT_BUCKETS = Histogram.DEFAULT_BUCKETS
 
 
-def _cleanup_old_prometheus_files(path, file_pattern, cleanup_delay, logger):
+def _cleanup_old_prometheus_files(
+        path: str,
+        file_pattern: str,
+        cleanup_delay: float,
+        logger: "LoggerFunction"
+) -> None:
     """cleanup behind processes which didn't finish gracefully."""
 
     oldest_accepted_mtime = datetime.now() - timedelta(seconds=cleanup_delay)
@@ -120,7 +126,7 @@ def _cleanup_old_prometheus_files(path, file_pattern, cleanup_delay, logger):
                 pass
 
 
-def cleanup_old_prometheus_files(logger=logging.log):
+def cleanup_old_prometheus_files(logger: "LoggerFunction" = logging.log) -> None:
     path = PROMETHEUS_MULTIPROC_DIR
     if path:
         _cleanup_old_prometheus_files(path, file_pattern='gauge_live*.db', cleanup_delay=timedelta(hours=1).total_seconds(), logger=logger)
@@ -130,7 +136,7 @@ def cleanup_old_prometheus_files(logger=logging.log):
 @retrying(retry_on_exception=lambda _: True,
           wait_fixed=500,
           stop_max_attempt_number=2)
-def generate_prometheus_metrics():
+def generate_prometheus_metrics() -> bytes:
     cleanup_old_prometheus_files()
 
     registry = CollectorRegistry()
@@ -209,7 +215,7 @@ class _MultiCounter(_MultiMetric):
         if STATSD_CLIENT:
             STATSD_CLIENT.incr(self._statsd, delta)
 
-    def init_prometheus_metric(self, name: str, documentation: Optional[str], labelnames: Sequence[str] = ()):
+    def init_prometheus_metric(self, name: str, documentation: str, labelnames: Sequence[str] = ()) -> Counter:
         return Counter(name, documentation, labelnames=labelnames, registry=self._registry)
 
 
@@ -220,7 +226,7 @@ class _MultiGauge(_MultiMetric):
         if STATSD_CLIENT:
             STATSD_CLIENT.gauge(self._statsd, value)
 
-    def init_prometheus_metric(self, name: str, documentation: Optional[str], labelnames: Sequence[str] = ()):
+    def init_prometheus_metric(self, name: str, documentation: str, labelnames: Sequence[str] = ()) -> Gauge:
         return Gauge(name, documentation, labelnames=labelnames, registry=self._registry)
 
 
@@ -244,7 +250,7 @@ class _MultiTiming(_MultiMetric):
         if STATSD_CLIENT:
             STATSD_CLIENT.timing(self._statsd, value * 1000)
 
-    def init_prometheus_metric(self, name: str, documentation: Optional[str], labelnames: Sequence[str] = ()):
+    def init_prometheus_metric(self, name: str, documentation: str, labelnames: Sequence[str] = ()) -> Histogram:
         return Histogram(name, documentation, labelnames=labelnames, registry=self._registry, buckets=self._histogram_buckets)
 
     def __enter__(self):
@@ -336,7 +342,7 @@ class MetricManager:
         self.registry = registry or REGISTRY
         self.push_gateways = push_gateways or []
 
-    def full_name(self, name: str):
+    def full_name(self, name: str) -> str:
         if self.prefix:
             return f'{self.prefix}.{name}'
         return name