Refactor core/environment&stack module - part 2

backend/dataall/core/environment/api/resolvers.py

@@ -5,21 +5,20 @@
 import boto3
 from botocore.config import Config
 from botocore.exceptions import ClientError
-from sqlalchemy import and_
 
 from dataall.base.aws.iam import IAM
 from dataall.base.aws.parameter_store import ParameterStoreManager
 from dataall.base.aws.sts import SessionHelper
 from dataall.base.utils import Parameter
-from dataall.core.environment.db.environment_models import Environment, EnvironmentGroup
+from dataall.core.environment.db.environment_models import Environment
 from dataall.core.environment.services.managed_iam_policies import PolicyManager
 from dataall.core.environment.services.environment_resource_manager import EnvironmentResourceManager
 from dataall.core.environment.services.environment_service import EnvironmentService
 from dataall.core.environment.api.enums import EnvironmentPermission
 from dataall.core.permissions.services.resource_policy_service import ResourcePolicyService
 from dataall.core.stacks.api import stack_helper
 from dataall.core.stacks.aws.cloudformation import CloudFormation
-from dataall.core.stacks.db.stack_repositories import Stack
+from dataall.core.stacks.db.stack_repositories import StackRepository
 from dataall.core.vpc.services.vpc_service import VpcService
 from dataall.base.aws.ec2_client import EC2
 from dataall.base.feature_toggle_checker import is_feature_enabled
@@ -38,7 +37,6 @@
     LINK_ENVIRONMENT,
 )
 
-
 log = logging.getLogger()
 
 
@@ -116,7 +114,7 @@ def create_environment(context: Context, source, input={}):
             uri=input.get('organizationUri'),
             data=input,
         )
-        Stack.create_stack(
+        StackRepository.create_stack(
             session=session,
             environment_uri=env.environmentUri,
             target_type='environment',
@@ -376,20 +374,8 @@ def resolve_user_role(context: Context, source: Environment):
         return EnvironmentPermission.Owner.value
     elif source.SamlGroupName in context.groups:
         return EnvironmentPermission.Admin.value
-    else:
-        with context.engine.scoped_session() as session:
-            env_group = (
-                session.query(EnvironmentGroup)
-                .filter(
-                    and_(
-                        EnvironmentGroup.environmentUri == source.environmentUri,
-                        EnvironmentGroup.groupUri.in_(context.groups),
-                    )
-                )
-                .first()
-            )
-            if env_group:
-                return EnvironmentPermission.Invited.value
+    elif EnvironmentService.is_user_invited(source.environmentUri):
+        return EnvironmentPermission.Invited.value
     return EnvironmentPermission.NotInvited.value
 
 
@@ -418,14 +404,7 @@ def _get_environment_group_aws_session(session, username, groups, environment, g
                 message=f'User: {username} is not member of the environment admins team {environment.SamlGroupName}',
             )
     else:
-        env_group: EnvironmentGroup = (
-            session.query(EnvironmentGroup)
-            .filter(
-                EnvironmentGroup.environmentUri == environment.environmentUri,
-                EnvironmentGroup.groupUri == groupUri,
-            )
-            .first()
-        )
+        env_group = EnvironmentService.get_environment_group(session, environment.environmentUri, groupUri)
         if not env_group:
             raise exceptions.UnauthorizedOperation(
                 action='ENVIRONMENT_AWS_ACCESS',
@@ -700,8 +679,9 @@ def resolve_environment(context, source, **kwargs):
     """Resolves the environment for a environmental resource"""
     if not source:
         return None
+
     with context.engine.scoped_session() as session:
-        return session.query(Environment).get(source.environmentUri)
+        return EnvironmentService.get_environment_by_uri(source.environmentUri)
 
 
 def resolve_parameters(context, source: Environment, **kwargs):

backend/dataall/core/environment/db/environment_repositories.py

@@ -1,5 +1,11 @@
-from dataall.core.environment.db.environment_models import EnvironmentParameter, Environment
-from sqlalchemy.sql import and_
+from dataall.core.environment.db.environment_models import (
+    EnvironmentParameter,
+    Environment,
+    ConsumptionRole,
+    EnvironmentGroup,
+)
+from sqlalchemy.sql import and_, or_
+from sqlalchemy.orm import Query
 
 from dataall.base.db import exceptions
 
@@ -67,3 +73,231 @@ def find_environment_by_account_region(session, account_id, region):
         if not environment:
             return None
         return environment
+
+    @staticmethod
+    def get_environment_consumption_role(session, role_uri, environment_uri) -> ConsumptionRole:
+        return (
+            session.query(ConsumptionRole)
+            .filter(
+                (
+                    and_(
+                        ConsumptionRole.consumptionRoleUri == role_uri,
+                        ConsumptionRole.environmentUri == environment_uri,
+                    )
+                )
+            )
+            .first()
+        )
+
+    @staticmethod
+    def get_environment_group(session, group_uri, environment_uri):
+        return (
+            session.query(EnvironmentGroup)
+            .filter(
+                (
+                    and_(
+                        EnvironmentGroup.groupUri == group_uri,
+                        EnvironmentGroup.environmentUri == environment_uri,
+                    )
+                )
+            )
+            .first()
+        )
+
+    @staticmethod
+    def get_consumption_role(session, uri):
+        return (
+            session.query(ConsumptionRole)
+            .filter(
+                and_(
+                    ConsumptionRole.consumptionRoleUri == uri,
+                )
+            )
+            .first()
+        )
+
+    @staticmethod
+    def find_consumption_roles_by_IAMArn(session, uri, arn):
+        return (
+            session.query(ConsumptionRole)
+            .filter(and_(ConsumptionRole.environmentUri == uri, ConsumptionRole.IAMRoleArn == arn))
+            .first()
+        )
+
+    @staticmethod
+    def query_all_environment_consumption_roles(session, uri, filter) -> Query:
+        query = session.query(ConsumptionRole).filter(ConsumptionRole.environmentUri == uri)
+        if filter and filter.get('term'):
+            term = filter['term']
+            query = query.filter(
+                or_(
+                    ConsumptionRole.consumptionRoleName.ilike('%' + term + '%'),
+                )
+            )
+        if filter and filter.get('groupUri'):
+            group = filter['groupUri']
+            query = query.filter(
+                or_(
+                    ConsumptionRole.groupUri == group,
+                )
+            )
+        return query.order_by(ConsumptionRole.consumptionRoleUri)
+
+    @staticmethod
+    def query_user_environment_consumption_roles(session, groups, uri, filter) -> Query:
+        query = (
+            session.query(ConsumptionRole)
+            .filter(ConsumptionRole.environmentUri == uri)
+            .filter(ConsumptionRole.groupUri.in_(groups))
+        )
+        if filter and filter.get('term'):
+            term = filter['term']
+            query = query.filter(
+                or_(
+                    ConsumptionRole.consumptionRoleName.ilike('%' + term + '%'),
+                )
+            )
+        if filter and filter.get('groupUri'):
+            print('filter group')
+            group = filter['groupUri']
+            query = query.filter(
+                or_(
+                    ConsumptionRole.groupUri == group,
+                )
+            )
+        return query.order_by(ConsumptionRole.consumptionRoleUri)
+
+    @staticmethod
+    def query_environment_invited_groups(session, uri, filter) -> Query:
+        query = (
+            session.query(EnvironmentGroup)
+            .join(
+                Environment,
+                EnvironmentGroup.environmentUri == Environment.environmentUri,
+            )
+            .filter(
+                and_(
+                    Environment.environmentUri == uri,
+                    EnvironmentGroup.groupUri != Environment.SamlGroupName,
+                )
+            )
+        )
+        if filter and filter.get('term'):
+            term = filter['term']
+            query = query.filter(
+                or_(
+                    EnvironmentGroup.groupUri.ilike('%' + term + '%'),
+                )
+            )
+        return query
+
+    @staticmethod
+    def query_user_environment_groups(session, groups, uri, filter) -> Query:
+        query = (
+            session.query(EnvironmentGroup)
+            .filter(EnvironmentGroup.environmentUri == uri)
+            .filter(EnvironmentGroup.groupUri.in_(groups))
+        )
+        if filter and filter.get('term'):
+            term = filter['term']
+            query = query.filter(
+                or_(
+                    EnvironmentGroup.groupUri.ilike('%' + term + '%'),
+                )
+            )
+        return query
+
+    @staticmethod
+    def query_all_environment_groups(session, uri, filter) -> Query:
+        query = session.query(EnvironmentGroup).filter(EnvironmentGroup.environmentUri == uri)
+        if filter and filter.get('term'):
+            term = filter['term']
+            query = query.filter(
+                or_(
+                    EnvironmentGroup.groupUri.ilike('%' + term + '%'),
+                )
+            )
+        return query
+
+    @staticmethod
+    def query_user_consumption_roles(session, username, groups, filter) -> Query:
+        query = (
+            session.query(ConsumptionRole)
+            .filter(ConsumptionRole.groupUri.in_(groups))
+            .distinct(ConsumptionRole.consumptionRoleName)
+        )
+        if filter and filter.get('term'):
+            term = filter['term']
+            query = query.filter(
+                or_(
+                    ConsumptionRole.consumptionRoleName.ilike('%' + term + '%'),
+                )
+            )
+        if filter and filter.get('groupUri'):
+            print('filter group')
+            group = filter['groupUri']
+            query = query.filter(
+                or_(
+                    ConsumptionRole.groupUri == group,
+                )
+            )
+        return query
+
+    @staticmethod
+    def query_user_groups(session, username, groups, filter) -> Query:
+        query = (
+            session.query(EnvironmentGroup)
+            .filter(EnvironmentGroup.groupUri.in_(groups))
+            .distinct(EnvironmentGroup.groupUri)
+        )
+        if filter and filter.get('term'):
+            term = filter['term']
+            query = query.filter(
+                or_(
+                    EnvironmentGroup.groupUri.ilike('%' + term + '%'),
+                )
+            )
+        return query
+
+    @staticmethod
+    def query_user_environments(session, username, groups, filter) -> Query:
+        query = (
+            session.query(Environment)
+            .outerjoin(
+                EnvironmentGroup,
+                Environment.environmentUri == EnvironmentGroup.environmentUri,
+            )
+            .filter(
+                or_(
+                    Environment.owner == username,
+                    EnvironmentGroup.groupUri.in_(groups),
+                )
+            )
+        )
+        if filter and filter.get('term'):
+            term = filter['term']
+            query = query.filter(
+                or_(
+                    Environment.label.ilike('%' + term + '%'),
+                    Environment.description.ilike('%' + term + '%'),
+                    Environment.tags.contains(f'{{{term}}}'),
+                    Environment.region.ilike('%' + term + '%'),
+                )
+            )
+        if filter and filter.get('SamlGroupName') and filter.get('SamlGroupName') in groups:
+            query = query.filter(EnvironmentGroup.groupUri == filter.get('SamlGroupName'))
+        return query
+
+    @staticmethod
+    def is_user_invited_to_environment(session, groups, uri):
+        env_group = (
+            session.query(EnvironmentGroup)
+            .filter(
+                and_(
+                    EnvironmentGroup.environmentUri == uri,
+                    EnvironmentGroup.groupUri.in_(groups),
+                )
+            )
+            .first()
+        )
+        return env_group is not None