eveAI/common/services/entitlements/license_period_services.py

from dateutil.relativedelta import relativedelta
from datetime import datetime as dt, timezone as tz, timedelta
from flask import current_app

from sqlalchemy.exc import SQLAlchemyError
from sqlalchemy.sql.expression import and_

from common.extensions import db
from common.models.entitlements import LicensePeriod, License, PeriodStatus, LicenseUsage
from common.utils.eveai_exceptions import EveAILicensePeriodsExceeded, EveAIPendingLicensePeriod, EveAINoActiveLicense
from common.utils.model_logging_utils import set_logging_information, update_logging_information


class LicensePeriodServices:
    @staticmethod
    def find_current_license_period_for_usage(tenant_id: int) -> LicensePeriod:
        """
        Find the current license period for a tenant. It ensures the status of the different license periods are adapted
        when required, and a LicenseUsage object is created if required.

        Args:
            tenant_id: The ID of the tenant to find the license period for

        Raises:
            EveAIException: and derived classes
        """
        try:
            current_app.logger.debug(f"Finding current license period for tenant {tenant_id}")
            current_date = dt.now(tz.utc).date()
            license_period = (db.session.query(LicensePeriod)
                              .filter_by(tenant_id=tenant_id)
                              .filter(and_(LicensePeriod.period_start <= current_date,
                                           LicensePeriod.period_end >= current_date))
                              .first())
            current_app.logger.debug(f"End searching for license period for tenant {tenant_id} ")
            if not license_period:
                current_app.logger.debug(f"No license period found for tenant {tenant_id} on date {current_date}")
                license_period = LicensePeriodServices._create_next_license_period_for_usage(tenant_id)
                current_app.logger.debug(f"Created license period {license_period.id} for tenant {tenant_id}")
            if license_period:
                current_app.logger.debug(f"Found license period {license_period.id} for tenant {tenant_id} "
                                         f"with status {license_period.status}")
                match license_period.status:
                    case PeriodStatus.UPCOMING | PeriodStatus.PENDING:
                        current_app.logger.debug(f"In upcoming state")
                        LicensePeriodServices._complete_last_license_period(tenant_id=tenant_id)
                        current_app.logger.debug(f"Completed last license period for tenant {tenant_id}")
                        LicensePeriodServices._activate_license_period(license_period=license_period)
                        current_app.logger.debug(f"Activated license period {license_period.id} for tenant {tenant_id}")
                        if not license_period.license_usage:
                            new_license_usage = LicenseUsage(
                                tenant_id=tenant_id,
                            )
                            new_license_usage.license_period = license_period
                            try:
                                db.session.add(new_license_usage)
                                db.session.commit()

                            except SQLAlchemyError as e:
                                db.session.rollback()
                                current_app.logger.error(
                                    f"Error creating new license usage for license period "
                                    f"{license_period.id}: {str(e)}")
                                raise e
                        if license_period.status == PeriodStatus.ACTIVE:
                            return license_period
                        else:
                            # Status is PENDING, so no prepaid payment received. There is no license period we can use.
                            # We allow for a delay of 5 days before raising an exception.
                            current_date = dt.now(tz.utc).date()
                            delta = abs(current_date - license_period.period_start)
                            if delta > timedelta(days=current_app.config.get('ENTITLEMENTS_MAX_PENDING_DAYS', 5)):
                                raise EveAIPendingLicensePeriod()
                    case PeriodStatus.ACTIVE:
                        return license_period
            else:
                raise EveAILicensePeriodsExceeded(license_id=None)
        except SQLAlchemyError as e:
            db.session.rollback()
            current_app.logger.error(f"Error finding current license period for tenant {tenant_id}: {str(e)}")
            raise e
        except Exception as e:
            raise e

    @staticmethod
    def _create_next_license_period_for_usage(tenant_id) -> LicensePeriod:
        """
        Create a new period for this license using the current license configuration

        Args:
            tenant_id: The ID of the tenant to create the period for

        Returns:
            LicensePeriod: The newly created license period
        """
        current_date = dt.now(tz.utc).date()

        # Zoek de actieve licentie voor deze tenant op de huidige datum
        the_license = (db.session.query(License)
                      .filter_by(tenant_id=tenant_id)
                      .filter(License.start_date <= current_date)
                      .filter(License.end_date >= current_date)
                      .first())

        if not the_license:
            current_app.logger.error(f"No active license found for tenant {tenant_id} on date {current_date}")
            raise EveAINoActiveLicense(tenant_id=tenant_id)
        else:
            current_app.logger.debug(f"Found active license {the_license.id} for tenant {tenant_id} "
                                    f"on date {current_date}")

        next_period_number = 1
        if the_license.periods:
            # If there are existing periods, get the next sequential number
            next_period_number = max(p.period_number for p in the_license.periods) + 1
        current_app.logger.debug(f"Next period number for tenant {tenant_id} is {next_period_number}")

        if next_period_number > the_license.nr_of_periods:
            raise EveAILicensePeriodsExceeded(license_id=the_license.id)

        new_license_period = LicensePeriod(
            license_id=the_license.id,
            tenant_id=tenant_id,
            period_number=next_period_number,
            period_start=the_license.start_date + relativedelta(months=next_period_number-1),
            period_end=the_license.start_date + relativedelta(months=next_period_number, days=-1),
            status=PeriodStatus.UPCOMING,
            upcoming_at=dt.now(tz.utc),
        )
        set_logging_information(new_license_period, dt.now(tz.utc))

        try:
            current_app.logger.debug(f"Creating next license period for tenant {tenant_id} ")
            db.session.add(new_license_period)
            db.session.commit()
            current_app.logger.info(f"Created next license period for tenant {tenant_id} "
                                    f"with id {new_license_period.id}")
            return new_license_period
        except SQLAlchemyError as e:
            db.session.rollback()
            current_app.logger.error(f"Error creating next license period for tenant {tenant_id}: {str(e)}")
            raise e

    @staticmethod
    def _activate_license_period(license_period_id: int = None, license_period: LicensePeriod = None) -> LicensePeriod:
        """
        Activate a license period

        Args:
            license_period_id: The ID of the license period to activate (optional if license_period is provided)
            license_period: The LicensePeriod object to activate (optional if license_period_id is provided)

        Returns:
            LicensePeriod: The activated license period object

        Raises:
            ValueError: If neither license_period_id nor license_period is provided
        """
        current_app.logger.debug(f"Activating license period")
        if license_period is None and license_period_id is None:
            raise ValueError("Either license_period_id or license_period must be provided")

        # Get a license period object if only ID was provided
        if license_period is None:
            current_app.logger.debug(f"Getting license period {license_period_id} to activate")
            license_period = LicensePeriod.query.get_or_404(license_period_id)

        if license_period.pending_at is not None:
            license_period.pending_at = dt.now(tz.utc)
        license_period.status = PeriodStatus.PENDING
        if license_period.prepaid_payment:
            # There is a payment received for the given period
            license_period.active_at = dt.now(tz.utc)
            license_period.status = PeriodStatus.ACTIVE

        # Copy snapshot fields from the license to the period
        the_license = License.query.get_or_404(license_period.license_id)
        license_period.currency = the_license.currency
        license_period.basic_fee = the_license.basic_fee
        license_period.max_storage_mb = the_license.max_storage_mb
        license_period.additional_storage_price = the_license.additional_storage_price
        license_period.additional_storage_bucket = the_license.additional_storage_bucket
        license_period.included_embedding_mb = the_license.included_embedding_mb
        license_period.additional_embedding_price = the_license.additional_embedding_price
        license_period.additional_embedding_bucket = the_license.additional_embedding_bucket
        license_period.included_interaction_tokens = the_license.included_interaction_tokens
        license_period.additional_interaction_token_price = the_license.additional_interaction_token_price
        license_period.additional_interaction_bucket = the_license.additional_interaction_bucket
        license_period.additional_storage_allowed = the_license.additional_storage_allowed
        license_period.additional_embedding_allowed = the_license.additional_embedding_allowed
        license_period.additional_interaction_allowed = the_license.additional_interaction_allowed

        update_logging_information(license_period, dt.now(tz.utc))

        if not license_period.license_usage:
            license_period.license_usage = LicenseUsage(
                tenant_id=license_period.tenant_id,
                license_period_id=license_period.id,
            )

        license_period.license_usage.recalculate_storage()

        try:
            db.session.add(license_period)
            db.session.add(license_period.license_usage)
            db.session.commit()
        except SQLAlchemyError as e:
            db.session.rollback()
            current_app.logger.error(f"Error activating license period {license_period_id}: {str(e)}")
            raise e

        return license_period

    @staticmethod
    def _complete_last_license_period(tenant_id) -> None:
        """
        Complete the active or pending license period for a tenant. This is done by setting the status to COMPLETED.

        Args:
            tenant_id: De ID van de tenant
        """
        # Zoek de licenseperiode voor deze tenant met status ACTIVE of PENDING
        active_period = (
            db.session.query(LicensePeriod)
            .filter_by(tenant_id=tenant_id)
            .filter(LicensePeriod.status.in_([PeriodStatus.ACTIVE, PeriodStatus.PENDING]))
            .first()
        )

        # Als er geen actieve periode gevonden is, hoeven we niets te doen
        if not active_period:
            return

        # Zet de gevonden periode op COMPLETED
        active_period.status = PeriodStatus.COMPLETED
        active_period.completed_at = dt.now(tz.utc)
        update_logging_information(active_period, dt.now(tz.utc))

        try:
            db.session.add(active_period)
            db.session.commit()
        except SQLAlchemyError as e:
            db.session.rollback()
            current_app.logger.error(f"Error completing period {active_period.id} for {tenant_id}: {str(e)}")
            raise e