rucio 32.8.6__py3-none-any.whl

rucio/core/account_limit.py

@@ -0,0 +1,357 @@
+# -*- coding: utf-8 -*-
+# 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.
+
+from typing import TYPE_CHECKING
+
+from sqlalchemy.orm.exc import NoResultFound
+from sqlalchemy.sql import func, select, literal
+from sqlalchemy.sql.expression import and_, or_
+
+from rucio.core.account import get_all_rse_usages_per_account
+from rucio.core.rse import get_rse_name
+from rucio.core.rse_expression_parser import parse_expression
+from rucio.db.sqla import models
+from rucio.db.sqla.session import read_session, transactional_session
+
+if TYPE_CHECKING:
+    from sqlalchemy.orm import Session
+
+
+@read_session
+def get_rse_account_usage(rse_id, *, session: "Session"):
+    """
+    Returns the account limit and usage for all accounts on a RSE.
+
+    :param rse_id:   The id of the RSE.
+    :param session:  Database session in use.
+    :return:         List of dictionaries.
+    """
+    result = []
+    stmt = select(
+        models.RSE.id.label('rse_id'),
+        models.RSE.vo.label('rse_vo'),
+        models.RSE.rse.label('rse_name'),
+        models.Account.account,
+        models.AccountUsage.files.label('used_files'),
+        models.AccountUsage.bytes.label('used_bytes'),
+        models.AccountLimit.bytes.label('quota_bytes'),
+    ).where(
+        models.RSE.id == rse_id
+    ).join_from(
+        models.RSE,
+        models.Account,
+        or_(
+            and_(
+                models.RSE.vo == 'def',
+                models.Account.account.notlike('%@%')
+            ),
+            and_(
+                models.RSE.vo != 'def',
+                models.Account.account.like(literal('%@') + models.RSE.vo)
+            )
+        )
+    ).outerjoin(
+        models.AccountUsage,
+        and_(
+            models.AccountUsage.account == models.Account.account,
+            models.AccountUsage.rse_id == models.RSE.id
+        )
+    ).outerjoin(
+        models.AccountLimit,
+        and_(
+            models.AccountLimit.account == models.Account.account,
+            models.AccountLimit.rse_id == models.RSE.id
+        )
+    )
+    for row in session.execute(stmt):
+        if row.rse_vo != row.account.vo:
+            continue
+        result.append({
+            'rse_id': row.rse_id,
+            'rse': row.rse_name,
+            'account': row.account,
+            'used_files': row.used_files if row.used_files is not None else 0,
+            'used_bytes': row.used_bytes if row.used_bytes is not None else 0,
+            'quota_bytes': row.quota_bytes
+        })
+    return result
+
+
+@read_session
+def get_global_account_limits(account=None, *, session: "Session"):
+    """
+    Returns the global account limits for the account.
+
+    :param account:  Account to check the limit for.
+    :param session:  Database session in use.
+    :return:         Dict {'MOCK': {'resolved_rses': ['MOCK'], 'limit': 10, 'resolved_rse_ids': [123]}}.
+    """
+    if account:
+        global_account_limits = session.query(models.AccountGlobalLimit).filter_by(account=account).all()
+    else:
+        global_account_limits = session.query(models.AccountGlobalLimit).all()
+
+    resolved_global_account_limits = {}
+    for limit in global_account_limits:
+        if account:
+            resolved_rses = parse_expression(limit['rse_expression'], filter_={'vo': account.vo}, session=session)
+        else:
+            resolved_rses = parse_expression(limit['rse_expression'], session=session)
+        limit_in_bytes = limit['bytes']
+        if limit_in_bytes == -1:
+            limit_in_bytes = float('inf')
+        resolved_global_account_limits[limit['rse_expression']] = {
+            'resolved_rses': [resolved_rse['rse'] for resolved_rse in resolved_rses],
+            'resolved_rse_ids': [resolved_rse['id'] for resolved_rse in resolved_rses],
+            'limit': limit_in_bytes
+        }
+    return resolved_global_account_limits
+
+
+@read_session
+def get_global_account_limit(account, rse_expression, *, session: "Session"):
+    """
+    Returns the global account limit for the account on the rse expression.
+
+    :param account:         Account to check the limit for.
+    :param rse_expression:  RSE expression to check the limit for.
+    :param session:         Database session in use.
+    :return:                Limit in Bytes.
+    """
+    try:
+        global_account_limit = session.query(models.AccountGlobalLimit).filter_by(account=account, rse_expression=rse_expression).one()
+        if global_account_limit.bytes == -1:
+            return float("inf")
+        else:
+            return global_account_limit.bytes
+    except NoResultFound:
+        return None
+
+
+@read_session
+def get_local_account_limit(account, rse_id, *, session: "Session"):
+    """
+    Returns the account limit for the account on the rse.
+
+    :param account:  Account to check the limit for.
+    :param rse_id:   RSE id to check the limit for.
+    :param session:  Database session in use.
+    :return:         Limit in Bytes.
+    """
+    try:
+        account_limit = session.query(models.AccountLimit).filter(models.AccountLimit.account == account,
+                                                                  models.AccountLimit.rse_id == rse_id).one()
+        if account_limit.bytes == -1:
+            return float("inf")
+        else:
+            return account_limit.bytes
+    except NoResultFound:
+        return None
+
+
+@read_session
+def get_local_account_limits(account, rse_ids=None, *, session: "Session"):
+    """
+    Returns the account limits for the account on the list of rses.
+
+    :param account:  Account to check the limit for.
+    :param rse_ids:  List of RSE ids to check the limit for.
+    :param session:  Database session in use.
+    :return:         Dictionary {'rse_id': bytes, ...}.
+    """
+
+    account_limits = {}
+    if rse_ids:
+        rse_id_clauses = []
+        for rse_id in rse_ids:
+            rse_id_clauses.append(and_(models.AccountLimit.rse_id == rse_id, models.AccountLimit.account == account))
+        rse_id_clause_chunks = [rse_id_clauses[x:x + 10] for x in range(0, len(rse_id_clauses), 10)]
+        for rse_id_chunk in rse_id_clause_chunks:
+            tmp_limits = session.query(models.AccountLimit).filter(or_(*rse_id_chunk)).all()
+            for limit in tmp_limits:
+                if limit.bytes == -1:
+                    account_limits[limit.rse_id] = float("inf")
+                else:
+                    account_limits[limit.rse_id] = limit.bytes
+    else:
+        account_limits_tmp = session.query(models.AccountLimit).filter(models.AccountLimit.account == account).all()
+        for limit in account_limits_tmp:
+            if limit.bytes == -1:
+                account_limits[limit.rse_id] = float("inf")
+            else:
+                account_limits[limit.rse_id] = limit.bytes
+    return account_limits
+
+
+@transactional_session
+def set_local_account_limit(account, rse_id, bytes_, *, session: "Session"):
+    """
+    Returns the limits for the account on the rse.
+
+    :param account:  Account to check the limit for.
+    :param rse_id:   RSE id to check the limit for.
+    :param bytes_:    The limit value, in bytes, to set.
+    :param session:  Database session in use.
+    """
+    try:
+        account_limit = session.query(models.AccountLimit).filter(models.AccountLimit.account == account,
+                                                                  models.AccountLimit.rse_id == rse_id).one()
+        account_limit.bytes = bytes_
+    except NoResultFound:
+        models.AccountLimit(account=account, rse_id=rse_id, bytes=bytes_).save(session=session)
+
+
+@transactional_session
+def set_global_account_limit(account, rse_expression, bytes_, *, session: "Session"):
+    """
+    Sets the global limit for the account on a RSE expression.
+
+    :param account:         Account to check the limit for.
+    :param rse_expression:  RSE expression to check the limit for.
+    :param bytes_:           The limit value, in bytes, to set.
+    :param session:         Database session in use.
+    """
+    try:
+        account_limit = session.query(models.AccountGlobalLimit).filter(models.AccountGlobalLimit.account == account,
+                                                                        models.AccountGlobalLimit.rse_expression == rse_expression).one()
+        account_limit.bytes = bytes_
+    except NoResultFound:
+        models.AccountGlobalLimit(account=account, rse_expression=rse_expression, bytes=bytes_).save(session=session)
+
+
+@transactional_session
+def delete_local_account_limit(account, rse_id, *, session: "Session"):
+    """
+    Deletes a local account limit.
+
+    :param account:  Account to delete the limit for.
+    :param rse_id:   RSE id to delete the limit for.
+    :param session:  Database session in use.
+    :returns:        True if something was deleted; False otherwise.
+    """
+    try:
+        session.query(models.AccountLimit).filter(models.AccountLimit.account == account,
+                                                  models.AccountLimit.rse_id == rse_id).one().delete(session=session)
+        return True
+    except NoResultFound:
+        return False
+
+
+@transactional_session
+def delete_global_account_limit(account, rse_expression, *, session: "Session"):
+    """
+    Deletes a global account limit.
+
+    :param account:         Account to delete the limit for.
+    :param rse_expression:  RSE expression to delete the limit for.
+    :param session:         Database session in use.
+    :returns:               True if something was deleted; False otherwise.
+    """
+    try:
+        session.query(models.AccountGlobalLimit).filter(models.AccountGlobalLimit.account == account,
+                                                        models.AccountGlobalLimit.rse_expression == rse_expression).one().delete(session=session)
+        return True
+    except NoResultFound:
+        return False
+
+
+@transactional_session
+def get_local_account_usage(account, rse_id=None, *, session: "Session"):
+    """
+    Read the account usage and connect it with (if available) the account limits of the account.
+
+    :param account:  The account to read.
+    :param rse_id:   The rse_id to read (If none, get all).
+    :param session:  Database session in use.
+
+    :returns:        List of dicts {'rse_id', 'rse', 'bytes', 'files', 'bytes_limit', 'bytes_remaining'}
+    """
+
+    if not rse_id:
+        # All RSESs
+        limits = get_local_account_limits(account=account, session=session)
+        counters = {c.rse_id: c for c in session.query(models.AccountUsage).filter_by(account=account).all()}
+    else:
+        # One RSE
+        limits = get_local_account_limits(account=account, rse_ids=[rse_id], session=session)
+        counters = {c.rse_id: c for c in session.query(models.AccountUsage).filter_by(account=account, rse_id=rse_id).all()}
+    result_list = []
+
+    for rse_id in set(limits).union(counters):
+        counter = counters.get(rse_id)
+        if counter:
+            counter_files = counter.files
+            counter_bytes = counter.bytes
+        else:
+            counter_files = 0
+            counter_bytes = 0
+
+        if counter_bytes > 0 or counter_files > 0 or rse_id in limits.keys():
+            result_list.append({
+                'rse_id': rse_id,
+                'rse': get_rse_name(rse_id=rse_id, session=session),
+                'bytes': counter_bytes,
+                'files': counter_files,
+                'bytes_limit': limits.get(rse_id, 0),
+                'bytes_remaining': limits.get(rse_id, 0) - counter_bytes,
+            })
+    return result_list
+
+
+@transactional_session
+def get_global_account_usage(account, rse_expression=None, *, session: "Session"):
+    """
+    Read the account usage and connect it with the global account limits of the account.
+
+    :param account:          The account to read.
+    :param rse_expression:   The RSE expression (If none, get all).
+    :param session:          Database session in use.
+
+    :returns:                List of dicts {'rse_expression', 'bytes', 'files' 'bytes_limit', 'bytes_remaining'}
+    """
+    result_list = []
+    if not rse_expression:
+        # All RSE Expressions
+        limits = get_global_account_limits(account=account, session=session)
+        all_rse_usages = {usage['rse_id']: (usage['bytes'], usage['files']) for usage in get_all_rse_usages_per_account(account=account, session=session)}
+        for rse_expression, limit in limits.items():
+            usage = 0
+            files = 0
+            for rse in limit['resolved_rse_ids']:
+                usage += all_rse_usages.get(rse, [0])[0]
+                files += all_rse_usages.get(rse, [0, 0])[1]
+            result_list.append({'rse_expression': rse_expression,
+                                'bytes': usage, 'files': files,
+                                'bytes_limit': limit['limit'],
+                                'bytes_remaining': limit['limit'] - usage})
+    else:
+        # One RSE Expression
+        limit = get_global_account_limit(account=account, rse_expression=rse_expression, session=session)
+        vo = account.vo
+        resolved_rses = [resolved_rse['id'] for resolved_rse in parse_expression(rse_expression, filter_={'vo': vo}, session=session)]
+        usage = session.query(func.sum(models.AccountUsage.bytes), func.sum(models.AccountUsage.files))\
+                       .filter(models.AccountUsage.account == account, models.AccountUsage.rse_id.in_(resolved_rses))\
+                       .group_by(models.AccountUsage.account).first()
+        if limit is None:
+            limit = 0
+        if usage is None:
+            usage = 0, 0
+        result_list.append({
+            'rse_expression': rse_expression,
+            'bytes': usage[0], 'files': usage[1],
+            'bytes_limit': limit,
+            'bytes_remaining': limit - usage[0]
+        })
+    return result_list