diff --git a/iam.tf b/iam.tf index 57b3367..1b0a082 100644 --- a/iam.tf +++ b/iam.tf @@ -328,7 +328,10 @@ data "aws_iam_policy_document" "dynamodb-onto-write-access" { data "aws_iam_policy_document" "admin-lambda-access" { statement { actions = [ - "cognito-idp:*" + "cognito-idp:*", + "dynamodb:GetItem", + "dynamodb:Query", + "dynamodb:Scan", ] resources = [ var.cognito-user-pool-arn, @@ -341,6 +344,7 @@ data "aws_iam_policy_document" "admin-lambda-access" { resources = [ "arn:aws:ses:${var.region}:${data.aws_caller_identity.this.account_id}:identity/*", aws_ses_configuration_set.ses_feedback_config.arn, + aws_dynamodb_table.sbeacon-dataportal-users-quota.arn, ] } } diff --git a/lambda/admin/admin_functions.py b/lambda/admin/admin_functions.py index a4eb272..f30fb23 100644 --- a/lambda/admin/admin_functions.py +++ b/lambda/admin/admin_functions.py @@ -8,6 +8,7 @@ from shared.cognitoutils import authenticate_admin from shared.apiutils import BeaconError, LambdaRouter from shared.utils.lambda_utils import ENV_COGNITO, ENV_BEACON, ENV_SES +from shared.dynamodb import Quota, UsageMap USER_POOL_ID = ENV_COGNITO.COGNITO_USER_POOL_ID BEACON_UI_URL = ENV_BEACON.BEACON_UI_URL @@ -178,9 +179,18 @@ def get_users(event, context): response = cognito_client.list_users(**kwargs) # Extract users and next pagination token users = response.get("Users", []) + + data = [] + for user in users: + user_sub = next(attr["Value"] for attr in user["User"]["Attributes"] if attr["Name"] == "sub") + try: + myQuota = Quota.get(user_sub) + user["Usage"] = myQuota.to_dict().get("Usage", UsageMap().as_dict()) + except Quota.DoesNotExist: + user["Usage"] = UsageMap().as_dict() next_pagination_token = response.get("PaginationToken", None) - return {"users": users, "pagination_token": next_pagination_token} + return {"users": data, "pagination_token": next_pagination_token} @router.attach("/admin/users/{email}", "delete", authenticate_admin) diff --git a/lambda/dataPortal/admin_dportal_functions.py b/lambda/dataPortal/admin_dportal_functions.py index 6aded36..5daff82 100644 --- a/lambda/dataPortal/admin_dportal_functions.py +++ b/lambda/dataPortal/admin_dportal_functions.py @@ -234,7 +234,7 @@ def list_projects(event, context): projects = Projects.scan(**params) data = [project.to_dict() for project in projects] - last_evaluated_key = json.dumps(projects.last_evaluated_key) + last_evaluated_key = json.dumps(projects.last_evaluated_key) if projects.last_evaluated_key else projects.last_evaluated_key return {"success":True, "data": data, "last_evaluated_key": last_evaluated_key} diff --git a/lambda/dataPortal/quota_function.py b/lambda/dataPortal/quota_function.py index 6240bf0..24a520b 100644 --- a/lambda/dataPortal/quota_function.py +++ b/lambda/dataPortal/quota_function.py @@ -1,7 +1,7 @@ import json from shared.apiutils import PortalError, LambdaRouter -from utils.models import Quota +from shared.dynamodb import Quota router = LambdaRouter() diff --git a/lambda/dataPortal/user_functions.py b/lambda/dataPortal/user_functions.py index 5923e7c..9c0a8e1 100644 --- a/lambda/dataPortal/user_functions.py +++ b/lambda/dataPortal/user_functions.py @@ -29,7 +29,7 @@ def list_my_projects(event, context): data = [ Projects.get(user_project.name).to_dict() for user_project in user_projects ] - last_evaluated_key = json.dumps(user_projects.last_evaluated_key) + last_evaluated_key = json.dumps(user_projects.last_evaluated_key) if user_projects.last_evaluated_key else user_projects.last_evaluated_key return {"success":True, "data": data, "last_evaluated_key": last_evaluated_key} diff --git a/lambda/dataPortal/utils/models.py b/lambda/dataPortal/utils/models.py index a2803e8..8fb65a7 100644 --- a/lambda/dataPortal/utils/models.py +++ b/lambda/dataPortal/utils/models.py @@ -77,31 +77,6 @@ class Meta: uid = UnicodeAttribute(hash_key=True) instanceName = UnicodeAttribute(range_key=True) - -class UsageMap(MapAttribute): - quotaSize = NumberAttribute(attr_name="quotaSize") - quotaQueryCount = NumberAttribute(attr_name="quotaQueryCount") - usageSize = NumberAttribute(attr_name="usageSize") - usageCount = NumberAttribute(attr_name="usageCount") - - -class Quota(Model): - class Meta: - table_name = os.environ.get("DYNAMO_QUOTA_USER_TABLE") - region = REGION - - uid = UnicodeAttribute(hash_key=True) - CostEstimation = NumberAttribute() - Usage = UsageMap() - - def to_dict(self): - return { - "uid": self.uid, - "CostEstimation": self.CostEstimation, - "Usage": self.Usage.as_dict(), - } - - class SavedQueries(Model): class Meta: table_name = os.environ.get("DYNAMO_SAVED_QUERIES_TABLE") diff --git a/main.tf b/main.tf index abf8e10..ba07c6f 100644 --- a/main.tf +++ b/main.tf @@ -74,6 +74,7 @@ locals { DYNAMO_DESCENDANTS_TABLE = aws_dynamodb_table.descendant_terms.name DYNAMO_PROJECT_USERS_TABLE = aws_dynamodb_table.project_users.name DYNAMO_PROJECT_USERS_UID_INDEX = local.project_users_uid_index + DYNAMO_QUOTA_USER_TABLE = aws_dynamodb_table.sbeacon-dataportal-users-quota.name } # layers binaries_layer = "${aws_lambda_layer_version.binaries_layer.layer_arn}:${aws_lambda_layer_version.binaries_layer.version}" @@ -721,7 +722,6 @@ module "lambda-data-portal" { DYNAMO_PROJECTS_TABLE = aws_dynamodb_table.projects.name, DYNAMO_PROJECT_USERS_TABLE = aws_dynamodb_table.project_users.name, DYNAMO_JUPYTER_INSTANCES_TABLE = aws_dynamodb_table.juptyer_notebooks.name, - DYNAMO_QUOTA_USER_TABLE = aws_dynamodb_table.sbeacon-dataportal-users-quota.name, DYNAMO_SAVED_QUERIES_TABLE = aws_dynamodb_table.saved_queries.name JUPYTER_INSTACE_ROLE_ARN = aws_iam_role.sagemaker_jupyter_instance_role.arn, USER_POOL_ID = var.cognito-user-pool-id, diff --git a/shared_resources/python-modules/python/shared/dynamodb/__init__.py b/shared_resources/python-modules/python/shared/dynamodb/__init__.py index 2451d04..4d28fdf 100644 --- a/shared_resources/python-modules/python/shared/dynamodb/__init__.py +++ b/shared_resources/python-modules/python/shared/dynamodb/__init__.py @@ -1 +1,2 @@ from .ontologies import Anscestors, Descendants, Ontology +from .quota import Quota, UsageMap diff --git a/shared_resources/python-modules/python/shared/dynamodb/quota.py b/shared_resources/python-modules/python/shared/dynamodb/quota.py new file mode 100644 index 0000000..3948cdb --- /dev/null +++ b/shared_resources/python-modules/python/shared/dynamodb/quota.py @@ -0,0 +1,33 @@ +import os + +import boto3 +from pynamodb.models import Model +from pynamodb.attributes import NumberAttribute, UnicodeAttribute, MapAttribute +from shared.utils import ENV_DYNAMO + + +SESSION = boto3.session.Session() +REGION = SESSION.region_name + +class UsageMap(MapAttribute): + quotaSize = NumberAttribute(attr_name='quotaSize',default=0) + quotaQueryCount = NumberAttribute(attr_name='quotaQueryCount',default=0) + usageSize = NumberAttribute(attr_name='usageSize',default=0) + usageCount = NumberAttribute(attr_name='usageCount',default=0) + + +class Quota(Model): + class Meta: + table_name = ENV_DYNAMO.DYNAMO_QUOTA_USER_TABLE + region = REGION + + uid = UnicodeAttribute(hash_key=True) + CostEstimation = NumberAttribute() + Usage = UsageMap() + + def to_dict(self): + return { + "uid": self.uid, + "CostEstimation": self.CostEstimation, + "Usage": self.Usage.as_dict(), + } diff --git a/shared_resources/python-modules/python/shared/utils/lambda_utils.py b/shared_resources/python-modules/python/shared/utils/lambda_utils.py index d043a45..3b75037 100644 --- a/shared_resources/python-modules/python/shared/utils/lambda_utils.py +++ b/shared_resources/python-modules/python/shared/utils/lambda_utils.py @@ -263,6 +263,10 @@ def DYNAMO_PROJECT_USERS_TABLE(self): @property def DYNAMO_PROJECT_USERS_UID_INDEX(self): return os.environ["DYNAMO_PROJECT_USERS_UID_INDEX"] + + @property + def DYNAMO_QUOTA_USER_TABLE(self): + return os.environ["DYNAMO_QUOTA_USER_TABLE"] # class SnsEnvironment: