feat: Enhance tenant registration service and improve error handling

- Update API documentation and class names for clarity
- Implement pagination for listing tenant registrations
- Add input validation and improve error responses
- Update test scripts to cover new functionality
- Refactor lambda function for better error handling and response
  formatting

resources/functions/tenant-registrations/index.py

@@ -4,14 +4,22 @@
 from http import HTTPStatus
 from boto3.dynamodb.conditions import Attr
 from aws_lambda_powertools import Logger, Tracer
-from aws_lambda_powertools.event_handler import APIGatewayHttpResolver
 from aws_lambda_powertools.utilities.typing import LambdaContext
 from aws_requests_auth.boto_utils import BotoAWSRequestsAuth
+from aws_lambda_powertools.event_handler import APIGatewayHttpResolver
+from aws_lambda_powertools.logging import correlation_paths
+from aws_lambda_powertools.event_handler.openapi.params import Query, Path
+from aws_lambda_powertools.shared.types import Annotated
+from aws_lambda_powertools.event_handler.exceptions import (
+    InternalServerError,
+    NotFoundError,
+)
+from typing import Optional
 import requests
 import uuid
 import json
 
-app = APIGatewayHttpResolver()
+app = APIGatewayHttpResolver(enable_validation=True)
 tracer = Tracer()
 logger = Logger()
 dynamodb = boto3.resource("dynamodb")
@@ -53,11 +61,11 @@ def create_tenant_registration():
         )
     except ClientError as e:
         logger.error(f"Error creating tenant registration: {str(e)}")
-        return {"message": "Internal server error"}, HTTPStatus.INTERNAL_SERVER_ERROR
+        raise InternalServerError("Unknown error during processing!")
 
     response = requests.post(f"{tenant_api_url}tenants", json=tenant_data, auth=auth)
     if response.status_code != HTTPStatus.CREATED:
-        return {"message": "Failed to create tenant"}, HTTPStatus.INTERNAL_SERVER_ERROR
+        raise InternalServerError("Failed to create tenant")
 
     tenant_id = response.json()["data"]["tenantId"]
     # add tenant id to the tenant-registration after creating it to ensure
@@ -74,113 +82,134 @@ def create_tenant_registration():
     )
 
     return {
-        "tenantRegistrationId": tenant_registration_id,
-        "tenantId": tenant_id,
-        "message": "Tenant registration initiated",
+        "data": {
+            "tenantRegistrationId": tenant_registration_id,
+            "tenantId": tenant_id,
+            "message": "Tenant registration initiated",
+        }
     }, HTTPStatus.CREATED
 
 
 @app.get("/tenant-registrations/<tenant_registration_id>")
 @tracer.capture_method
-def get_tenant_registration(tenant_registration_id: str):
+def get_tenant_registration(tenant_registration_id: Annotated[str, Path(min_length=1)]):
     try:
         response = tenant_registration_table.get_item(
             Key={"tenantRegistrationId": tenant_registration_id}
         )
     except ClientError as e:
-        logger.error(f"Error deleting tenant registration: {str(e)}")
-        return {"message": "Internal server error"}, HTTPStatus.INTERNAL_SERVER_ERROR
+        logger.error(f"Error getting tenant registration: {str(e)}")
+        raise InternalServerError("Unknown error during processing!")
 
     if "Item" not in response:
-        return {"message": "Tenant registration not found."}, HTTPStatus.NOT_FOUND
+        raise NotFoundError(
+            f"Tenant registration not found for id {tenant_registration_id}"
+        )
 
-    return response["Item"]
+    return {"data": response["Item"]}, HTTPStatus.OK
 
 
-def __update_tenant_registration(tenant_registration_id, update_data):
-    update_expression = "SET " + ", ".join(f"#{k}=:{k}" for k in update_data.keys())
-    expression_attribute_names = {f"#{k}": k for k in update_data.keys()}
-    expression_attribute_values = {f":{k}": v for k, v in update_data.items()}
+@app.get("/tenant-registrations")
+@tracer.capture_method
+def list_tenant_registrations(
+    limit: Annotated[Optional[int], Query(gt=0)] = 10,
+    next_token: Annotated[Optional[str], Query(min_length=0)] = None,
+):
+    logger.info("Request received to list all tenant registrations")
+    registrations = None
+    last_evaluated_key = None
+
+    kwargs = {"Limit": limit}
+    if next_token:
+        kwargs["ExclusiveStartKey"] = {"tenantRegistrationId": next_token}
 
-    return tenant_registration_table.update_item(
-        Key={"tenantRegistrationId": tenant_registration_id},
-        UpdateExpression=update_expression,
-        ExpressionAttributeNames=expression_attribute_names,
-        ExpressionAttributeValues=expression_attribute_values,
-        ConditionExpression=Attr("tenantRegistrationId").exists(),
-        ReturnValues="ALL_NEW",
-    )
+    try:
+        response = tenant_registration_table.scan(**kwargs)
+        registrations = response["Items"]
+        last_evaluated_key = response.get("LastEvaluatedKey")
+
+    except ClientError as error:
+        logger.error(f"Error listing tenant registrations: {str(error)}")
+        raise InternalServerError("Unknown error during processing!")
+
+    return_response = {"data": registrations}
+    if last_evaluated_key:
+        return_response["next_token"] = last_evaluated_key["tenantRegistrationId"]
+
+    return return_response, HTTPStatus.OK
 
 
 @app.patch("/tenant-registrations/<tenant_registration_id>")
 @tracer.capture_method
-def update_tenant_registration(tenant_registration_id: str):
+def update_tenant_registration(
+    tenant_registration_id: Annotated[str, Path(min_length=1)]
+):
     update_data = app.current_event.json_body
-    logger.info(f"Update data: {update_data}")
-    logger.info(f"tenantRegistrationData: {update_data.get("tenantRegistrationData", {})}")
-    logger.info(f"tenantData: {update_data.get("tenantData", {})}")
     try:
         tenant_id = None
+        tenant_registration = None
+        tenant = {}
         if update_data.get("tenantRegistrationData", {}) == {}:
-            tenant_registration = tenant_registration_table.get_item(
+            tenant_registration_response = tenant_registration_table.get_item(
                 Key={"tenantRegistrationId": tenant_registration_id}
             )
-            tenant_id = tenant_registration["Item"].get("tenantId")
+            tenant_id = tenant_registration_response["Item"].get("tenantId")
+            tenant_registration = tenant_registration_response["Item"]
             if not tenant_id:
-                return {
-                    "message": "Tenant ID not found for this registration"
-                }, HTTPStatus.BAD_REQUEST
+                raise NotFoundError("Tenant ID not found for this registration")
         else:
-            tenant_registration = __update_tenant_registration(
+            tenant_registration_response = __update_tenant_registration(
                 tenant_registration_id, update_data.get("tenantRegistrationData", {})
             )
-            tenant_id = tenant_registration['Attributes']['tenantId']
+            tenant_id = tenant_registration_response["Attributes"]["tenantId"]
+            tenant_registration = tenant_registration_response["Attributes"]
 
-        logger.info(f"url: {tenant_api_url}tenants/{tenant_id}")
         if update_data.get("tenantData", {}) != {}:
             update_response = requests.put(
                 f"{tenant_api_url}tenants/{tenant_id}",
                 json=update_data.get("tenantData"),
                 auth=auth,
             )
+            tenant = update_response.json()
             if update_response.status_code != 200:
-                logger.error(
-                    f"Error updating tenant: {update_response.json()}"
-                )
-                return {
-                    "message": "Failed to update tenant"
-                }, HTTPStatus.INTERNAL_SERVER_ERROR
-
-        return {"message": "Tenant registration updated successfully"}
+                logger.error(f"Error updating tenant: {update_response.json()}")
+                raise InternalServerError("Failed to update tenant")
+
+        return {
+            "data": {
+                "tenantRegistration": tenant_registration,
+                "tenant": tenant.get('data'),
+                "message": "Tenant registration updated successfully",
+            }
+        }, HTTPStatus.OK
     except ClientError as e:
         logger.error(f"Error updating tenant registration: {str(e)}")
-        return {"message": "Internal server error"}, HTTPStatus.INTERNAL_SERVER_ERROR
+        raise InternalServerError("Unknown error during processing!")
 
 
 @app.delete("/tenant-registrations/<tenant_registration_id>")
 @tracer.capture_method
-def delete_tenant_registration(tenant_registration_id: str):
+def delete_tenant_registration(
+    tenant_registration_id: Annotated[str, Path(min_length=1)]
+):
     try:
         response = tenant_registration_table.get_item(
             Key={"tenantRegistrationId": tenant_registration_id}
         )
         if "Item" not in response:
-            return {"message": "Tenant registration not found."}, HTTPStatus.NOT_FOUND
+            raise NotFoundError(
+                f"Tenant registration not found for id {tenant_registration_id}"
+            )
 
         tenant_id = response["Item"].get("tenantId")
         if not tenant_id:
-            return {
-                "message": "Tenant ID not found for this registration"
-            }, HTTPStatus.BAD_REQUEST
+            raise NotFoundError("Tenant ID not found for this registration")
 
-        # Delete tenant using the tenant management API
         delete_response = requests.delete(
             f"{tenant_api_url}tenants/{tenant_id}", auth=auth
         )
         if delete_response.status_code != 200:
-            return {
-                "message": "Failed to delete tenant"
-            }, HTTPStatus.INTERNAL_SERVER_ERROR
+            raise InternalServerError("Failed to delete tenant")
 
         __update_tenant_registration(
             tenant_registration_id,
@@ -194,10 +223,27 @@ def delete_tenant_registration(tenant_registration_id: str):
             offboarding_detail_type,
         )
 
-        return {"message": "Tenant registration deletion initiated"}
+        return {
+            "data": {"message": "Tenant registration deletion initiated"}
+        }, HTTPStatus.OK
     except ClientError as e:
         logger.error(f"Error deleting tenant registration: {str(e)}")
-        return {"message": "Internal server error"}, HTTPStatus.INTERNAL_SERVER_ERROR
+        raise InternalServerError("Unknown error during processing!")
+
+
+def __update_tenant_registration(tenant_registration_id, update_data):
+    update_expression = "SET " + ", ".join(f"#{k}=:{k}" for k in update_data.keys())
+    expression_attribute_names = {f"#{k}": k for k in update_data.keys()}
+    expression_attribute_values = {f":{k}": v for k, v in update_data.items()}
+
+    return tenant_registration_table.update_item(
+        Key={"tenantRegistrationId": tenant_registration_id},
+        UpdateExpression=update_expression,
+        ExpressionAttributeNames=expression_attribute_names,
+        ExpressionAttributeValues=expression_attribute_values,
+        ConditionExpression=Attr("tenantRegistrationId").exists(),
+        ReturnValues="ALL_NEW",
+    )
 
 
 def __create_control_plane_event(event_details, detail_type):
@@ -214,7 +260,9 @@ def __create_control_plane_event(event_details, detail_type):
     logger.info(response)
 
 
-@logger.inject_lambda_context
+@logger.inject_lambda_context(
+    correlation_id_path=correlation_paths.API_GATEWAY_HTTP, log_event=True
+)
 @tracer.capture_lambda_handler
 def handler(event: dict, context: LambdaContext) -> dict:
     logger.info(event)

scripts/sbt-aws.sh

@@ -15,6 +15,7 @@ help() {
   echo "  refresh-tokens"
   echo "  create-tenant-registration"
   echo "  get-tenant-registration <tenant_registration_id>"
+  echo "  get-all-tenant-registrations <limit> <next_token>"
   echo "  update-tenant-registration <tenant_registration_id> <key> <value>"
   echo "  delete-tenant-registration <tenant_registration_id>"
   echo "  get-tenant <tenant_id>"
@@ -218,6 +219,29 @@ get_tenant_registration() {
   fi
 }
 
+get_all_tenant_registrations() {
+  source_config
+  MY_LIMIT="${1:-10}"
+  NEXT_TOKEN="${2:-}"
+
+  if $DEBUG; then
+    echo "Getting all tenant registrations"
+  fi
+
+  RESPONSE=$(curl -G --request GET \
+    --url "${CONTROL_PLANE_API_ENDPOINT}tenant-registrations?limit=${MY_LIMIT}" \
+    --data-urlencode "next_token=${NEXT_TOKEN}" \
+    --header "Authorization: Bearer $ACCESS_TOKEN" \
+    --silent)
+
+  if $DEBUG; then
+    echo "Response: $RESPONSE"
+  else
+    echo "$RESPONSE"
+  fi
+}
+
+
 update_tenant_registration() {
   source_config
   TENANT_REGISTRATION_ID="$1"
@@ -474,6 +498,10 @@ case "$1" in
   get_tenant_registration "$2"
   ;;
 
+"get-all-tenant-registrations")
+  get_all_tenant_registrations "$2" "$3"
+  ;;
+
 "update-tenant-registration")
   if [ $# -ne 4 ]; then
     echo "Error: update-tenant-registration requires tenant registration id, key, and value"