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

rucio/common/dumper/__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");
@@ -22,6 +21,7 @@ import os
 import re
 import sys
 import tempfile
+from typing import TYPE_CHECKING, Any, Literal, Optional, Union
 
 import gfal2
 import magic
@@ -30,29 +30,39 @@ import requests
 from rucio.common import config
 from rucio.core.rse import get_rse_id, get_rse_protocols
 
+if TYPE_CHECKING:
+    from collections.abc import Iterator
+    from multiprocessing.connection import Connection
+    from types import ModuleType
+    from typing import IO, TextIO
+
+    from _typeshed import FileDescriptorOrPath, GenericPath, StrOrBytesPath
+
+    from rucio.common.types import RSEProtocolDict
+
 
 class HTTPDownloadFailed(Exception):
-    def __init__(self, msg='', code=None):
+    def __init__(self, msg: str = '', code: Optional[str] = None):
         self.code = code
         if code is not None:
             msg = '{0} (Status {1})'.format(msg, code)
         super(HTTPDownloadFailed, self).__init__(msg)
 
 
-class LogPipeHandler(logging.Handler, object):
-    def __init__(self, pipe):
+class LogPipeHandler(logging.Handler):
+    def __init__(self, pipe: "Connection"):
         super(LogPipeHandler, self).__init__()
         self.pipe = pipe
 
-    def emit(self, record):
+    def emit(self, record: logging.LogRecord) -> None:
         self.pipe.send(self.format(record))
 
-    def close(self):
+    def close(self) -> None:
         super(LogPipeHandler, self).close()
         self.pipe.close()
 
 
-def error(text, exit_code=1):
+def error(text: str, exit_code: int = 1) -> None:
     '''
     Log and print `text` error. This function ends the execution of the program with exit code
     `exit_code` (defaults to 1).
@@ -63,7 +73,7 @@ def error(text, exit_code=1):
     exit(1)
 
 
-def mkdir(dir_):
+def mkdir(dir_: "StrOrBytesPath") -> None:
     '''
     This functions creates the `dir_` directory if it doesn't exist. If `dir_`
     already exists this function does nothing.
@@ -71,10 +81,11 @@ def mkdir(dir_):
     try:
         os.mkdir(dir_)
     except OSError as error:
-        assert error.errno == 17
+        if error.errno != 17:
+            raise error
 
 
-def cacert_config(config, rucio_home):
+def cacert_config(config: "ModuleType", rucio_home: str) -> Optional[Union["FileDescriptorOrPath", Literal[False]]]:
     logger = logging.getLogger('dumper.__init__')
     try:
         cacert = config.config_get('client', 'ca_cert').replace('$RUCIO_HOME', rucio_home)
@@ -88,13 +99,13 @@ def cacert_config(config, rucio_home):
     return cacert
 
 
-def rucio_home():
+def rucio_home() -> str:
     return os.environ.get('RUCIO_HOME', '/opt/rucio')
 
 
-def get_requests_session():
+def get_requests_session() -> requests.Session:
     requests_session = requests.Session()
-    requests_session.verify = cacert_config(config, rucio_home())
+    requests_session.verify = cacert_config(config, rucio_home())  # type: ignore
     requests_session.stream = True
     return requests_session
 
@@ -116,7 +127,7 @@ else:
 # pylint: enable=no-member
 
 
-def isplaintext(filename):
+def isplaintext(filename: "GenericPath") -> bool:
     '''
     Returns True if `filename` has mimetype == 'text/plain'.
     '''
@@ -125,7 +136,7 @@ def isplaintext(filename):
     return mimetype(filename).split(';')[0] == 'text/plain'
 
 
-def smart_open(filename):
+def smart_open(filename: "GenericPath") -> Optional[Union["TextIO", gzip.GzipFile]]:
     '''
     Returns an open file object if `filename` is plain text, else assumes
     it is a bzip2 compressed file and returns a file-like object to
@@ -146,14 +157,18 @@ def smart_open(filename):
 
 
 @contextlib.contextmanager
-def temp_file(directory, final_name=None, binary=False):
+def temp_file(
+    directory: str,
+    final_name: Optional[str] = None,
+    binary: bool = False
+) -> "Iterator[tuple[IO[Any], StrOrBytesPath]]":
     '''
     Allows to create a temporal file to store partial results, when the
     file is complete it is renamed to `final_name`.
 
     - `directory`: working path to create the temporal and the final file.
     - `final_name`: Path of the final file, relative to `directory`.
-       If the `final_name` is omitted or None the renaming step is ommited,
+       If the `final_name` is omitted or None the renaming step is omitted,
        leaving the temporal file with the results.
     - `binary`: whether to open the file in binary mode (default: False).
 
@@ -193,7 +208,7 @@ DATETIME_FORMAT_FULL = '%Y-%m-%dT%H:%M:%S'
 MILLISECONDS_RE = re.compile(r'\.(\d{3})Z$')
 
 
-def to_datetime(str_or_datetime):
+def to_datetime(str_or_datetime: Union[datetime.datetime, str]) -> Optional[datetime.datetime]:
     """
     Convert string to datetime. The format is somewhat flexible.
     Timezone information is ignored.
@@ -221,51 +236,65 @@ def to_datetime(str_or_datetime):
             'Trying to parse "%s" date with resolution of milliseconds',
             str_or_datetime,
         )
-        miliseconds = int(MILLISECONDS_RE.search(str_or_datetime).group(1))
+        milliseconds = int(MILLISECONDS_RE.search(str_or_datetime).group(1))
         str_or_datetime = MILLISECONDS_RE.sub('', str_or_datetime)
         date = datetime.datetime.strptime(
             str_or_datetime,
             DATETIME_FORMAT,
         )
-        date = date + datetime.timedelta(microseconds=miliseconds * 1000)
+        date = date + datetime.timedelta(microseconds=milliseconds * 1000)
     return date
 
 
-def ddmendpoint_preferred_protocol(ddmendpoint):
+def ddmendpoint_preferred_protocol(ddmendpoint: str) -> "RSEProtocolDict":
     return next(p for p in get_rse_protocols(get_rse_id(ddmendpoint))['protocols'] if p['domains']['wan']['read'] == 1)
 
 
-def ddmendpoint_url(ddmendpoint):
+def ddmendpoint_url(ddmendpoint: str) -> str:
     preferred_protocol = ddmendpoint_preferred_protocol(ddmendpoint)
     prefix = re.sub(r'rucio/$', '', preferred_protocol['prefix'])
     return '{scheme}://{hostname}:{port}'.format(**preferred_protocol) + prefix
 
 
-def http_download_to_file(url, file_, session=None):
+def http_download_to_file(url: str, file_: "IO", session: Optional[requests.Session] = None) -> None:
     '''
     Download the file in `url` storing it in the `file_` file-like
     object.
     If given `session` must be a requests.Session instance, and will be
     used to download the file, otherwise requests.get() will be used.
     '''
-    if session is None:
-        response = requests.get(url, stream=True)
-    else:
-        response = session.get(url)
-
-    if response.status_code != 200:
-        logging.error(
-            'Retrieving %s returned %d status code',
-            url,
-            response.status_code,
-        )
-        raise HTTPDownloadFailed('Error downloading ' + url, response.status_code)
+    def _do_download(url, file_, session, try_decode=False):
+        if session is None:
+            response = requests.get(url, stream=True)
+        else:
+            response = session.get(url)
+
+        if response.status_code != 200:
+            logging.error(
+                'Retrieving %s returned %d status code',
+                url,
+                response.status_code,
+            )
+            raise HTTPDownloadFailed('Error downloading ' + url, str(response.status_code))
+
+        if try_decode:
+            if response.encoding is None:
+                response.encoding = 'utf-8'
+            for chunk in response.iter_content(CHUNK_SIZE, decode_unicode=True):
+                file_.write(chunk)
+        else:
+            for chunk in response.iter_content(CHUNK_SIZE):
+                file_.write(chunk)
 
-    for chunk in response.iter_content(CHUNK_SIZE):
-        file_.write(chunk)
+    try:
+        # try without decoding first
+        _do_download(url, file_, session, False)
+    except TypeError:
+        # if that fails due to writing binary data to text file, try to force decode
+        _do_download(url, file_, session, True)
 
 
-def http_download(url, filename):
+def http_download(url: str, filename: "FileDescriptorOrPath") -> None:
     '''
     Download the file in `url` storing it in the path given by `filename`.
     '''
@@ -273,7 +302,7 @@ def http_download(url, filename):
         http_download_to_file(url, f)
 
 
-def gfal_download_to_file(url, file_):
+def gfal_download_to_file(url: str, file_: "IO") -> None:
     '''
     Download the file in `url` storing it in the `file_` file-like
     object.
@@ -286,7 +315,7 @@ def gfal_download_to_file(url, file_):
         chunk = infile.read(CHUNK_SIZE)
     except gfal2.GError as e:
         if e.code == 70:
-            logger.debug('GError(70) raised, using GRIDFTP PLUGIN:STAT_ON_OPEN=False workarround to download %s', url)
+            logger.debug('GError(70) raised, using GRIDFTP PLUGIN:STAT_ON_OPEN=False workaround to download %s', url)
             ctx.set_opt_boolean('GRIDFTP PLUGIN', 'STAT_ON_OPEN', False)
             infile = ctx.open(url, 'r')
             chunk = infile.read(CHUNK_SIZE)
@@ -298,7 +327,7 @@ def gfal_download_to_file(url, file_):
         chunk = infile.read(CHUNK_SIZE)
 
 
-def gfal_download(url, filename):
+def gfal_download(url: str, filename: "FileDescriptorOrPath") -> None:
     '''
     Download the file in `url` storing it in the path given by `filename`.
     '''

rucio/common/dumper/consistency.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");
@@ -19,9 +18,10 @@ import os
 import re
 import subprocess
 import tempfile
+from typing import cast
 
 from rucio.common import dumper
-from rucio.common.dumper import error, DUMPS_CACHE_DIR, data_models, path_parsing
+from rucio.common.dumper import DUMPS_CACHE_DIR, data_models, error, path_parsing
 
 subcommands = ['consistency', 'consistency-manual']
 
@@ -42,10 +42,12 @@ class Consistency(data_models.DataModel):
                 ddm_endpoint, prev_date, cache_dir=cache_dir)
             next_date_fname = data_models.Replica.download(
                 ddm_endpoint, next_date, cache_dir=cache_dir)
-            assert prev_date_fname is not None
-            assert next_date_fname is not None
+            if (prev_date_fname is None) or (next_date_fname is None):
+                raise ValueError("Both prev_date_fname and next_date_fname are required for subcommand='consistency'; found prev_data_fname=%s and next_date_fname=%s" % (prev_date_fname, next_date_fname))
+        elif subcommand == 'consistency-manual':
+            pass
         else:
-            assert subcommand == 'consistency-manual'
+            raise ValueError("subcommand %s not accepted, choice from ('consistency','consistency-manual')" % subcommand)
 
         prefix_components = path_parsing.components(dumper.ddmendpoint_url(ddm_endpoint))
 
@@ -67,7 +69,7 @@ class Consistency(data_models.DataModel):
             Parser to have consistent paths in storage dumps.
 
             :param line: String with one line of a dump.
-            :returns: Path formated as in the Rucio Replica Dumps.
+            :returns: Path formatted as in the Rucio Replica Dumps.
             '''
             relative = path_parsing.remove_prefix(
                 prefix_components,
@@ -162,13 +164,14 @@ def _try_to_advance(it, default=None):
     return el.strip()
 
 
-def min3(*values):
+def min_value(*values):
     '''
-    Minimum between the 3 values ignoring None
+    Minimum between the input values, ignoring None
     '''
-    values = [value for value in values if value is not None]
-    assert len(values) > 0
-    return min(values)
+    values_without_none = cast('list[str]', [value for value in values if value is not None])
+    if len(values_without_none) == 0:
+        raise ValueError("Input contains 0 non-null values.")
+    return min(values_without_none)
 
 
 def split_if_not_none(value, sep=',', fields=2):
@@ -199,7 +202,7 @@ def compare3(it0, it1, it2):
         path0, status0 = split_if_not_none(v0)
         path2, status2 = split_if_not_none(v2)
 
-        vmin = min3(path0, v1, path2)
+        vmin = min_value(path0, v1, path2)
         in0 = in1 = in2 = False
         in0_status = in2_status = None
 
@@ -292,7 +295,7 @@ def gnu_sort(file_path, prefix=None, delimiter=None, fieldspec=None, cache_dir=D
 
     :param prefix: If given the output file will be named <prefix>_sorted.
     Otherwise the prefix is the name of the input file.
-    :param delimiter: Delimiter character if the data is formated in
+    :param delimiter: Delimiter character if the data is formatted in
     columns (argument of -t in the sort command).
     :param fieldspec: String with the specification of column or columns
     to be used to sort (argument -k in the sort command).
@@ -302,7 +305,8 @@ def gnu_sort(file_path, prefix=None, delimiter=None, fieldspec=None, cache_dir=D
     memory and it is relatively fast if used with the environment variable
     LC_ALL set to C as in this function.
     '''
-    assert (delimiter is None and fieldspec is None) or (delimiter is not None and fieldspec is not None)
+    if (delimiter is not None) ^ (fieldspec is not None):
+        raise ValueError("Either both delimiter and fieldspec is set, or neither are.")
     if delimiter is None:
         cmd_line = 'LC_ALL=C sort {0} > {1}'
     else:
@@ -316,23 +320,21 @@ def gnu_sort(file_path, prefix=None, delimiter=None, fieldspec=None, cache_dir=D
     if os.path.exists(sorted_path):
         return sorted_path
 
-    # FIXME: mktemp() is an insecure function and this may be a security
-    # threat in some scenarios. Find another way to do it.
-    tfile = tempfile.mktemp(dir=cache_dir)
+    tfile = tempfile.NamedTemporaryFile(dir=cache_dir, delete=False)
 
     subprocess.check_call(
-        cmd_line.format(file_path, tfile),
+        cmd_line.format(file_path, tfile.name),
         shell=True,
     )
 
-    os.link(tfile, sorted_path)
-    os.unlink(tfile)
+    os.link(tfile.name, sorted_path)
+    os.unlink(tfile.name)
 
     return sorted_path
 
 
 def populate_args(argparser):
-    # Option to download the rucio replica dumps automaticaly
+    # Option to download the rucio replica dumps automatically
     parser = argparser.add_parser(
         'consistency',
         help='Consistency check to verify possible lost files and dark data '
@@ -391,8 +393,10 @@ def _parse_args_consistency(args):
         error('The storage dump filename must be of the form '
               '"dump_YYYYMMDD" where the date correspond to the date '
               'of the newest files included')
+
     date_str = date_str.group(1)
-    assert date_str is not None
+    if date_str is None:
+        error('Invalid date {0}'.format(date_str))
     try:
         args_dict['date'] = date = datetime.datetime.strptime(date_str, '%Y%m%d')
     except ValueError:

rucio/common/dumper/data_models.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");
@@ -27,16 +26,10 @@ import re
 
 from tabulate import tabulate
 
-from rucio.common.dumper import DUMPS_CACHE_DIR
-from rucio.common.dumper import HTTPDownloadFailed
-from rucio.common.dumper import get_requests_session
-from rucio.common.dumper import http_download_to_file
-from rucio.common.dumper import smart_open
-from rucio.common.dumper import temp_file
-from rucio.common.dumper import to_datetime
+from rucio.common.dumper import DUMPS_CACHE_DIR, HTTPDownloadFailed, get_requests_session, http_download_to_file, smart_open, temp_file, to_datetime
 
 
-class DataModel(object):
+class DataModel:
     """
     Data model for the dumps
     """
@@ -50,7 +43,7 @@ class DataModel(object):
     def __init__(self, *args):
         if len(args) != len(self.SCHEMA):
             raise TypeError(
-                'Wrong number of parameters (fields) to initalize {0} '
+                'Wrong number of parameters (fields) to initialize {0} '
                 'instance. {1} given, {2} expected:\n{3}'.format(
                     type(self).__name__,
                     len(args),
@@ -157,20 +150,17 @@ class DataModel(object):
         logger = logging.getLogger('auditor.data_models')
         requests_session = get_requests_session()
         if date == 'latest':
-            url = ''.join((cls.BASE_URL, cls.URI, '?rse={0}'.format(rse)))
+            url = ''.join((cls.BASE_URL, cls.URI, '?rse={0}'.format(rse)))  # type: ignore
             request_headers = requests_session.head(url)
             for field in request_headers.headers['content-disposition'].split(';'):
                 if field.startswith('filename='):
                     date = field.split('=')[1].split('_')[-1].split('.')[0]
 
+        elif isinstance(date, datetime.datetime):
+            date = date.strftime('%d-%m-%Y')
+            url = ''.join((cls.BASE_URL, cls.URI, '?rse={0}&date={1}'.format(rse, date)))  # type: ignore
         else:
-            assert isinstance(date, datetime.datetime)
-            date = date.strftime('%d-%m-%Y')  # pylint: disable=no-member
-            url = ''.join((
-                cls.BASE_URL,
-                cls.URI,
-                '?rse={0}&date={1}'.format(rse, date),
-            ))
+            raise ValueError("Passed date (%s) must be a datetime object or 'latest'." % date)
 
         if not os.path.isdir(cache_dir):
             os.mkdir(cache_dir)
@@ -193,7 +183,7 @@ class DataModel(object):
                     url,
                     response.status_code,
                 )
-                raise HTTPDownloadFailed('Downloading {0} dump'.format(cls.__name__), code=response.status_code)
+                raise HTTPDownloadFailed('Downloading {0} dump'.format(cls.__name__), code=str(response.status_code))
 
             with temp_file(cache_dir, final_name=filename) as (tfile, _):
                 http_download_to_file(url, tfile, session=requests_session)
@@ -244,7 +234,8 @@ class CompleteDataset(DataModel):
             self.state = args[7]
         else:
             self.state = None
-        assert len(args) <= 8
+        if len(args) > 8:
+            raise ValueError("Too many arguments, must be 8 or less. Instead passed %s" % len(args))
 
 
 class Replica(DataModel):
@@ -271,10 +262,11 @@ class Replica(DataModel):
 
         if len(args) == 8:
             logger.warning('Missing parameter\nrse: %s\ndataset: %s\n', self.rse, self.name)
-        assert len(args) <= 9
+        elif len(args) > 9:
+            raise ValueError("Too many arguments. Must be 9 or less, instead passed %s" % len(args))
 
 
-class Filter(object):
+class Filter:
     _Condition = collections.namedtuple('_Condition', ('comparator', 'attribute', 'expected'))
 
     def __init__(self, filter_str, record_class):
@@ -299,7 +291,8 @@ class Filter(object):
         for expr in filter_str.split(','):
             key, expected = expr.split('=')
             # Better checks required
-            assert key in record_class.get_fieldnames()
+            if key not in record_class.get_fieldnames():
+                raise ValueError("Key %s not supported." % key)
             parser = list(filter(lambda t: t[0] == key, record_class.SCHEMA))[0][1]
             self.conditions.append(self._Condition(
                 comparator=operator.eq,

rucio/common/dumper/path_parsing.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");