rucio 37.0.0rc1__py3-none-any.whl

rucio/core/did_meta_plugins/json_meta.py

@@ -0,0 +1,240 @@
+# 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.
+
+import json as json_lib
+import operator
+from typing import TYPE_CHECKING, Any, cast
+
+from sqlalchemy import and_, select
+from sqlalchemy.exc import DataError, NoResultFound
+
+from rucio.common import exception
+from rucio.core.did_meta_plugins.did_meta_plugin_interface import DidMetaPlugin
+from rucio.core.did_meta_plugins.filter_engine import FilterEngine
+from rucio.db.sqla import models
+from rucio.db.sqla.constants import DIDType
+from rucio.db.sqla.session import read_session, stream_session, transactional_session
+from rucio.db.sqla.util import json_implemented
+
+if TYPE_CHECKING:
+    from sqlalchemy.orm import Session
+
+
+class JSONDidMeta(DidMetaPlugin):
+    """
+    A plugin to store DID metadata on a table on the relational database, using JSON blobs
+    """
+
+    def __init__(self):
+        super(JSONDidMeta, self).__init__()
+        self.plugin_name = "JSON"
+
+    @read_session
+    def get_metadata(self, scope, name, *, session: "Session"):
+        """
+        Get data identifier metadata (JSON)
+
+        :param scope: The scope name.
+        :param name: The data identifier name.
+        :param session: The database session in use.
+        """
+        if not json_implemented(session=session):
+            raise NotImplementedError
+
+        try:
+            stmt = select(
+                models.DidMeta
+            ).where(
+                and_(models.DidMeta.scope == scope,
+                     models.DidMeta.name == name)
+            )
+            row = session.execute(stmt).scalar_one()
+            meta = getattr(row, 'meta')
+            return json_lib.loads(meta) if session.bind.dialect.name in ['oracle', 'sqlite'] else meta
+        except NoResultFound:
+            return {}
+
+    @transactional_session
+    def set_metadata(self, scope, name, key, value, recursive=False, *, session: "Session"):
+        self.set_metadata_bulk(scope=scope, name=name, metadata={key: value}, recursive=recursive, session=session)
+
+    @transactional_session
+    def set_metadata_bulk(self, scope, name, metadata, recursive=False, *, session: "Session"):
+        if not json_implemented(session=session):
+            raise NotImplementedError
+
+        stmt = select(
+            models.DataIdentifier
+        ).where(
+            and_(models.DataIdentifier.scope == scope,
+                 models.DataIdentifier.name == name)
+        )
+        if session.execute(stmt).one_or_none() is None:
+            raise exception.DataIdentifierNotFound("Data identifier '%s:%s' not found" % (scope, name))
+
+        stmt = select(
+            models.DidMeta
+        ).where(
+            and_(models.DidMeta.scope == scope,
+                 models.DidMeta.name == name)
+        )
+        row_did_meta = session.execute(stmt).scalar_one_or_none()
+        if row_did_meta is None:
+            # Add metadata column to new table (if not already present)
+            row_did_meta = models.DidMeta(scope=scope, name=name)
+            row_did_meta.save(session=session, flush=False)
+
+        existing_meta = {}
+        if hasattr(row_did_meta, 'meta'):
+            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))
+                else:
+                    existing_meta = cast("dict[str, Any]", row_did_meta.meta)
+
+        for key, value in metadata.items():
+            existing_meta[key] = value
+
+        row_did_meta.meta = None
+        session.flush()
+
+        # Oracle insert takes a string as input
+        if session.bind.dialect.name in ['oracle', 'sqlite']:
+            existing_meta = json_lib.dumps(existing_meta)
+
+        row_did_meta.meta = existing_meta
+        row_did_meta.save(session=session, flush=True)
+
+    @transactional_session
+    def delete_metadata(self, scope, name, key, *, session: "Session"):
+        """
+        Delete a key from the metadata column
+
+        :param scope: the scope of did
+        :param name: the name of the did
+        :param key: the key to be deleted
+        :param session: The database session in use.
+        """
+        if not json_implemented(session=session):
+            raise NotImplementedError
+
+        try:
+            stmt = select(
+                models.DidMeta
+            ).where(
+                and_(models.DidMeta.scope == scope,
+                     models.DidMeta.name == name)
+            )
+            row = session.execute(stmt).scalar_one()
+            existing_meta = getattr(row, 'meta')
+            # Oracle returns a string instead of a dict
+            if session.bind.dialect.name in ['oracle', 'sqlite'] and existing_meta is not None:
+                existing_meta = json_lib.loads(existing_meta)
+
+            if key not in existing_meta:
+                raise exception.KeyNotFound(key)
+
+            existing_meta.pop(key, None)
+
+            row.meta = None
+            session.flush()
+
+            # Oracle insert takes a string as input
+            if session.bind.dialect.name in ['oracle', 'sqlite']:
+                existing_meta = json_lib.dumps(existing_meta)
+
+            row.meta = existing_meta
+        except NoResultFound:
+            raise exception.DataIdentifierNotFound(f"Key not found for data identifier '{scope}:{name}'")
+
+    @stream_session
+    def list_dids(self, scope, filters, did_type='collection', ignore_case=False, limit=None,
+                  offset=None, long=False, recursive=False, ignore_dids=None, *, session: "Session"):
+        if not json_implemented(session=session):
+            raise NotImplementedError
+
+        if not ignore_dids:
+            ignore_dids = set()
+
+        # backwards compatibility for filters as single {}.
+        if isinstance(filters, dict):
+            filters = [filters]
+
+        # instantiate fe and create sqla query, note that coercion to a model keyword
+        # is not appropriate here as the filter words are stored in a single json column.
+        fe = FilterEngine(filters, model_class=models.DidMeta, strict_coerce=False)
+        stmt = fe.create_sqla_query(
+            additional_model_attributes=[
+                models.DidMeta.scope,
+                models.DidMeta.name
+            ], additional_filters=[
+                (models.DidMeta.scope, operator.eq, scope)
+            ],
+            json_column=models.DidMeta.meta,
+            session=session
+        )
+
+        if limit:
+            stmt = stmt.limit(
+                limit
+            )
+        if recursive:
+            from rucio.core.did import list_content
+
+            # Get attached DIDs and save in list because query has to be finished before starting a new one in the recursion
+            collections_content = []
+            for did in session.execute(stmt).yield_per(100):
+                if (did.did_type == DIDType.CONTAINER or did.did_type == DIDType.DATASET):
+                    collections_content += [d for d in list_content(scope=did.scope, name=did.name)]
+
+            # Replace any name filtering with recursed DID names.
+            for did in collections_content:
+                for or_group in filters:
+                    or_group['name'] = did['name']
+                for result in self.list_dids(scope=did['scope'], filters=filters, recursive=True, did_type=did_type, limit=limit, offset=offset,
+                                             long=long, ignore_dids=ignore_dids, session=session):
+                    yield result
+
+        try:
+            for did in session.execute(stmt).yield_per(5):                  # don't unpack this as it makes it dependent on query return order!
+                if long:
+                    did_full = "{}:{}".format(did.scope, did.name)
+                    if did_full not in ignore_dids:         # concatenating results of OR clauses may contain duplicate DIDs if query result sets not mutually exclusive.
+                        ignore_dids.add(did_full)
+                        yield {
+                            'scope': did.scope,
+                            'name': did.name,
+                            'did_type': None,               # not available with JSON plugin
+                            'bytes': None,                  # not available with JSON plugin
+                            'length': None                  # not available with JSON plugin
+                        }
+                else:
+                    did_full = "{}:{}".format(did.scope, did.name)
+                    if did_full not in ignore_dids:         # concatenating results of OR clauses may contain duplicate DIDs if query result sets not mutually exclusive.
+                        ignore_dids.add(did_full)
+                        yield did.name
+        except DataError as e:
+            raise exception.InvalidMetadata("Database query failed: {}. This can be raised when the datatype of a key is inconsistent between dids.".format(e))
+
+    @read_session
+    def manages_key(self, key, *, session: "Session"):
+        return json_implemented(session=session)
+
+    def get_plugin_name(self):
+        """
+        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/mongo_meta.py

@@ -0,0 +1,229 @@
+# 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.
+
+import operator
+from typing import TYPE_CHECKING
+
+import pymongo
+
+from rucio.common import config, exception
+from rucio.common.types import InternalScope
+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 typing import Optional
+
+    from sqlalchemy.orm import Session
+
+IMMUTABLE_KEYS = [
+    '_id',  # index key
+    'scope',  # generated on insert
+    'name',  # generated on insert
+    'vo'  # generated on insert
+]
+
+
+class MongoDidMeta(DidMetaPlugin):
+    def __init__(
+        self,
+        host: "Optional[str]" = None,
+        port: "Optional[int]" = None,
+        db: "Optional[str]" = None,
+        collection: "Optional[str]" = None,
+        user: "Optional[str]" = None,
+        password: "Optional[str]" = None,
+    ):
+        super(MongoDidMeta, self).__init__()
+
+        # 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"
+
+    def drop_database(self):
+        self.client.drop_database(self.db.name)
+
+    def get_metadata(self, scope, name, *, session: "Optional[Session]" = None):
+        """
+        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
+        """
+        # get first document with this did == _id
+        doc = self.col.find_one({
+            "_id": "{}:{}".format(scope.internal, name)
+        })
+
+        # pop immutable keys
+        for key in IMMUTABLE_KEYS:
+            if key in doc:
+                doc.pop(key)
+
+        if not doc:
+            raise exception.DataIdentifierNotFound(f"No metadata found for did '{scope}:{name}'")
+        return doc
+
+    def set_metadata(self, scope, name, key, value, recursive=False, *, session: "Optional[Session]" = 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
+        """
+        self.set_metadata_bulk(scope=scope, name=name, metadata={key: value}, recursive=recursive, session=session)
+
+    def set_metadata_bulk(self, scope, name, metadata, recursive=False, *, session: "Optional[Session]" = None):
+        """
+        Bulk set metadata keys.
+
+        :param scope: the scope of did
+        :param name: the name of the did
+        :param metadata: dictionary of metadata keypairs to be added
+        :param recursive: recurse into DIDs (not supported)
+        :param session: The database session in use
+        """
+        # pop immutable keys
+        for key in IMMUTABLE_KEYS:
+            if key in metadata:
+                metadata.pop(key)
+
+        # set first document with did == _id
+        self.col.update_one(
+            {
+                "_id": "{}:{}".format(scope.internal, name)
+            },
+            {
+                '$set': metadata,
+                '$setOnInsert': {
+                    'scope': "{}".format(scope.external),
+                    'vo': "{}".format(scope.vo),
+                    'name': "{}".format(name)
+                }
+            },
+            upsert=True
+        )
+
+    def delete_metadata(self, scope, name, key, *, session: "Optional[Session]" = 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
+        """
+        meta = {key: ""}
+        try:
+            self.col.update_one({"_id": "{}:{}".format(scope.internal, name)}, {'$unset': meta})
+        except Exception as e:
+            raise exception.DataIdentifierNotFound(e)
+
+    def list_dids(self, scope, filters, did_type='collection', ignore_case=False, limit=None,
+                  offset=None, long=False, recursive=False, ignore_dids=None, *, session: "Optional[Session]" = None):
+        if not ignore_dids:
+            ignore_dids = set()
+
+        # backwards compatibility for filters as single {}.
+        if isinstance(filters, dict):
+            filters = [filters]
+
+        # instantiate fe and create mongo query
+        fe = FilterEngine(filters, model_class=None, strict_coerce=False)
+        mongo_query_str = fe.create_mongo_query(
+            additional_filters=[
+                ('scope', operator.eq, scope.internal),
+                ('vo', operator.eq, scope.vo)
+            ]
+        )
+
+        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()
+                ))
+
+        if long:
+            query_result = self.col.find(mongo_query_str)
+            if limit:
+                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
+                    ignore_dids.add(did_full)
+                    yield {
+                        'scope': InternalScope(did['scope']),
+                        'name': did['name'],
+                        'did_type': "N/A",
+                        'bytes': "N/A",
+                        'length': "N/A"
+                    }
+        else:
+            query_result = self.col.find(mongo_query_str)
+            if limit:
+                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
+                    ignore_dids.add(did_full)
+                    yield did['name']
+
+    def manages_key(self, key, *, session: "Optional[Session]" = None):
+        return True
+
+    def get_plugin_name(self):
+        """
+        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