add resource_threshold

Author: petrkalos

backend/dataall/base/db/exceptions.py

@@ -193,3 +193,26 @@ def __init__(self, action, message):
 
     def __str__(self):
         return f'{self.message}'
+
+
+class ResourceThresholdExceeded(Exception):
+    def __init__(self, username, action):
+        self.username = username
+        self.action = action
+        self.message = f"""
+                    An error occurred (ResourceThresholdExceeded) when calling {self.action} operation:
+                    Requests exceeded max daily invocation count for User: {self.username}
+                """
+
+    def __str__(self):
+        return f'{self.message}'
+
+
+class ModelGuardrailException(Exception):
+    def __init__(self, message):
+        self.message = f"""
+                    An error occurred (ModelGuardrailException) when invoking the model: {message}
+                """
+
+    def __str__(self):
+        return f'{self.message}'

backend/dataall/core/resource_threshold/__init__.py

@@ -0,0 +1,12 @@
+from dataall.base.db import Base, utils
+from sqlalchemy import String, Integer, Column, Date
+from datetime import date
+
+
+class ResourceThreshold(Base):
+    __tablename__ = 'resource_threshold'
+    actionUri = Column(String(64), primary_key=True, default=utils.uuid('resource_threshold'))
+    username = Column(String(64), nullable=False)
+    actionType = Column(String(64), nullable=False)
+    date = Column(Date, default=date.today, nullable=False)
+    count = Column(Integer, default=1, nullable=False)

backend/dataall/core/resource_threshold/db/resource_threshold_models.py

@@ -0,0 +1,56 @@
+from dataall.core.resource_threshold.db.resource_threshold_models import ResourceThreshold
+from sqlalchemy import and_
+from datetime import date
+
+
+class ResourceThresholdRepository:
+    @staticmethod
+    def get_count_today(session, username, action_type):
+        amount = (
+            session.query(ResourceThreshold.count)
+            .filter(
+                and_(
+                    ResourceThreshold.username == username,
+                    ResourceThreshold.actionType == action_type,
+                    ResourceThreshold.date == date.today(),
+                )
+            )
+            .scalar()
+        )
+        return amount if amount else 0
+
+    @staticmethod
+    def add_entry(session, username, action_type):
+        user_entry = ResourceThresholdRepository._get_user_entry(session, username, action_type)
+        if user_entry:
+            session.query(ResourceThreshold).filter(
+                and_(
+                    ResourceThreshold.username == username,
+                    ResourceThreshold.actionType == action_type,
+                )
+            ).update({ResourceThreshold.count: 1, ResourceThreshold.date: date.today()}, synchronize_session=False)
+            session.commit()
+        else:
+            action_entry = ResourceThreshold(username=username, actionType=action_type)
+            session.add(action_entry)
+            session.commit()
+
+    @staticmethod
+    def increment_count(session, username, action_type):
+        session.query(ResourceThreshold).filter(
+            and_(
+                ResourceThreshold.username == username,
+                ResourceThreshold.actionType == action_type,
+                ResourceThreshold.date == date.today(),
+            )
+        ).update({ResourceThreshold.count: ResourceThreshold.count + 1}, synchronize_session=False)
+        session.commit()
+
+    @staticmethod
+    def _get_user_entry(session, username, action_type):
+        entry = (
+            session.query(ResourceThreshold)
+            .filter(and_(ResourceThreshold.username == username, ResourceThreshold.actionType == action_type))
+            .first()
+        )
+        return entry

backend/dataall/core/resource_threshold/db/resource_threshold_repositories.py

@@ -0,0 +1,42 @@
+from dataall.core.resource_threshold.db.resource_threshold_repositories import ResourceThresholdRepository
+from dataall.base.db import exceptions
+from functools import wraps
+from dataall.base.config import config
+from dataall.base.context import get_context
+
+import logging
+
+log = logging.getLogger(__name__)
+
+
+class ResourceThresholdService:
+    @staticmethod
+    def check_invocation_count(action_type, max_count_config_path):
+        def decorator(func):
+            @wraps(func)
+            def wrapper(*args, **kwargs):
+                context = get_context()
+                with context.db_engine.scoped_session() as session:
+                    count = ResourceThresholdRepository.get_count_today(
+                        session=session, username=context.username, action_type=action_type
+                    )
+                    max_count = config.get_property(max_count_config_path, 10)
+                    log.info(
+                        f'User {context.username} has invoked {action_type} {count} times today of max {max_count}'
+                    )
+                    if count < max_count:
+                        if count == 0:
+                            ResourceThresholdRepository.add_entry(
+                                session=session, username=context.username, action_type=action_type
+                            )
+                        else:
+                            ResourceThresholdRepository.increment_count(
+                                session=session, username=context.username, action_type=action_type
+                            )
+                        return func(*args, **kwargs)
+                    else:
+                        raise exceptions.ResourceThresholdExceeded(username=context.username, action=action_type)
+
+            return wrapper
+
+        return decorator

backend/dataall/core/resource_threshold/services/__init__.py

@@ -20,6 +20,7 @@
 from dataall.modules.omics.db.omics_models import OmicsWorkflow, OmicsRun
 from dataall.modules.metadata_forms.db.metadata_form_models import *
 from dataall.modules.redshift_datasets.db.redshift_models import RedshiftDataset, RedshiftTable, RedshiftConnection
+from dataall.core.resource_threshold.db.resource_threshold_models import ResourceThreshold
 # fmt: on
 # enable ruff-format back
 

backend/dataall/core/resource_threshold/services/resource_threshold_service.py

@@ -0,0 +1,37 @@
+"""resource_threshholds_added
+
+Revision ID: 2258cd8d6e9f
+Revises: 5a798acc6282
+Create Date: 2024-08-22 12:31:38.465650
+
+"""
+
+from alembic import op
+import sqlalchemy as sa
+
+
+# revision identifiers, used by Alembic.
+revision = '2258cd8d6e9f'
+down_revision = '5a798acc6282'
+branch_labels = None
+depends_on = None
+
+
+def upgrade():
+    # ### commands auto generated by Alembic - please adjust! ###
+    op.create_table(
+        'resource_threshold',
+        sa.Column('actionUri', sa.String(length=64), nullable=False),
+        sa.Column('username', sa.String(length=64), nullable=False),
+        sa.Column('actionType', sa.String(length=64), nullable=False),
+        sa.Column('date', sa.Date(), nullable=False),
+        sa.Column('count', sa.Integer(), nullable=False),
+        sa.PrimaryKeyConstraint('actionUri'),
+    )
+    # ### end Alembic commands ###
+
+
+def downgrade():
+    # ### commands auto generated by Alembic - please adjust! ###
+    op.drop_table('resource_threshold')
+    # ### end Alembic commands ###