saladeaula.digital/enrollments-events/app/enrollment.py

from abc import ABC
from dataclasses import dataclass
from datetime import datetime, timedelta
from enum import Enum
from typing import Any, Literal, TypedDict
from uuid import uuid4

from layercake.dateutils import now, ttl
from layercake.dynamodb import DynamoDBPersistenceLayer
from layercake.extra_types import CpfStr, NameStr
from layercake.strutils import md5_hash
from pydantic import (
    UUID4,
    BaseModel,
    ConfigDict,
    EmailStr,
    Field,
)
from typing_extensions import NotRequired

from config import DEDUP_WINDOW_OFFSET_DAYS


class User(BaseModel):
    model_config = ConfigDict(arbitrary_types_allowed=True)

    id: UUID4 | str
    name: NameStr
    email: EmailStr
    email_verified: bool = False
    cpf: CpfStr | None = None


class Course(BaseModel):
    id: UUID4 | str
    name: str
    access_period: int = 90  # 3 months


class Enrollment(BaseModel):
    id: UUID4 | str = Field(default_factory=uuid4)
    user: User
    course: Course
    progress: int = Field(default=0, ge=0, le=100)
    status: Literal['PENDING'] = 'PENDING'

    def model_dump(
        self,
        exclude=None,
        *args,
        **kwargs,
    ) -> dict[str, Any]:
        return super().model_dump(
            exclude={'user': {'email_verified'}},
            *args,
            **kwargs,
        )


Org = TypedDict('Org', {'org_id': str, 'name': str})

CreatedBy = TypedDict('CreatedBy', {'id': str, 'name': str})

DeduplicationWindow = TypedDict('DeduplicationWindow', {'offset_days': int})

Subscription = TypedDict(
    'Subscription',
    {
        'org_id': str,
        'billing_day': int,
        'billing_period': NotRequired[str],
    },
)


class Kind(str, Enum):
    ORDER = 'ORDER'
    ENROLLMENT = 'ENROLLMENT'


@dataclass(frozen=True)
class LinkedEntity(ABC):
    id: str
    kind: Kind
    table_name: str | None = None


class DeduplicationConflictError(Exception):
    def __init__(self, *args):
        super().__init__('Enrollment already exists')


def enroll(
    enrollment: Enrollment,
    *,
    org: Org | None = None,
    subscription: Subscription | None = None,
    created_by: CreatedBy | None = None,
    scheduled_at: datetime | None = None,
    linked_entities: frozenset[LinkedEntity] = frozenset(),
    deduplication_window: DeduplicationWindow | None = None,
    persistence_layer: DynamoDBPersistenceLayer,
) -> bool:
    now_ = now()
    user = enrollment.user
    course = enrollment.course
    lock_hash = md5_hash(f'{user.id}{course.id}')
    access_expires_at = now_ + timedelta(days=course.access_period)

    with persistence_layer.transact_writer() as transact:
        transact.put(
            item={
                'sk': '0',
                'created_at': now_,
                'access_expires_at': access_expires_at,
                **enrollment.model_dump(),
            }
            | ({'org_id': org['org_id']} if org else {})
            | ({'subscription_covered': True} if subscription else {})
            | ({'scheduled_at': scheduled_at} if scheduled_at else {})
        )

        # Relationships between this enrollment and its related entities
        for entity in linked_entities:
            # Parent knows the child
            transact.put(
                item={
                    'id': entity.id,
                    'sk': f'LINKED_ENTITIES#CHILD#ENROLLMENT#{enrollment.id}',
                    'created_at': now_,
                },
                cond_expr='attribute_not_exists(sk)',
                table_name=entity.table_name,
            )

            # Child knows the parent
            transact.put(
                item={
                    'id': enrollment.id,
                    'sk': f'LINKED_ENTITIES#PARENT#{entity.kind.value}#{entity.id}',
                    'created_at': now_,
                },
                cond_expr='attribute_not_exists(sk)',
            )

        if org:
            transact.put(
                item={
                    'id': enrollment.id,
                    'sk': 'ORG',
                    'created_at': now_,
                }
                | org
            )

        if subscription:
            transact.put(
                item={
                    'id': enrollment.id,
                    'sk': 'METADATA#SUBSCRIPTION_COVERED',
                    'created_at': now_,
                }
                | subscription,
            )

        if created_by:
            transact.put(
                item={
                    'id': enrollment.id,
                    'sk': 'CREATED_BY',
                    'created_by': created_by,
                    'created_at': now_,
                }
            )

        # Prevents the user from enrolling in the same course again until
        # the deduplication window expires or is removed.
        offset_days = (
            int(deduplication_window['offset_days'])
            if deduplication_window
            else DEDUP_WINDOW_OFFSET_DAYS
        )
        dedup_lock_ttl = ttl(
            start_dt=now_,
            days=course.access_period - offset_days,
        )
        transact.put(
            item={
                'id': 'LOCK',
                'sk': lock_hash,
                'enrollment_id': enrollment.id,
                'created_at': now_,
                'ttl': dedup_lock_ttl,
            },
            cond_expr='attribute_not_exists(sk)',
            exc_cls=DeduplicationConflictError,
        )
        transact.put(
            item={
                'id': enrollment.id,
                'sk': 'LOCK',
                'hash': lock_hash,
                'created_at': now_,
                'ttl': dedup_lock_ttl,
            },
        )

        # The deduplication window can be recalculated based on user settings.
        if deduplication_window:
            transact.put(
                item={
                    'id': enrollment.id,
                    'sk': 'METADATA#DEDUPLICATION_WINDOW',
                    'offset_days': offset_days,
                    'created_at': now_,
                },
            )

    return True