rucio 35.7.0__py3-none-any.whl → 37.0.0__py3-none-any.whl

rucio/core/did_meta_plugins/elasticsearch_meta.py

@@ -0,0 +1,407 @@
+# Copyright European Organization for Nuclear Research (CERN) since 2012
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#    http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+
+'''
+ Elasticsearch based metadata plugin.
+'''
+
+import datetime
+import operator
+from typing import TYPE_CHECKING, Any, Literal, Optional, Union
+
+from elasticsearch import Elasticsearch
+from elasticsearch import exceptions as elastic_exceptions
+
+from rucio.common import config, exception
+from rucio.core.did_meta_plugins.did_meta_plugin_interface import DidMetaPlugin
+from rucio.core.did_meta_plugins.filter_engine import FilterEngine
+
+if TYPE_CHECKING:
+    from collections.abc import Iterator
+
+    from sqlalchemy.orm import Session
+
+    from rucio.common.types import InternalScope
+
+IMMUTABLE_KEYS = [
+    'scope',            # generated on insert
+    'name',             # generated on insert
+    'vo'                # generated on insert
+]
+
+
+class ElasticDidMeta(DidMetaPlugin):
+    def __init__(
+        self,
+        hosts: Optional[list[str]] = None,
+        user: Optional[str] = None,
+        password: Optional[str] = None,
+        index: Optional[str] = None,
+        archive_index: Optional[str] = None,
+        use_ssl: Optional[bool] = False,
+        verify_certs: bool = True,
+        ca_certs: Optional[str] = None,
+        client_cert: Optional[str] = None,
+        client_key: Optional[str] = None,
+        request_timeout: int = 100,
+        max_retries: int = 3,
+        retry_on_timeout: bool = False
+    ) -> None:
+        super(ElasticDidMeta, self).__init__()
+        hosts = hosts or [config.config_get('metadata', 'elastic_service_hosts')]
+        user = user or config.config_get('metadata', 'elastic_user', False, None)
+        password = password or config.config_get('metadata', 'elastic_password', False, None)
+        self.index = index or config.config_get('metadata', 'meta_index', False, 'rucio_did_meta')
+        self.archive_index = archive_index or config.config_get('metadata', 'archive_index', False, 'archive_meta')
+        use_ssl = use_ssl or config.config_get_bool('metadata', 'use_ssl', False, False)
+        ca_certs = ca_certs or config.config_get('metadata', 'ca_certs', False, None)
+        client_cert = client_cert or config.config_get('metadata', 'client_cert', False, None)
+        client_key = client_key or config.config_get('metadata', 'client_key', False, None)
+
+        self.es_config = {
+            'hosts': hosts,
+            'timeout': request_timeout,
+            'max_retries': max_retries,
+            'retry_on_timeout': retry_on_timeout
+        }
+        if user and password:
+            self.es_config['basic_auth'] = (user, password)
+
+        if use_ssl:
+            self.es_config.update({
+                'ca_certs': ca_certs,
+                'verify_certs': verify_certs,
+            })
+            if client_cert and client_key:
+                self.es_config.update({
+                    'client_cert': client_cert,
+                    'client_key': client_key
+                })
+
+        self.client = Elasticsearch(**self.es_config)
+        self.plugin_name = "ELASTIC"
+
+    def drop_index(self) -> None:
+        self.client.indices.delete(index=self.index)
+
+    def get_metadata(
+        self,
+        scope: "InternalScope",
+        name: str,
+        *,
+        session: "Optional[Session]" = None
+    ) -> dict[str, Any]:
+        """
+        Get data identifier metadata.
+
+        :param scope: The scope name
+        :param name: The data identifier name
+        :param session: The database session in use
+        :returns: The metadata for the did
+        :raises DataIdentifierNotFound: If the DID metadata is not found.
+        :raises RucioException: If another error occurs during the process.
+        """
+
+        doc_id = f"{scope.internal}{name}"
+        try:
+            doc = self.client.get(index=self.index, id=doc_id)["_source"]
+        except elastic_exceptions.NotFoundError as err:
+            raise exception.DataIdentifierNotFound(f"No metadata found for DID '{scope}:{name}' not found") from err
+        except Exception as err:
+            raise exception.RucioException(err)
+        return doc
+
+    def set_metadata(
+        self,
+        scope: "InternalScope",
+        name: str,
+        key: str,
+        value: str,
+        recursive: bool = False,
+        *,
+        session: "Optional[Session]" = None
+    ) -> None:
+        """
+        Set single metadata key.
+
+        :param scope: the scope of did
+        :param name: the name of the did
+        :param key: the key to be added
+        :param value: the value of the key to be added
+        :param recursive: recurse into DIDs (not supported)
+        :param session: The database session in use
+        :raises DataIdentifierNotFound: If the DID is not found.
+        :raises RucioException: If an error occurs while setting the metadata.
+        """
+        self.set_metadata_bulk(scope=scope, name=name, meta={key: value}, recursive=recursive, session=session)
+
+    def set_metadata_bulk(
+        self,
+        scope: "InternalScope",
+        name: str,
+        meta: dict[str, Any],
+        recursive: bool = False,
+        *,
+        session: "Optional[Session]" = None
+    ) -> None:
+        """
+        Bulk set metadata keys.
+
+        :param scope: the scope of did
+        :param name: the name of the did
+        :param meta: dictionary of metadata keypairs to be added
+        :param recursive: recurse into DIDs (not supported)
+        :param session: The database session in use
+        :raises DataIdentifierNotFound: If the DID is not found.
+        :raises UnsupportedOperation: If recursive inserts are requested (currently unsupported).
+        :raises RucioException: If an error occurs while setting the metadata.
+        """
+        doc_id = f"{scope.internal}{name}"
+        try:
+            # Try to get existing metadata
+            existing_meta = self.get_metadata(scope, name)
+        except exception.DataIdentifierNotFound:
+            existing_meta = {
+                'scope': str(scope.external),
+                'name': name,
+                'vo': scope.vo
+            }
+        for key, value in meta.items():
+            if key not in IMMUTABLE_KEYS:
+                existing_meta[key] = value
+
+        try:
+            self.client.index(index=self.index, body=existing_meta, id=doc_id, refresh="true")
+        except Exception as err:
+            raise exception.RucioException(err)
+
+        if recursive:
+            raise exception.UnsupportedOperation(f"'{self.plugin_name.lower()}' metadata module does not currently support recursive inserts of metadata")
+
+    def delete_metadata(
+        self,
+        scope: "InternalScope",
+        name: str,
+        key: str,
+        *,
+        session: "Optional[Session]" = None
+    ) -> None:
+        """
+        Delete a key from metadata.
+
+        :param scope: the scope of did
+        :param name: the name of the did
+        :param key: the key to be deleted
+        :raises DataIdentifierNotFound: If the DID is not found.
+        :raises RucioException: If an error occurs while setting the metadata.
+        """
+        doc_id = f"{scope.internal}{name}"
+        try:
+            # First, get the current document
+            doc = self.client.get(index=self.index, id=doc_id)
+
+            # Check if the key exists in the document
+            if key in doc['_source']:
+                # Use script to remove the field
+                script = {
+                    "script": {
+                        "source": f"ctx._source.remove('{key}')",
+                        "lang": "painless"
+                    }
+                }
+                self.client.update(index=self.index, id=doc_id, body=script)
+        except elastic_exceptions.NotFoundError as err:
+            raise exception.DataIdentifierNotFound(f"No metadata found for DID '{scope}:{name}' not found") from err
+        except Exception as err:
+            raise exception.RucioException(err)
+
+    def list_dids(
+        self,
+        scope: "InternalScope",
+        filters: Union[list[dict[str, Any]], dict[str, Any]],
+        did_type: Literal['all', 'collection', 'dataset', 'container', 'file'] = 'collection',
+        ignore_case: bool = False,
+        limit: Optional[int] = None,
+        offset: Optional[int] = None,
+        long: bool = False,
+        recursive: bool = False,
+        ignore_dids: Optional[list] = None,
+        *,
+        session: "Optional[Session]" = None
+    ) -> "Iterator[dict[str, Any]]":
+        """
+        List DIDs (Data Identifier).
+
+        :param scope: The scope of the DIDs to search.
+        :param filters: The filters to apply to the DID search.
+        :param did_type: The type of DID (default is 'collection').
+        :param ignore_case: Whether to ignore case (default is False).
+        :param limit: The maximum number of DIDs to return.
+        :param offset: The starting point for the search (used for pagination).
+        :param long: Whether to return extended information (scope, name, did_type, bytes, length) (default is False).
+        :param recursive: Whether to search recursively (currently unsupported).
+        :param ignore_dids: A list of DIDs to ignore (default is an empty list).
+        :param session: The database session in use.
+        :returns: A generator yielding DIDs as strings (when `long` is False) or dictionaries (when `long` is True).
+        :raises UnsupportedOperation: If recursive searches are requested (currently unsupported).
+        :raises RucioException: If an error occurs during the search.
+        """
+
+        if not ignore_dids:
+            ignore_dids = []
+
+        # backwards compatibility for filters as single {}.
+        if isinstance(filters, dict):
+            filters = [filters]
+
+        # Create Elasticsearch query
+        fe = FilterEngine(filters, model_class=None, strict_coerce=False)
+        elastic_query_str = fe.create_elastic_query(
+            additional_filters=[
+                ('scope', operator.eq, str(scope.external)),
+                ('vo', operator.eq, scope.vo)
+            ]
+        )
+        pit = self.client.open_point_in_time(index=self.index, keep_alive="2m")
+        pit_id = pit["id"]
+        # Base query with point in time (pit) paramter.
+        # sort is needed for search_after, so we use scope sort (random choice)
+        query = {
+            "query": elastic_query_str,
+            "sort": [{"scope.keyword": "asc"}],
+            "_source": ["scope", "name"] if not long else ["scope", "name", "did_type", "bytes", "length"],
+            "pit": {"id": pit_id, "keep_alive": "2m"}
+        }
+
+        # Add sorting and pagination
+        if offset:
+            query["from"] = offset
+        size = limit if limit else 10000
+        query["size"] = size
+        search_after = None
+        total_processed = 0
+        try:
+            while True:
+                if search_after:
+                    query["search_after"] = search_after
+                    query.pop("from", None)
+                # Execute search
+                results = self.client.search(body=query)
+                hits = results['hits']['hits']
+                if not hits:
+                    break
+
+                for hit in hits:
+                    did_full = f"{hit['_source']['scope']}:{hit['_source']['name']}"
+                    if did_full not in ignore_dids:
+                        ignore_dids.append(did_full)
+                        if long:
+                            yield {
+                                'scope': (hit['_source']['scope']),
+                                'name': hit['_source']['name'],
+                                'did_type': hit['_source'].get('did_type', 'N/A'),
+                                'bytes': hit['_source'].get('bytes', 'N/A'),
+                                'length': hit['_source'].get('length', 'N/A')
+                            }
+                        else:
+                            yield hit['_source']['name']
+
+                    total_processed += 1
+                    if limit and total_processed >= limit:
+                        break
+
+                # Update search_after for the next iteration
+                search_after = hits[-1]["sort"]
+
+        finally:
+            # Always delete the point in time when done
+            self.client.close_point_in_time(body={"id": pit_id})
+
+        if recursive:
+            raise exception.UnsupportedOperation(f"'{self.plugin_name.lower()}' metadata module does not currently support recursive searches")
+
+    def on_delete(
+        self,
+        scope: "InternalScope",
+        name: str,
+        archive: bool = False,
+        session: "Optional[Session]" = None
+    ) -> None:
+        """
+        Delete a document and optionally archive it.
+
+        :param scope: The scope of the document
+        :param name: The name of the document
+        :param archive: Whether to archive the document before deletion
+        :raises DataIdentifierNotFound: If the DID is not found.
+        :raises RucioException: If an error occurs while setting the metadata.
+        """
+        doc_id = f"{scope}{name}"
+
+        try:
+            doc = self.client.get(index=self.index, id=doc_id)
+
+            if archive:
+                archived_doc = doc['_source']
+                archived_doc['deleted_at'] = datetime.datetime.now(datetime.timezone.utc).isoformat()
+                self.client.index(index=self.archive_index, id=doc_id, body=archived_doc)
+
+            self.client.delete(index=self.index, id=doc_id)
+
+        except elastic_exceptions.NotFoundError as err:
+            raise exception.DataIdentifierNotFound(f"No metadata found for DID '{scope}:{name}' not found") from err
+        except Exception as err:
+            raise exception.RucioException(err)
+
+    def get_metadata_archived(
+        self,
+        scope: "InternalScope",
+        name: str,
+        session: "Optional[Session]" = None
+    ) -> None:
+        """
+        Retrieve archived metadata for a given scope and name.
+
+        :param scope: The scope of the document
+        :param name: The name of the document
+        :return: The archived metadata or None if not found
+        :raises DataIdentifierNotFound: If the DID is not found.
+        :raises RucioException: If an error occurs while setting the metadata.
+        """
+        doc_id = f"{scope}{name}"
+
+        try:
+            doc = self.client.get(index=self.archive_index, id=doc_id)["_source"]
+            return doc
+        except elastic_exceptions.NotFoundError as err:
+            raise exception.DataIdentifierNotFound(f"No metadata found for DID '{scope}:{name}' not found") from err
+        except Exception as err:
+            raise exception.RucioException(err)
+
+    def manages_key(
+        self,
+        key: str,
+        *,
+        session: "Optional[Session]" = None
+    ) -> bool:
+        return True
+
+    def get_plugin_name(self) -> str:
+        """
+        Returns a unique identifier for this plugin. This can be later used for filtering down results to this plugin only.
+
+        :returns: The name of the plugin
+        """
+        return self.plugin_name

rucio/core/did_meta_plugins/filter_engine.py

@@ -49,6 +49,15 @@ OPERATORS_CONVERSION_LUT = {
     "": operator.eq
 }
 
+ELASTIC_OP_MAP = {
+    operator.eq: "=",
+    operator.ne: "!",
+    operator.gt: "gt",
+    operator.lt: "lt",
+    operator.ge: "gte",
+    operator.le: "lte"
+}
+
 # lookup table converting pythonic operators to oracle operators
 ORACLE_OP_MAP = {
     operator.eq: "==",
@@ -348,6 +357,56 @@ class FilterEngine:
 
         return query_str
 
+    def create_elastic_query(
+        self,
+        additional_filters: Optional["Iterable[FilterTuple]"] = None
+    ) -> dict[str, Any]:
+        """
+        Returns a single elastic query dictionary describing the filters expression.
+
+        :param additional_filters: additional filters to be applied to all clauses.
+        :returns: an elastic query dictionary describing the filters expression.
+        """
+
+        additional_filters = additional_filters or []
+        for or_group in self._filters:
+            for _filter in additional_filters:
+                or_group.append(list(_filter))  # type: ignore
+
+        should_clauses = []
+        for or_group in self._filters:
+            bool_query = {
+                "must": [],
+                "must_not": []
+            }
+
+            for and_group in or_group:
+                key, oper, value = and_group
+                key = str(key)
+
+                if isinstance(value, str) and any(char in value for char in ['*', '%']):
+                    if value in ('*', '%', '*', '%'):
+                        bool_query["must"].append({"wildcard": {key: value}})
+                    else:
+                        wildcard_query = {"wildcard": {key: value}}
+                        if oper == operator.eq:
+                            bool_query["must"].append(wildcard_query)
+                        elif oper == operator.ne:
+                            bool_query["must_not"].append(wildcard_query)
+                else:
+                    if oper in [operator.lt, operator.gt, operator.ge, operator.le]:
+                        elsop = ELASTIC_OP_MAP[oper]
+                        bool_query["must"].append({"range": {key: {elsop: value}}})
+                    elif oper == operator.eq:
+                        bool_query["must"].append({"term": {key: value}})
+                    elif oper == operator.ne:
+                        bool_query["must_not"].append({"term": {key: value}})
+
+            should_clauses.append({"bool": bool_query})
+
+        query_expression = {"bool": {"should": should_clauses}}
+        return query_expression
+
     def create_postgres_query(
         self,
         additional_filters: Optional["Iterable[FilterTuple]"] = None,
@@ -551,7 +610,7 @@ class FilterEngine:
         """
         A (more) human readable format of <filters>.
         """
-        operators_conversion_LUT_inv = {op2: op1 for op1, op2 in OPERATORS_CONVERSION_LUT.items()}
+        operators_conversion_lut_inv = {op2: op1 for op1, op2 in OPERATORS_CONVERSION_LUT.items()}
 
         filters = '\n'
         for or_group in self._filters:
@@ -559,10 +618,10 @@ class FilterEngine:
                 key, oper, value = and_group
                 if isinstance(key, InstrumentedAttribute):
                     key = and_group[0].key
-                if operators_conversion_LUT_inv[oper] == "":
+                if operators_conversion_lut_inv[oper] == "":
                     oper = "eq"
                 else:
-                    oper = operators_conversion_LUT_inv[oper]
+                    oper = operators_conversion_lut_inv[oper]
                 if isinstance(value, InstrumentedAttribute):
                     value = and_group[2].key  # type: ignore
                 elif isinstance(value, DIDType):

rucio/core/did_meta_plugins/json_meta.py

@@ -100,9 +100,9 @@ class JSONDidMeta(DidMetaPlugin):
             if row_did_meta.meta:
                 if session.bind.dialect.name in ['oracle', 'sqlite']:
                     # Oracle and sqlite returns a string instead of a dict
-                    existing_meta = json_lib.loads(cast(str, row_did_meta.meta))
+                    existing_meta = json_lib.loads(cast("str", row_did_meta.meta))
                 else:
-                    existing_meta = cast(dict[str, Any], row_did_meta.meta)
+                    existing_meta = cast("dict[str, Any]", row_did_meta.meta)
 
         for key, value in metadata.items():
             existing_meta[key] = value

rucio/core/did_meta_plugins/mongo_meta.py

@@ -28,10 +28,10 @@ if TYPE_CHECKING:
     from sqlalchemy.orm import Session
 
 IMMUTABLE_KEYS = [
-    '_id',              # index key
-    'scope',            # generated on insert
-    'name',             # generated on insert
-    'vo'                # generated on insert
+    '_id',  # index key
+    'scope',  # generated on insert
+    'name',  # generated on insert
+    'vo'  # generated on insert
 ]
 
 
@@ -46,28 +46,40 @@ class MongoDidMeta(DidMetaPlugin):
         password: "Optional[str]" = None,
     ):
         super(MongoDidMeta, self).__init__()
-        if host is None:
-            host = config.config_get('metadata', 'mongo_service_host')
-        if port is None:
-            port = config.config_get_int('metadata', 'mongo_service_port')
-        if db is None:
-            db = config.config_get('metadata', 'mongo_db')
-        if collection is None:
-            collection = config.config_get('metadata', 'mongo_collection')
-
-        if user is None and config.config_has_section("metadata"):
-            user = config.config_get("metadata", "mongo_user", default=None)
-
-        if user is not None:
-            if password is None:
-                password = config.config_get("metadata", "mongo_password")
-            auth = "{user}:{password}@".format(user=user, password=password)
-        else:
-            auth = ""
 
-        self.client = pymongo.MongoClient("mongodb://{auth}{host}:{port}/".format(auth=auth, host=host, port=port))
-        self.db = self.client[db]
-        self.col = self.db[collection]
+        # Validate required parameters.
+        con_params = {
+            'mongo_service_host': host,
+            'mongo_service_port': port,
+            'mongo_db': db,
+            'mongo_collection': collection,
+        }
+
+        for param in con_params:
+            if con_params[param] is None:
+                if config.config_has_option('metadata', param):
+                    con_params[param] = (
+                        config.config_get_int('metadata', param)
+                        if param == 'mongo_service_port'
+                        else config.config_get('metadata', param)
+                    )
+                else:
+                    raise exception.ConnectionParameterNotFound(param)
+
+        if user is None and config.config_has_option('metadata', 'mongo_user'):
+            user = config.config_get('metadata', 'mongo_user', default=None)
+
+        if password is None and config.config_has_option('metadata', 'mongo_password'):
+            password = config.config_get('metadata', 'mongo_password')
+
+        # Set the auth (fallback to an anonymous connection if either user or password is not defined).
+        auth = "" if not user or not password else f"{user}:{password}@"
+
+        self.client = pymongo.MongoClient(
+            f"mongodb://{auth}{con_params['mongo_service_host']}:{con_params['mongo_service_port']}/"
+        )
+        self.db = self.client[con_params['mongo_db']]
+        self.col = self.db[con_params['mongo_collection']]
 
         self.plugin_name = "MONGO"
 
@@ -175,9 +187,10 @@ class MongoDidMeta(DidMetaPlugin):
 
         if recursive:
             # TODO: possible, but requires retrieving the results of a concurrent sqla query to call list_content on for datasets and containers
-            raise exception.UnsupportedOperation("'{}' metadata module does not currently support recursive searches".format(
-                self.plugin_name.lower()
-            ))
+            raise exception.UnsupportedOperation(
+                "'{}' metadata module does not currently support recursive searches".format(
+                    self.plugin_name.lower()
+                ))
 
         if long:
             query_result = self.col.find(mongo_query_str)
@@ -185,7 +198,7 @@ class MongoDidMeta(DidMetaPlugin):
                 query_result = query_result.limit(limit)
             for did in query_result:
                 did_full = "{}:{}".format(did['scope'], did['name'])
-                if did_full not in ignore_dids:         # aggregating recursive queries may contain duplicate DIDs
+                if did_full not in ignore_dids:  # aggregating recursive queries may contain duplicate DIDs
                     ignore_dids.add(did_full)
                     yield {
                         'scope': InternalScope(did['scope']),
@@ -200,7 +213,7 @@ class MongoDidMeta(DidMetaPlugin):
                 query_result = query_result.limit(limit)
             for did in query_result:
                 did_full = "{}:{}".format(did['scope'], did['name'])
-                if did_full not in ignore_dids:         # aggregating recursive queries may contain duplicate DIDs
+                if did_full not in ignore_dids:  # aggregating recursive queries may contain duplicate DIDs
                     ignore_dids.add(did_full)
                     yield did['name']