diff --git a/backend/dataall/modules/worksheets/api/resolvers.py b/backend/dataall/modules/worksheets/api/resolvers.py index 1c72ff4fe..97e4c9edf 100644 --- a/backend/dataall/modules/worksheets/api/resolvers.py +++ b/backend/dataall/modules/worksheets/api/resolvers.py @@ -13,31 +13,18 @@ def create_worksheet(context: Context, source, input: dict = None): raise exceptions.RequiredParameter('groupUri') if input.get('SamlAdminGroupName') not in context.groups: raise exceptions.InvalidInput('groupUri', input.get('SamlAdminGroupName'), " a user's groups") - if not input.get('label'): raise exceptions.RequiredParameter('label') - with context.engine.scoped_session() as session: - return WorksheetService.create_worksheet( - session=session, - username=context.username, - data=input, - ) + return WorksheetService.create_worksheet(data=input) def update_worksheet(context: Context, source, worksheetUri: str = None, input: dict = None): - with context.engine.scoped_session() as session: - return WorksheetService.update_worksheet( - session=session, username=context.username, uri=worksheetUri, data=input - ) + return WorksheetService.update_worksheet(uri=worksheetUri, data=input) def get_worksheet(context: Context, source, worksheetUri: str = None): - with context.engine.scoped_session() as session: - return WorksheetService.get_worksheet( - session=session, - uri=worksheetUri, - ) + return WorksheetService.get_worksheet(uri=worksheetUri) def resolve_user_role(context: Context, source: Worksheet): @@ -51,24 +38,12 @@ def resolve_user_role(context: Context, source: Worksheet): def list_worksheets(context, source, filter: dict = None): if not filter: filter = {} - with context.engine.scoped_session() as session: - return WorksheetRepository.paginated_user_worksheets( - session=session, - username=context.username, - groups=context.groups, - uri=None, - data=filter, - check_perm=True, - ) + return WorksheetService.list_user_worksheets(filter) def run_sql_query(context: Context, source, environmentUri: str = None, worksheetUri: str = None, sqlQuery: str = None): - with context.engine.scoped_session() as session: - return WorksheetService.run_sql_query( - session=session, uri=environmentUri, worksheetUri=worksheetUri, sqlQuery=sqlQuery - ) + return WorksheetService.run_sql_query(uri=environmentUri, worksheetUri=worksheetUri, sqlQuery=sqlQuery) def delete_worksheet(context, source, worksheetUri: str = None): - with context.engine.scoped_session() as session: - return WorksheetService.delete_worksheet(session=session, uri=worksheetUri) + return WorksheetService.delete_worksheet(uri=worksheetUri) diff --git a/backend/dataall/modules/worksheets/services/worksheet_service.py b/backend/dataall/modules/worksheets/services/worksheet_service.py index 128f2af94..e07ce63ce 100644 --- a/backend/dataall/modules/worksheets/services/worksheet_service.py +++ b/backend/dataall/modules/worksheets/services/worksheet_service.py @@ -3,6 +3,7 @@ from dataall.core.activity.db.activity_models import Activity from dataall.core.environment.services.environment_service import EnvironmentService from dataall.base.db import exceptions +from dataall.base.context import get_context from dataall.core.permissions.services.resource_policy_service import ResourcePolicyService from dataall.core.permissions.services.tenant_policy_service import TenantPolicyService from dataall.modules.worksheets.aws.athena_client import AthenaClient @@ -23,7 +24,7 @@ class WorksheetService: @staticmethod - def get_worksheet_by_uri(session, uri: str) -> Worksheet: + def _get_worksheet_by_uri(session, uri: str) -> Worksheet: if not uri: raise exceptions.RequiredParameter(param_name='worksheetUri') worksheet = WorksheetRepository.find_worksheet_by_uri(session, uri) @@ -33,94 +34,114 @@ def get_worksheet_by_uri(session, uri: str) -> Worksheet: @staticmethod @TenantPolicyService.has_tenant_permission(MANAGE_WORKSHEETS) - def create_worksheet(session, username, data=None) -> Worksheet: - worksheet = Worksheet( - owner=username, - label=data.get('label'), - description=data.get('description', 'No description provided'), - tags=data.get('tags'), - chartConfig={'dimensions': [], 'measures': [], 'chartType': 'bar'}, - SamlAdminGroupName=data['SamlAdminGroupName'], - ) - - session.add(worksheet) - session.commit() - - activity = Activity( - action='WORKSHEET:CREATE', - label='WORKSHEET:CREATE', - owner=username, - summary=f'{username} created worksheet {worksheet.name} ', - targetUri=worksheet.worksheetUri, - targetType='worksheet', - ) - session.add(activity) - - ResourcePolicyService.attach_resource_policy( - session=session, - group=data['SamlAdminGroupName'], - permissions=WORKSHEET_ALL, - resource_uri=worksheet.worksheetUri, - resource_type=Worksheet.__name__, - ) + def create_worksheet(data=None) -> Worksheet: + context = get_context() + with context.db_engine.scoped_session() as session: + worksheet = Worksheet( + owner=context.username, + label=data.get('label'), + description=data.get('description', 'No description provided'), + tags=data.get('tags'), + chartConfig={'dimensions': [], 'measures': [], 'chartType': 'bar'}, + SamlAdminGroupName=data['SamlAdminGroupName'], + ) + + session.add(worksheet) + session.commit() + + activity = Activity( + action='WORKSHEET:CREATE', + label='WORKSHEET:CREATE', + owner=context.username, + summary=f'{context.username} created worksheet {worksheet.name} ', + targetUri=worksheet.worksheetUri, + targetType='worksheet', + ) + session.add(activity) + + ResourcePolicyService.attach_resource_policy( + session=session, + group=data['SamlAdminGroupName'], + permissions=WORKSHEET_ALL, + resource_uri=worksheet.worksheetUri, + resource_type=Worksheet.__name__, + ) return worksheet @staticmethod @TenantPolicyService.has_tenant_permission(MANAGE_WORKSHEETS) @ResourcePolicyService.has_resource_permission(UPDATE_WORKSHEET) - def update_worksheet(session, username, uri, data=None): - worksheet = WorksheetService.get_worksheet_by_uri(session, uri) - for field in data.keys(): - setattr(worksheet, field, data.get(field)) - session.commit() - - activity = Activity( - action='WORKSHEET:UPDATE', - label='WORKSHEET:UPDATE', - owner=username, - summary=f'{username} updated worksheet {worksheet.name} ', - targetUri=worksheet.worksheetUri, - targetType='worksheet', - ) - session.add(activity) - return worksheet + def update_worksheet(uri, data=None): + context = get_context() + with context.db_engine.scoped_session() as session: + worksheet = WorksheetService._get_worksheet_by_uri(session, uri) + for field in data.keys(): + setattr(worksheet, field, data.get(field)) + session.commit() + + activity = Activity( + action='WORKSHEET:UPDATE', + label='WORKSHEET:UPDATE', + owner=context.username, + summary=f'{context.username} updated worksheet {worksheet.name} ', + targetUri=worksheet.worksheetUri, + targetType='worksheet', + ) + session.add(activity) + return worksheet @staticmethod @ResourcePolicyService.has_resource_permission(GET_WORKSHEET) - def get_worksheet(session, uri): - worksheet = WorksheetService.get_worksheet_by_uri(session, uri) - return worksheet + def get_worksheet(uri): + with get_context().db_engine.scoped_session() as session: + worksheet = WorksheetService._get_worksheet_by_uri(session, uri) + return worksheet + + @staticmethod + def list_user_worksheets(filter): + context = get_context() + with context.db_engine.scoped_session() as session: + return WorksheetRepository.paginated_user_worksheets( + session=session, + username=context.username, + groups=context.groups, + uri=None, + data=filter, + check_perm=True, + ) @staticmethod @TenantPolicyService.has_tenant_permission(MANAGE_WORKSHEETS) @ResourcePolicyService.has_resource_permission(DELETE_WORKSHEET) - def delete_worksheet(session, uri) -> bool: - worksheet = WorksheetService.get_worksheet_by_uri(session, uri) - session.delete(worksheet) - ResourcePolicyService.delete_resource_policy( - session=session, - group=worksheet.SamlAdminGroupName, - resource_uri=uri, - resource_type=Worksheet.__name__, - ) - return True + def delete_worksheet(uri) -> bool: + with get_context().db_engine.scoped_session() as session: + worksheet = WorksheetService._get_worksheet_by_uri(session, uri) + session.delete(worksheet) + ResourcePolicyService.delete_resource_policy( + session=session, + group=worksheet.SamlAdminGroupName, + resource_uri=uri, + resource_type=Worksheet.__name__, + ) + return True @staticmethod @ResourcePolicyService.has_resource_permission(RUN_ATHENA_QUERY) - def run_sql_query(session, uri, worksheetUri, sqlQuery): - environment = EnvironmentService.get_environment_by_uri(session, uri) - worksheet = WorksheetService.get_worksheet_by_uri(session, worksheetUri) - - env_group = EnvironmentService.get_environment_group( - session, worksheet.SamlAdminGroupName, environment.environmentUri - ) - - cursor = AthenaClient.run_athena_query( - aws_account_id=environment.AwsAccountId, - env_group=env_group, - s3_staging_dir=f's3://{environment.EnvironmentDefaultBucketName}/athenaqueries/{env_group.environmentAthenaWorkGroup}/', - region=environment.region, - sql=sqlQuery, - ) - - return AthenaClient.convert_query_output(cursor) + def run_sql_query(uri, worksheetUri, sqlQuery): + with get_context().db_engine.scoped_session() as session: + environment = EnvironmentService.get_environment_by_uri(session, uri) + worksheet = WorksheetService._get_worksheet_by_uri(session, worksheetUri) + + env_group = EnvironmentService.get_environment_group( + session, worksheet.SamlAdminGroupName, environment.environmentUri + ) + + cursor = AthenaClient.run_athena_query( + aws_account_id=environment.AwsAccountId, + env_group=env_group, + s3_staging_dir=f's3://{environment.EnvironmentDefaultBucketName}/athenaqueries/{env_group.environmentAthenaWorkGroup}/', + region=environment.region, + sql=sqlQuery, + ) + + return AthenaClient.convert_query_output(cursor)