rucio 35.7.0__py3-none-any.whl

rucio/core/did_meta_plugins/postgres_meta.py

@@ -0,0 +1,316 @@
+# 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
+import operator
+from typing import TYPE_CHECKING
+
+import psycopg2
+import psycopg2.extras
+
+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
+
+
+class ExternalPostgresDidMeta(DidMetaPlugin):
+    # TODO: column-based plugin? mixed-mode (json & columns)?
+    pass
+
+
+class ExternalPostgresJSONDidMeta(DidMetaPlugin):
+    def __init__(self, host=None, port=None, db=None, user=None, password=None, db_schema=None, table=None,
+                 table_is_managed=None, table_column_vo=None, table_column_scope=None, table_column_name=None,
+                 table_column_data=None):
+        super(ExternalPostgresJSONDidMeta, self).__init__()
+        if host is None:
+            host = config.config_get('metadata', 'postgres_service_host')
+        if port is None:
+            port = config.config_get_int('metadata', 'postgres_service_port')
+        if db is None:
+            db = config.config_get('metadata', 'postgres_db')
+        if user is None:
+            user = config.config_get('metadata', 'postgres_user')
+        if password is None:
+            password = config.config_get('metadata', 'postgres_password')
+        if db_schema is None:
+            db_schema = config.config_get('metadata', 'postgres_db_schema', default='public')
+        if table is None:
+            table = config.config_get('metadata', 'postgres_table', default='dids')
+        if table_is_managed is None:
+            table_is_managed = config.config_get_bool('metadata', 'postgres_table_is_managed', default=False)
+        if table_column_vo is None:
+            table_column_vo = config.config_get('metadata', 'postgres_table_column_vo', default='vo')
+        if table_column_scope is None:
+            table_column_scope = config.config_get('metadata', 'postgres_table_column_scope', default='scope')
+        if table_column_name is None:
+            table_column_name = config.config_get('metadata', 'postgres_table_column_name', default='name')
+        if table_column_data is None:
+            table_column_data = config.config_get('metadata', 'postgres_table_column_data', default='data')
+
+        self.fixed_table_columns = {
+            'vo': table_column_vo,
+            'scope': table_column_scope,
+            'name': table_column_name
+        }
+        self.jsonb_column = table_column_data
+
+        self.table = table
+        self.client = psycopg2.connect(
+            host=host,
+            port=port,
+            database=db,
+            user=user,
+            password=password)
+
+        # set search_path to include database schema by default
+        cur = self.client.cursor()
+        statement = "SET search_path TO {};".format(db_schema)
+        cur.execute(statement)
+        cur.close()
+
+        if not table_is_managed:                    # not managed by Rucio, so just verify table schema
+            self._verify_table_schema(table_column_vo, table_column_scope, table_column_name, table_column_data)
+        else:                                       # managed by Rucio, create a metadata table if it doesn't exist
+            self._try_create_metadata_table()
+
+        self.plugin_name = "POSTGRES_JSON"
+
+    def _try_create_metadata_table(self):
+        """
+        Try to create a metadata table.
+        """
+        table_clauses = (
+            ("id", "bigint", "NOT NULL", "GENERATED ALWAYS AS IDENTITY"),
+            ("vo", "varchar", "NOT NULL"),
+            ("scope", "varchar", "NOT NULL"),
+            ("name", "varchar", "NOT NULL"),
+            ("data", "jsonb", "DEFAULT", "'{}'::jsonb"),
+            ("UNIQUE", "(scope, name)")  # unique scope+name table constraint, required for ON CONFLICT
+        )
+        statement = "CREATE TABLE IF NOT EXISTS {} ({})".format(
+            self.table,
+            ', '.join([' '.join(clause) for clause in table_clauses]))
+
+        cur = self.client.cursor()
+        cur.execute(statement)
+        cur.close()
+        self.client.commit()
+
+    def _verify_table_schema(self, table_column_vo, table_column_scope, table_column_name, table_column_data):
+        """
+        Rudimentary verification that the metadata table schema meets the requirements for the plugin.
+
+        Should be called when using externally managed database tables as a sanity check.
+
+        :param table_column_vo: The table column used for the vo
+        :param table_column_scope: The table column used for the scope
+        :param table_column_name: The table column used for the name
+        :param table_column_data: The table column used for the data
+        :raises: MetadataSchemaMismatchError
+        """
+        # Check mandatory columns are of right data type and have the right nullable qualifier.
+        statement = "SELECT column_name, data_type, is_nullable " \
+                    "FROM INFORMATION_SCHEMA.COLUMNS where table_name = '{}';".format(self.table)
+        cur = self.client.cursor()
+        cur.execute(statement)
+        existing_table_columns = cur.fetchall()
+        cur.close()
+
+        mandatory_column_specifications = [
+            (table_column_vo, "character varying", "NO"),
+            (table_column_scope, "character varying", "NO"),
+            (table_column_name, "character varying", "NO"),
+            (table_column_data, "jsonb", "YES")
+        ]
+        for specification in mandatory_column_specifications:
+            if specification not in existing_table_columns:
+                raise exception.MetadataSchemaMismatchError(
+                    "mandatory table column {} does not match that defined in the required table schema {}".format(
+                        specification, existing_table_columns))
+
+        # Check required table constraints exist.
+        statement = "SELECT con.contype AS constraint_type, " \
+                    "(SELECT array_agg(att.attname) FROM pg_attribute att " \
+                    " INNER JOIN unnest(con.conkey) unnest(conkey) ON unnest.conkey = att.attnum " \
+                    " WHERE att.attrelid = con.conrelid) AS columns " \
+                    "FROM pg_constraint con " \
+                    "INNER JOIN pg_class rel ON rel.oid = con.conrelid " \
+                    "INNER JOIN pg_namespace nsp ON nsp.oid = rel.relnamespace " \
+                    "WHERE rel.relname = '{}';".format(self.table)
+        cur = self.client.cursor()
+        cur.execute(statement)
+        existing_table_constraints = cur.fetchall()  # list of (constraint_type, [columns])
+        cur.close()
+
+        mandatory_table_constraints = [
+            ("u", [table_column_scope, table_column_name]),  # unique scope+name table constraint
+        ]
+        for constraint in mandatory_table_constraints:
+            if constraint not in existing_table_constraints:
+                raise exception.MetadataSchemaMismatchError(
+                    "mandatory table constraint {} does not match that defined in the required table schema {}".format(
+                        constraint, len(existing_table_constraints)))
+
+    def _drop_metadata_table(self):
+        statement = "DROP TABLE IF EXISTS {};".format(self.table)
+        cur = self.client.cursor()
+        cur.execute(statement)
+        cur.close()
+        self.client.commit()
+
+    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
+        """
+        statement = "SELECT data from {} ".format(self.table) + \
+                    "WHERE scope='{}' AND name='{}';".format(scope.internal, name)
+        cur = self.client.cursor()
+        cur.execute(statement)
+        metadata = cur.fetchone()
+        cur.close()
+
+        if not metadata:
+            raise exception.DataIdentifierNotFound("No metadata found for did '{}:{}".format(scope, name))
+
+        return metadata[0]
+
+    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
+        """
+        # upsert metadata
+        statement = "INSERT INTO {} (scope, name, vo, data) ".format(self.table) + \
+                    "VALUES ('{}', '{}', '{}', '{}') ".format(scope.external, name, scope.vo, json.dumps(metadata)) + \
+                    "ON CONFLICT (scope, name) DO UPDATE set data = {}.data || EXCLUDED.data;".format(self.table)
+        cur = self.client.cursor()
+        cur.execute(statement)
+        cur.close()
+        self.client.commit()
+
+    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
+        :param session: the database session in use
+        """
+        statement = "UPDATE {} ".format(self.table) + \
+                    "SET data = {}.data - '{}';".format(self.table, key)
+        cur = self.client.cursor()
+        cur.execute(statement)
+        cur.close()
+        self.client.commit()
+
+    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]
+
+        try:
+            # instantiate fe and create postgres query
+            fe = FilterEngine(filters, model_class=None, strict_coerce=False)
+            postgres_query_str = fe.create_postgres_query(
+                additional_filters=[
+                    ('scope', operator.eq, scope.internal),
+                    ('vo', operator.eq, scope.vo)
+                ],
+                fixed_table_columns=self.fixed_table_columns,
+                jsonb_column=self.jsonb_column
+            )
+        except Exception as e:
+            raise exception.DataIdentifierNotFound(e)
+
+        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())
+            )
+
+        statement = "SELECT * FROM {} WHERE {} ".format(self.table, postgres_query_str)
+        if limit:
+            statement += "LIMIT {}".format(limit)
+        cur = self.client.cursor(cursor_factory=psycopg2.extras.RealDictCursor)
+        cur.execute(statement)
+        query_result = cur.fetchall()
+        cur.close()
+
+        if long:
+            for row in query_result:
+                did = "{}:{}".format(row['scope'], row['name'])
+                if did not in ignore_dids:         # aggregating recursive queries may contain duplicate DIDs
+                    ignore_dids.add(did)
+                    yield {
+                        'scope': InternalScope(row['scope']),
+                        'name': row['name'],
+                        'did_type': "N/A",
+                        'bytes': "N/A",
+                        'length': "N/A"
+                    }
+        else:
+            for row in query_result:
+                did = "{}:{}".format(row['scope'], row['name'])
+                if did not in ignore_dids:         # aggregating recursive queries may contain duplicate DIDs
+                    ignore_dids.add(did)
+                    yield row['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

rucio/core/dirac.py

@@ -0,0 +1,237 @@
+# 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 re
+from json import loads
+from json.decoder import JSONDecodeError
+from typing import TYPE_CHECKING, Any, Optional
+
+from sqlalchemy import and_, select
+from sqlalchemy.exc import NoResultFound
+
+from rucio.common.config import config_get
+from rucio.common.exception import ConfigNotFound, InvalidType, RucioException, UnsupportedOperation
+from rucio.common.types import InternalAccount, InternalScope
+from rucio.common.utils import extract_scope
+from rucio.core.did import add_did, attach_dids_to_dids
+from rucio.core.replica import add_replicas
+from rucio.core.rule import add_rule, list_rules, update_rule
+from rucio.core.scope import list_scopes
+from rucio.db.sqla import models
+from rucio.db.sqla.constants import DIDType
+from rucio.db.sqla.session import read_session, transactional_session
+
+if TYPE_CHECKING:
+    from collections.abc import Iterable
+
+    from sqlalchemy.orm import Session
+
+
+@read_session
+def _exists(
+    scope: str,
+    name: str,
+    *,
+    session: "Session"
+) -> tuple[bool, Optional[DIDType]]:
+    """
+    Check if the did exists
+
+    :scope: The scope
+    :name: The name
+    :session: The session used
+    """
+    try:
+        stmt = select(
+            models.DataIdentifier.did_type
+        ).with_hint(
+            models.DataIdentifier,
+            "INDEX(DIDS DIDS_PK)",
+            'oracle'
+        ).where(
+            and_(models.DataIdentifier.scope == scope,
+                 models.DataIdentifier.name == name)
+        )
+        return True, session.execute(stmt).scalar_one()
+    except NoResultFound:
+        return False, None
+
+
+@transactional_session
+def add_files(
+    lfns: "Iterable[dict[str, Any]]",
+    account: str,
+    ignore_availability: bool,
+    parents_metadata: Optional[dict[str, Any]] = None,
+    vo: str = 'def',
+    *,
+    session: "Session"
+) -> None:
+    """
+    Bulk add files :
+    - Create the file and replica.
+    - If doesn't exist create the dataset containing the file as well as a rule on the dataset on ANY sites.
+    - Create all the ascendants of the dataset if they do not exist
+
+    :param lfns: List of lfn (dictionary {'lfn': <lfn>, 'rse': <rse>, 'bytes': <bytes>, 'adler32': <adler32>, 'guid': <guid>, 'pfn': <pfn>}
+    :param issuer: The issuer account.
+    :param ignore_availability: A boolean to ignore blocklisted sites.
+    :param parents_metadata: Metadata for selected hierarchy DIDs. (dictionary {'lpn': {key : value}})
+    :param vo: The VO to act on
+    :param session: The session used
+    """
+    rule_extension_list = []
+    attachments = []
+    if not parents_metadata:
+        parents_metadata = {}
+    # The list of scopes is necessary for the extract_scope
+    filter_ = {'scope': InternalScope(scope='*', vo=vo)}
+    scopes = list_scopes(filter_=filter_, session=session)
+    scopes = [scope.external for scope in scopes]
+    exist_lfn = []
+    try:
+        config_lifetime: str = config_get(section='dirac', option='lifetime', default='{}', session=session)
+        lifetime_dict: dict = loads(config_lifetime)
+    except ConfigNotFound:
+        lifetime_dict = {}
+    except JSONDecodeError as err:
+        raise InvalidType('Problem parsing lifetime option in dirac section : %s' % str(err))
+    except Exception as err:
+        raise RucioException(str(err))
+
+    for lfn in lfns:
+        # First check if the file exists
+        filename = lfn['lfn']
+        lfn_scope, _ = extract_scope(filename, scopes)
+        lfn_scope = InternalScope(lfn_scope, vo=vo)
+
+        exists, did_type = _exists(lfn_scope, filename)
+        if exists:
+            continue
+
+        # Get all the ascendants of the file
+        lfn_split = filename.split('/')
+        lpns = ["/".join(lfn_split[:idx]) for idx in range(2, len(lfn_split))]
+        lpns.reverse()
+        print(lpns)
+
+        # The parent must be a dataset. Register it as well as the rule
+        dsn_name = lpns[0]
+        dsn_scope, _ = extract_scope(dsn_name, scopes)
+        dsn_scope = InternalScope(dsn_scope, vo=vo)
+        dsn_meta = parents_metadata.get(dsn_name, {})
+
+        # Compute lifetime
+        lifetime = None
+        if dsn_scope in lifetime_dict:
+            lifetime = lifetime_dict[dsn_scope.external]
+        else:
+            for pattern in lifetime_dict:
+                if dsn_scope.external and re.match(pattern, str(dsn_scope.external)):
+                    lifetime = lifetime_dict[pattern]
+                    break
+
+        exists, did_type = _exists(dsn_scope, dsn_name)
+        if exists and did_type == DIDType.CONTAINER:
+            raise UnsupportedOperation('Cannot create %s as dataset' % dsn_name)
+        if (dsn_name not in exist_lfn) and not exists:
+            print('Will create %s' % dsn_name)
+            # to maintain a compatibility between master and LTS-1.26 branches remove keywords for first 3 arguments
+            add_did(dsn_scope,
+                    dsn_name,
+                    DIDType.DATASET,
+                    account=InternalAccount(account, vo=vo),
+                    statuses=None,
+                    meta=dsn_meta,
+                    rules=[{'copies': 1, 'rse_expression': 'ANY=true', 'weight': None, 'account': InternalAccount(account, vo=vo), 'lifetime': lifetime, 'grouping': 'NONE'}],
+                    lifetime=None,
+                    dids=None,
+                    rse_id=None,
+                    session=session)
+            exist_lfn.append(dsn_name)
+            parent_name = lpns[1]
+            parent_scope, _ = extract_scope(parent_name, scopes)
+            parent_scope = InternalScope(parent_scope, vo=vo)
+            attachments.append({'scope': parent_scope, 'name': parent_name, 'dids': [{'scope': dsn_scope, 'name': dsn_name}]})
+            rule_extension_list.append((dsn_scope, dsn_name))
+        if lifetime and (dsn_scope, dsn_name) not in rule_extension_list:
+            # Reset the lifetime of the rule to the configured value
+            rule = [rul for rul in list_rules({'scope': dsn_scope, 'name': dsn_name, 'account': InternalAccount(account, vo=vo)}, session=session) if rul['rse_expression'] == 'ANY=true']
+            if rule:
+                update_rule(rule[0]['id'], options={'lifetime': lifetime}, session=session)
+            rule_extension_list.append((dsn_scope, dsn_name))
+
+        # Register the file
+        rse_id = lfn.get('rse_id', None)
+        if not rse_id:
+            raise InvalidType('Missing rse_id')
+        bytes_ = lfn.get('bytes', None)
+        guid = lfn.get('guid', None)
+        adler32 = lfn.get('adler32', None)
+        pfn = lfn.get('pfn', None)
+        meta = lfn.get('meta', {})
+        files = {'scope': lfn_scope, 'name': filename, 'bytes': bytes_, 'adler32': adler32, 'meta': meta}
+        if pfn:
+            files['pfn'] = str(pfn)
+        if guid:
+            files['meta']['guid'] = guid
+        add_replicas(rse_id=rse_id,
+                     files=[files],
+                     dataset_meta=None,
+                     account=InternalAccount(account, vo=vo),
+                     ignore_availability=ignore_availability,
+                     session=session)
+        add_rule(dids=[{'scope': lfn_scope, 'name': filename}],
+                 account=InternalAccount(account, vo=vo),
+                 copies=1,
+                 rse_expression=lfn['rse'],
+                 grouping=None,
+                 weight=None,
+                 lifetime=86400,
+                 locked=None,
+                 subscription_id=None,
+                 session=session)
+        attachments.append({'scope': dsn_scope, 'name': dsn_name, 'dids': [{'scope': lfn_scope, 'name': filename}]})
+
+        # Now loop over the ascendants of the dataset and created them
+        for lpn in lpns[1:]:
+            child_scope, _ = extract_scope(lpn, scopes)
+            child_scope = InternalScope(child_scope, vo=vo)
+            exists, did_type = _exists(child_scope, lpn)
+            child_meta = parents_metadata.get(lpn, {})
+            if exists and did_type == DIDType.DATASET:
+                raise UnsupportedOperation('Cannot create %s as container' % lpn)
+            if (lpn not in exist_lfn) and not exists:
+                print('Will create %s' % lpn)
+                add_did(child_scope,
+                        lpn,
+                        DIDType.CONTAINER,
+                        account=InternalAccount(account, vo=vo),
+                        statuses=None,
+                        meta=child_meta,
+                        rules=None,
+                        lifetime=None,
+                        dids=None,
+                        rse_id=None,
+                        session=session)
+                exist_lfn.append(lpn)
+                parent_name = lpns[lpns.index(lpn) + 1]
+                parent_scope, _ = extract_scope(parent_name, scopes)
+                parent_scope = InternalScope(parent_scope, vo=vo)
+                attachments.append({'scope': parent_scope, 'name': parent_name, 'dids': [{'scope': child_scope, 'name': lpn}]})
+    # Finally attach everything
+    attach_dids_to_dids(attachments,
+                        account=InternalAccount(account, vo=vo),
+                        ignore_duplicate=True,
+                        session=session)