saladeaula.digital/api.saladeaula.digital/app/routes/orgs/users/__init__.py

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('/<org_id>/users')
def add_user(
    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)


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,
                    '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,
                    # 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_,
                }
            )
            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 org_id :org_id',
            expr_attr_values={
                ':org_id': {org.id},
            },
        )
        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