Add Usage Quota data on get List Cognito User

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,
     ]
   }
 }

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)

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}
 
 

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()
 

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}
 
 

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")

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,

shared_resources/python-modules/python/shared/dynamodb/__init__.py

@@ -1 +1,2 @@
 from .ontologies import Anscestors, Descendants, Ontology
+from .quota import Quota, UsageMap

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(),
+        }

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: