from http import HTTPStatus from typing import Annotated from uuid import uuid4 from aws_lambda_powertools.event_handler.api_gateway import Router from aws_lambda_powertools.event_handler.exceptions import ( NotFoundError, ServiceError, ) from aws_lambda_powertools.event_handler.openapi.params import Body from layercake.dateutils import now from layercake.dynamodb import DynamoDBPersistenceLayer, KeyPair, SortKey from layercake.extra_types import CnpjStr, CpfStr, NameStr from pydantic import BaseModel, EmailStr, Field from api_gateway import JSONResponse from boto3clients import dynamodb_client from config import USER_TABLE router = Router() dyn = DynamoDBPersistenceLayer(USER_TABLE, dynamodb_client) class Org(BaseModel): id: str | None = Field(default=None, exclude=True) name: str cnpj: CnpjStr class User(BaseModel): name: NameStr cpf: CpfStr email: EmailStr class CPFConflictError(Exception): ... class EmailConflictError(Exception): ... class UserConflictError(ServiceError): def __init__(self, msg: str | dict): super().__init__(HTTPStatus.CONFLICT, msg) class UserMissingError(NotFoundError): ... class OrgMissingError(NotFoundError): ... @router.post('//users') def add( org_id: str, user: Annotated[User, Body(embed=True)], org: Annotated[Org, Body(embed=True)], ): org.id = org_id if _create_user(user, org): return JSONResponse(HTTPStatus.CREATED) user_id = _get_user_id(user) _add_member(user_id, org) return JSONResponse(HTTPStatus.NO_CONTENT) @router.delete('//users/') def unlink(org_id: str, user_id: str): with dyn.transact_writer() as transact: transact.delete( key=KeyPair( pk=f'orgmembers#{org_id}', # Post-migration: uncomment the following line # pk=f'MEMBER#ORG#{org_id}', sk=user_id, ) ) transact.delete( key=KeyPair(org_id, f'admins#{user_id}'), # Post-migration: uncomment the following line # key=KeyPair(org_id, f'ADMIN#{user_id}'), ) transact.delete( key=KeyPair( pk=user_id, sk=f'orgs#{org_id}', # Post-migration: uncomment the following line # pk=f'ORG#{org_id}', ) ) transact.update( key=KeyPair(user_id, '0'), update_expr='DELETE tenant_id :org_id', expr_attr_values={':org_id': {org_id}}, ) return JSONResponse(HTTPStatus.NO_CONTENT) def _create_user(user: User, org: Org) -> bool: now_ = now() user_id = uuid4() try: with dyn.transact_writer() as transact: transact.put( item={ **user.model_dump(), 'id': user_id, 'sk': '0', 'email_verified': False, 'tenant_id': {org.id}, # Post-migration: uncomment the folloing line # 'org_id': {org.id}, 'created_at': now_, }, ) transact.put( item={ 'id': user_id, # Post-migration: rename `emails` to `EMAIL` 'sk': f'emails#{user.email}', 'email_verified': False, 'email_primary': True, 'created_at': now_, } ) transact.put( item={ # Post-migration: rename `cpf` to `CPF` 'id': 'cpf', 'sk': user.cpf, 'created_at': now_, }, cond_expr='attribute_not_exists(sk)', exc_cls=CPFConflictError, ) transact.put( item={ # Post-migration: rename `email` to `EMAIL` 'id': 'email', 'sk': user.email, 'created_at': now_, }, cond_expr='attribute_not_exists(sk)', exc_cls=EmailConflictError, ) transact.put( item={ 'id': user_id, 'sk': f'orgs#{org.id}', # Post-migration: uncomment the following line # pk=f'ORG#{org.id}', 'name': org.name, 'cnpj': org.cnpj, 'created_at': now_, } ) transact.put( item={ 'id': f'orgmembers#{org.id}', # Post-migration: uncomment the following line # pk=f'MEMBER#ORG#{org_id}', 'sk': user_id, 'created_at': now_, } ) transact.condition( key=KeyPair(org.id, '0'), # type: ignore cond_expr='attribute_exists(sk)', exc_cls=OrgMissingError, ) except (CPFConflictError, EmailConflictError): return False else: return True def _add_member(user_id: str, org: Org) -> None: now_ = now() with dyn.transact_writer() as transact: transact.update( key=KeyPair(user_id, '0'), update_expr='ADD tenant_id :org_id', # Post-migration: uncomment the following line # update_expr='ADD tenant_id :org_id', expr_attr_values={ ':org_id': {org.id}, }, cond_expr='attribute_exists(sk)', exc_cls=UserMissingError, ) transact.put( item={ 'id': user_id, # Post-migration: rename `orgs` to `ORG` 'sk': f'orgs#{org.id}', 'name': org.name, 'cnpj': org.cnpj, 'created_at': now_, } ) transact.put( item={ # Post-migration: rename `orgmembers` to `ORGMEMBER` 'id': f'orgmembers#{org.id}', 'sk': user_id, 'created_at': now_, }, cond_expr='attribute_not_exists(sk)', exc_cls=UserConflictError, ) transact.condition( key=KeyPair(org.id, '0'), # type: ignore cond_expr='attribute_exists(sk)', exc_cls=OrgMissingError, ) def _get_user_id(user: User) -> str: user_id = dyn.collection.get_items( KeyPair( pk='email', sk=SortKey(user.email, path_spec='user_id'), rename_key='id', ) + KeyPair( pk='cpf', sk=SortKey(user.cpf, path_spec='user_id'), rename_key='id', ), flatten_top=False, ).get('id') if not user_id: raise UserMissingError() return user_id