from http import HTTPStatus from typing import Annotated from aws_lambda_powertools import Logger from aws_lambda_powertools.event_handler.api_gateway import Router from aws_lambda_powertools.event_handler.exceptions import NotFoundError from aws_lambda_powertools.event_handler.openapi.params import Body, Query from layercake.dynamodb import DynamoDBPersistenceLayer, KeyPair, PartitionKey from pydantic import FutureDatetime from api_gateway import JSONResponse from boto3clients import dynamodb_client from config import ENROLLMENT_TABLE from ...enrollments.enroll import Enrollment, Org, Subscription, enroll_now logger = Logger(__name__) router = Router() dyn = DynamoDBPersistenceLayer(ENROLLMENT_TABLE, dynamodb_client) class ScheduledNotFoundError(NotFoundError): ... @router.get('//enrollments/scheduled') def scheduled( org_id: str, start_key: Annotated[str | None, Query] = None, ): return dyn.collection.query( key=PartitionKey(f'SCHEDULED#ORG#{org_id}'), start_key=start_key, limit=150, ) @router.delete('//enrollments/scheduled') def cancel( org_id: str, scheduled_for: Annotated[FutureDatetime, Body(embed=True)], lock_hash: Annotated[str, Body(embed=True)], ): with dyn.transact_writer() as transact: transact.delete( key=KeyPair( pk=f'SCHEDULED#ORG#{org_id}', sk=f'{scheduled_for.isoformat()}#{lock_hash}', ), cond_expr='attribute_exists(sk)', exc_cls=ScheduledNotFoundError, ) transact.delete( key=KeyPair( pk='LOCK#SCHEDULED', sk=lock_hash, ), ) return JSONResponse(status_code=HTTPStatus.NO_CONTENT) @router.post('//enrollments/scheduled/proceed') def proceed( org_id: str, scheduled_for: Annotated[FutureDatetime, Body(embed=True)], lock_hash: Annotated[str, Body(embed=True)], ): pk = f'SCHEDULED#ORG#{org_id}' sk = f'{scheduled_for.isoformat()}#{lock_hash}' scheduled = dyn.collection.get_item( KeyPair(pk, sk), exc_cls=ScheduledNotFoundError, ) org = Org( id=org_id, name=scheduled['org_name'], ) subscription = Subscription( billing_day=scheduled['subscription_billing_day'], ) try: enrollment = enroll_now( Enrollment( user=scheduled['user'], course=scheduled['course'], ), { 'org': org, 'subscription': subscription, 'created_by': router.context['user'], }, ) with dyn.transact_writer() as transact: transact.delete(key=KeyPair(pk, sk)) transact.delete(key=KeyPair('LOCK#SCHEDULED', lock_hash)) except Exception: raise else: return JSONResponse( status_code=HTTPStatus.CREATED, body=enrollment, )