From 8f41d398945b49306837b5030d54a8d6618cf072 Mon Sep 17 00:00:00 2001 From: Aleksandra Ovchinnikova Date: Thu, 30 Jul 2026 15:44:20 +0000 Subject: [PATCH] MPT-23244 Threading in redeem entitlement command + terminate entitlement changes --- backend/app/commands/redeem_entitlements.py | 251 ++++++++++-------- backend/app/db/handlers.py | 2 +- backend/app/db/models.py | 1 + backend/app/routers/entitlements.py | 11 +- backend/app/schemas/entitlements.py | 1 + ...27a2095_add_terminate_at_to_entitlement.py | 27 ++ backend/tests/api/test_entitlements_api.py | 117 ++++++-- .../commands/test_redeem_entitlements.py | 4 +- backend/tests/conftest.py | 2 + 9 files changed, 281 insertions(+), 135 deletions(-) create mode 100644 backend/migrations/versions/0df0027a2095_add_terminate_at_to_entitlement.py diff --git a/backend/app/commands/redeem_entitlements.py b/backend/app/commands/redeem_entitlements.py index af2a503..4c7b701 100644 --- a/backend/app/commands/redeem_entitlements.py +++ b/backend/app/commands/redeem_entitlements.py @@ -1,5 +1,7 @@ import asyncio import logging +from collections.abc import Sequence +from dataclasses import dataclass, field from datetime import UTC, datetime import httpx @@ -24,12 +26,23 @@ logger = logging.getLogger(__name__) -BATCH_SIZE = 100 + +@dataclass +class RedeemResult: + """Outcome of redeeming entitlements for a single organization, aggregated for notifications.""" + + organization_id: str + redeemed_rows: list[tuple[str, str, str, str]] = field(default_factory=list) + succeeded: bool = True + error: str | None = None async def fetch_datasources_for_organization(settings: Settings, organization_id: str) -> dict: - client = OptscaleClient(settings) - response = await client.fetch_datasources_for_organization(organization_id, details=False) + async with OptscaleClient(settings) as optscale_client: + response = await optscale_client.fetch_datasources_for_organization( + organization_id, # type: ignore[arg-type] + details=False, + ) return response.json()["cloud_accounts"] @@ -53,18 +66,17 @@ async def process_datasource( organization: Organization, entitlement_handler: EntitlementHandler, ffc_api_client: FFCAPIClient, -): +) -> tuple[str, str, str, str] | None: datasource_id = datasource["account_id"] datasource_type = datasource["type"] datasource_name = datasource["name"] - type_name = datasource_type.split("_")[0].capitalize() match datasource_type: case "azure_tenant" | "gcp_tenant": logger.debug( f"Found {datasource_id} {datasource_name} of type {datasource_type}, " "skip containers!" ) - return + return None case "azure_cnr" | "aws_cnr" | "gcp_cnr": type_name = datasource["type"].split("_")[0].capitalize() logger.info( @@ -75,119 +87,144 @@ async def process_datasource( f"Found {datasource_id} {datasource_name} of type {datasource_type}, " "unsupported type!" ) - return - try: - instance = await entitlement_handler.first( - where_clauses=[ - Entitlement.datasource_id == datasource_id, - Entitlement.status == EntitlementStatus.NEW, - ] + return None + + instance = await entitlement_handler.first( + where_clauses=[ + Entitlement.datasource_id == datasource_id, + Entitlement.status == EntitlementStatus.NEW, + ] + ) + if instance: + await entitlement_handler.update( + instance, + data={ + "status": EntitlementStatus.ACTIVE, + "redeemed_at": instance.redeem_at or datetime.now(UTC), + "redeemed_by": organization, + "linked_datasource_id": datasource["id"], + "linked_datasource_type": datasource["type"], + "linked_datasource_name": datasource["name"], + }, ) - if instance: - updated_entitlement = await entitlement_handler.update( - instance, - data={ - "status": EntitlementStatus.ACTIVE, - "redeemed_at": instance.redeem_at or datetime.now(UTC), - "redeemed_by": organization, - "linked_datasource_id": datasource["id"], - "linked_datasource_type": datasource["type"], - "linked_datasource_name": datasource["name"], - }, - ) - await create_entitlement_tag_for_datasource( - ffc_api_client=ffc_api_client, - entitlement_id=instance.id, - datasource_id=datasource["id"], + await create_entitlement_tag_for_datasource( + ffc_api_client=ffc_api_client, + entitlement_id=instance.id, + datasource_id=datasource["id"], + ) + logger.info( + f"The entitlement {instance.id} - {instance.name} " + f"owned by {instance.owner.id} - {instance.owner.name} " + f"has been redeemed by {organization.id} - {organization.name} " + f"for datasource {datasource_id} - {datasource_name}." + ) + return ( + f"{instance.id}\t/\t{instance.name}", + f"{instance.owner.id}\t/\t{instance.owner.name}", + f"{organization.id}\t/\t{organization.name}", + f"{datasource_id}\t/\t{datasource_name}", + ) + else: + logger.info(f"Entitlement not found for datasource {datasource_id} - {datasource_name}.") + return None + + +async def process_organization( + organization: Organization, + settings: Settings, + semaphore: asyncio.Semaphore, +) -> RedeemResult: + result = RedeemResult(organization_id=organization.id) + + async with semaphore: + logger.info( + f"Fetching datasources for organization: {organization.id} - {organization.name}..." + ) + try: + datasources = await fetch_datasources_for_organization( + settings, + organization.linked_organization_id, # type: ignore[arg-type] ) - msg = ( - f"The entitlement {instance.id} - {instance.name} " - f"owned by {instance.owner.id} - {instance.owner.name} " - f"has been redeemed by {organization.id} - {organization.name} " - f"for datasource {datasource_id} - {datasource_name}." + + async with session_factory() as session, FFCAPIClient(settings) as ffc_api_client: + entitlement_handler = EntitlementHandler(session) + async with session.begin(): + for datasource in datasources: + row = await process_datasource( + datasource, + organization, + entitlement_handler, + ffc_api_client, + ) + if row is not None: + result.redeemed_rows.append(row) + + except (httpx.HTTPError, httpx.ReadTimeout) as e: + message = ( + f"Failed to fetch datasources for organization {organization.id} " + f"({type(e).__name__}): {str(e) or repr(e)}" ) - logger.info(msg) - return updated_entitlement - else: - logger.info( - f"Entitlement not found for datasource {datasource_id} - {datasource_name}." + logger.error(message) + result.succeeded = False + result.error = str(message) + except DatabaseError as e: # pragma: no cover + message = ( + f"Failed to process or update datasources " + f"for organization {organization.id}: {str(e)}" ) - - except DatabaseError as e: # pragma: no cover - msg = ( - f"An error occurred while updating the entitlement for " - f"{datasource_id} - {datasource_name}: {e}" + logger.error(message) + result.succeeded = False + result.error = str(message) + + return result + + +async def notify_results(results: Sequence[RedeemResult]) -> None: + redeemed = [row for result in results if result.succeeded for row in result.redeemed_rows] + failed = [result for result in results if not result.succeeded] + + if redeemed: + msg = "Entitlement has" if len(redeemed) == 1 else "Entitlements have" + msg = f"{len(redeemed)} {msg} been successfully redeemed." + await send_info( + "Redeem Entitlements Success", + msg, + details=NotificationDetails( + header=( + ColumnHeader("Entitlement", width="stretch"), + ColumnHeader("Owner", width="stretch"), + ColumnHeader("Organization", width="stretch"), + ColumnHeader("Datasource", width="stretch"), + ), + rows=redeemed, + ), + ) + if failed: + detail = "; ".join(f"{result.organization_id}: {result.error}" for result in failed) + await send_exception( + "Redeem Entitlements Error", + f"{len(failed)} organizations failed to process: {detail}", ) - logger.error(msg) - await send_exception("Redeem Entitlements Error", msg) @capture_telemetry_cli_command(__name__, "Redeem Entitlements") async def redeem_entitlements(settings: Settings): - # FIXME: Long-lived DB transaction (making API calls inside the transaction) + semaphore = asyncio.Semaphore(settings.max_parallel_tasks) - async with session_factory.begin() as session: + async with session_factory() as session: organization_handler = OrganizationHandler(session) - entitlement_handler = EntitlementHandler(session) - ffc_api_client = FFCAPIClient(settings) - - async for organization in organization_handler.stream_scalars( - extra_conditions=[Organization.status == OrganizationStatus.ACTIVE], + organizations = await organization_handler.query_db( + where_clauses=[Organization.status == OrganizationStatus.ACTIVE], order_by=[Organization.created_at], - batch_size=BATCH_SIZE, - ): - logger.info( - f"Fetching datasources for organization: {organization.id} - {organization.name}..." - ) - datasources = None - try: - datasources = await fetch_datasources_for_organization( - settings, - organization.linked_organization_id, # type: ignore - ) - except (httpx.HTTPError, httpx.ReadTimeout) as e: - message = ( - f"Failed to fetch datasources for organization {organization.id} " - f"({type(e).__name__}): {str(e) or repr(e)}" - ) - logger.error(message) - await send_exception("Redeem Entitlements Error", message) - continue - redeemed_entitlements = [] - for datasource in datasources: - entitlement = await process_datasource( - datasource, - organization, - entitlement_handler, - ffc_api_client, - ) - if entitlement: - redeemed_entitlements.append(entitlement) - - if len(redeemed_entitlements) > 0: - msg = "Entitlement has" if len(redeemed_entitlements) == 1 else "Entitlements have" - msg = f"{len(redeemed_entitlements)} {msg} been successfully redeemed." - await send_info( - "Redeem Entitlements Success", - msg, - details=NotificationDetails( - header=( - ColumnHeader("Entitlement", width="stretch"), - ColumnHeader("Owner", width="stretch"), - ColumnHeader("Organization", width="stretch"), - ColumnHeader("Datasource", width="stretch"), - ), - rows=[ - ( - f"{ent.id}\t/\t{ent.name}", - f"{ent.owner.id}\t/\t{ent.owner.name}", - f"{ent.redeemed_by.id}\t/\t{ent.redeemed_by.name}", # type: ignore - f"{ent.datasource_id}\t/\t{ent.linked_datasource_name}", - ) - for ent in redeemed_entitlements - ], - ), - ) + ) + + tasks = [ + asyncio.create_task(process_organization(organization, settings, semaphore)) + for organization in organizations + ] + + results = await asyncio.gather(*tasks) + await notify_results(results) def command(ctx: typer.Context): diff --git a/backend/app/db/handlers.py b/backend/app/db/handlers.py index 3a77ffa..46c66a7 100644 --- a/backend/app/db/handlers.py +++ b/backend/app/db/handlers.py @@ -386,7 +386,7 @@ async def terminate(self, entitlement: Entitlement) -> Entitlement: entitlement, data={ "status": EntitlementStatus.TERMINATED, - "terminated_at": datetime.now(UTC), + "terminated_at": entitlement.terminate_at or datetime.now(UTC), "terminated_by": auth_context.get().get_actor(), }, ) diff --git a/backend/app/db/models.py b/backend/app/db/models.py index 0fb35d2..8e88044 100644 --- a/backend/app/db/models.py +++ b/backend/app/db/models.py @@ -339,6 +339,7 @@ class Entitlement(Base, HumanReadablePKMixin, AuditableMixin): terminated_by: Mapped[Actor | None] = relationship(foreign_keys=[terminated_by_id]) redeem_at: Mapped[datetime.datetime | None] = mapped_column(sa.DateTime(timezone=True)) + terminate_at: Mapped[datetime.datetime | None] = mapped_column(sa.DateTime(timezone=True)) class AdditionalAdminRequest(Base, HumanReadablePKMixin, AuditableMixin): diff --git a/backend/app/routers/entitlements.py b/backend/app/routers/entitlements.py index 1d16352..7384a23 100644 --- a/backend/app/routers/entitlements.py +++ b/backend/app/routers/entitlements.py @@ -94,6 +94,11 @@ async def create_entitlement( status_code=status.HTTP_400_BAD_REQUEST, detail="Affiliate accounts cannot provide a redeem_at for an Entitlement.", ) + if data.terminate_at: + raise HTTPException( + status_code=status.HTTP_400_BAD_REQUEST, + detail="Affiliate accounts cannot provide a terminate_at for an Entitlement.", + ) else: if not data.owner: raise HTTPException( @@ -124,7 +129,11 @@ async def get_entitlement_by_id( return convert_model_to_schema(EntitlementRead, entitlement) -@router.post("/{id}/terminate", response_model=EntitlementRead) +@router.post( + "/{id}/terminate", + response_model=EntitlementRead, + dependencies=[Depends(AuthorizedAccountTypes(AccountType.ADMIN))], +) async def terminate_entitlement( entitlement: Annotated[Entitlement, Depends(fetch_entitlement_or_404)], entitlement_repo: EntitlementRepository, diff --git a/backend/app/schemas/entitlements.py b/backend/app/schemas/entitlements.py index 0d5dbcc..60d04c8 100644 --- a/backend/app/schemas/entitlements.py +++ b/backend/app/schemas/entitlements.py @@ -24,6 +24,7 @@ class EntitlementBase(BaseSchema): str, Field(min_length=1, max_length=255, examples=["1098a2fa-07c0-4f40-96c7-3bf32a213e0e"]) ] redeem_at: datetime.datetime | None = None + terminate_at: datetime.datetime | None = None class EntitlementCreate(EntitlementBase): diff --git a/backend/migrations/versions/0df0027a2095_add_terminate_at_to_entitlement.py b/backend/migrations/versions/0df0027a2095_add_terminate_at_to_entitlement.py new file mode 100644 index 0000000..6af472c --- /dev/null +++ b/backend/migrations/versions/0df0027a2095_add_terminate_at_to_entitlement.py @@ -0,0 +1,27 @@ +"""add terminate_at to entitlement + +Revision ID: 0df0027a2095 +Revises: 91e67d9ca2be +Create Date: 2026-07-30 14:33:23.000064 + +""" +from typing import Sequence, Union + +from alembic import op +import sqlalchemy as sa +import sqlalchemy_utils + + +# revision identifiers, used by Alembic. +revision: str = '0df0027a2095' +down_revision: Union[str, None] = '91e67d9ca2be' +branch_labels: Union[str, Sequence[str], None] = None +depends_on: Union[str, Sequence[str], None] = None + + +def upgrade() -> None: + op.add_column('entitlements', sa.Column('terminate_at', sa.DateTime(timezone=True), nullable=True)) + + +def downgrade() -> None: + op.drop_column('entitlements', 'terminate_at') diff --git a/backend/tests/api/test_entitlements_api.py b/backend/tests/api/test_entitlements_api.py index 086ce87..1fb441f 100644 --- a/backend/tests/api/test_entitlements_api.py +++ b/backend/tests/api/test_entitlements_api.py @@ -1,4 +1,4 @@ -from datetime import UTC, datetime +from datetime import UTC, datetime, timedelta import pytest from httpx import AsyncClient @@ -7,7 +7,7 @@ from sqlalchemy.ext.asyncio import AsyncSession from app.conf import Settings -from app.db.models import Account, Entitlement, System +from app.db.models import Account, Entitlement, System, User from app.enums import AccountStatus, DatasourceType, EntitlementStatus, OrganizationStatus from tests.types import ModelFactory @@ -123,6 +123,27 @@ async def test_create_entitlement_by_affiliate_with_redeem_at( assert error_msg == "Affiliate accounts cannot provide a redeem_at for an Entitlement." +async def test_create_entitlement_by_affiliate_with_terminate_at( + api_client: AsyncClient, + gcp_jwt_token: str, +): + response = await api_client.post( + "/entitlements", + headers={"Authorization": f"Bearer {gcp_jwt_token}"}, + json={ + "name": "AWS", + "affiliate_external_id": "EXTERNAL_ID_987123", + "datasource_id": "ds-id", + "terminate_at": datetime.now(UTC).isoformat(), + }, + ) + + assert response.status_code == 400 + error_msg = response.json()["detail"] + + assert error_msg == "Affiliate accounts cannot provide a terminate_at for an Entitlement." + + async def test_create_entitlement_by_admin_with_owner( api_client: AsyncClient, ffc_jwt_token: str, @@ -137,6 +158,7 @@ async def test_create_entitlement_by_admin_with_owner( "datasource_id": "ds-id", "owner": {"id": affiliate_account.id}, "redeem_at": datetime.now(UTC).isoformat(), + "terminate_at": datetime.now(UTC).isoformat(), }, ) @@ -144,6 +166,7 @@ async def test_create_entitlement_by_admin_with_owner( entitlement = response.json() assert entitlement["owner"]["id"] == affiliate_account.id assert entitlement["redeem_at"] is not None + assert entitlement["terminate_at"] is not None async def test_create_entitlement_by_admin_without_owner( @@ -508,8 +531,9 @@ async def test_get_invalid_id_format(api_client: AsyncClient, gcp_jwt_token: str async def test_terminate_entitlement_success( entitlement_gcp: Entitlement, - api_client: AsyncClient, - gcp_jwt_token: str, + admin_client: AsyncClient, + admin_user_token: str, + admin_user: User, gcp_extension: System, db_session: AsyncSession, ): @@ -524,9 +548,9 @@ async def test_terminate_entitlement_success( await db_session.refresh(entitlement_gcp) request_start_dt = datetime.now(UTC) - response = await api_client.post( + response = await admin_client.post( f"/entitlements/{entitlement_gcp.id}/terminate", - headers={"Authorization": f"Bearer {gcp_jwt_token}"}, + headers={"Authorization": f"Bearer {admin_user_token}"}, ) request_end_dt = datetime.now(UTC) @@ -541,26 +565,59 @@ async def test_terminate_entitlement_success( assert entitlement_gcp.status == EntitlementStatus.TERMINATED assert entitlement_gcp.terminated_at is not None assert request_start_dt < entitlement_gcp.terminated_at < request_end_dt - assert entitlement_gcp.terminated_by_id == gcp_extension.id + assert entitlement_gcp.terminated_by_id == admin_user.id assert ( datetime.fromisoformat(data["events"]["terminated"]["at"]) == entitlement_gcp.terminated_at ) - assert data["events"]["terminated"]["by"]["id"] == gcp_extension.id - assert data["events"]["terminated"]["by"]["type"] == gcp_extension.type._value_ - assert data["events"]["terminated"]["by"]["name"] == gcp_extension.name + assert data["events"]["terminated"]["by"]["id"] == admin_user.id + assert data["events"]["terminated"]["by"]["type"] == admin_user.type._value_ + assert data["events"]["terminated"]["by"]["name"] == admin_user.name + + +async def test_terminate_entitlement_with_set_terminate_at( + entitlement_factory: ModelFactory[Entitlement], + admin_client: AsyncClient, + admin_user_token: str, + gcp_extension: System, + db_session: AsyncSession, +): + entitlement_gcp = await entitlement_factory( + name="GCP", + status=EntitlementStatus.ACTIVE, + owner=gcp_extension.owner, + created_by=gcp_extension, + updated_by=gcp_extension, + terminate_at=datetime.now(UTC) - timedelta(days=1), + ) + + response = await admin_client.post( + f"/entitlements/{entitlement_gcp.id}/terminate", + headers={"Authorization": f"Bearer {admin_user_token}"}, + ) + + assert response.status_code == 200 + data = response.json() + + assert data["id"] == str(entitlement_gcp.id) + assert data["status"] == "terminated" + + await db_session.refresh(entitlement_gcp) + terminated = data["events"]["terminated"]["at"] + assert datetime.fromisoformat(terminated) == entitlement_gcp.terminate_at + assert entitlement_gcp.status == EntitlementStatus.TERMINATED async def test_terminate_new_entitlement( entitlement_gcp: Entitlement, - api_client: AsyncClient, - gcp_jwt_token: str, + admin_client: AsyncClient, + admin_user_token: str, ): assert entitlement_gcp.status == EntitlementStatus.NEW - response = await api_client.post( + response = await admin_client.post( f"/entitlements/{entitlement_gcp.id}/terminate", - headers={"Authorization": f"Bearer {gcp_jwt_token}"}, + headers={"Authorization": f"Bearer {admin_user_token}"}, ) assert response.status_code == 400 @@ -571,8 +628,8 @@ async def test_terminate_new_entitlement( async def test_terminate_already_terminated_entitlement( entitlement_gcp: Entitlement, - api_client: AsyncClient, - gcp_jwt_token: str, + admin_client: AsyncClient, + admin_user_token: str, db_session: AsyncSession, ): entitlement_gcp.status = EntitlementStatus.TERMINATED @@ -580,9 +637,9 @@ async def test_terminate_already_terminated_entitlement( db_session.add(entitlement_gcp) await db_session.commit() - response = await api_client.post( + response = await admin_client.post( f"/entitlements/{entitlement_gcp.id}/terminate", - headers={"Authorization": f"Bearer {gcp_jwt_token}"}, + headers={"Authorization": f"Bearer {admin_user_token}"}, ) assert response.status_code == 400 @@ -591,16 +648,15 @@ async def test_terminate_already_terminated_entitlement( assert error_msg == "Entitlement is already terminated." -async def test_terminate_non_existant_entitlement( - api_client: AsyncClient, - gcp_jwt_token: str, - gcp_extension: System, +async def test_terminate_non_existing_entitlement( + admin_client: AsyncClient, + admin_user_token: str, db_session: AsyncSession, ): entitlement_id = "FENT-1234-5678-9012" - response = await api_client.post( + response = await admin_client.post( f"/entitlements/{entitlement_id}/terminate", - headers={"Authorization": f"Bearer {gcp_jwt_token}"}, + headers={"Authorization": f"Bearer {admin_user_token}"}, ) assert response.status_code == 404 @@ -609,6 +665,19 @@ async def test_terminate_non_existant_entitlement( assert error_msg == f"Entitlement with ID `{entitlement_id}` wasn't found." +async def test_terminate_entitlement_by_affiliate( + affiliate_client: AsyncClient, + entitlement_gcp: Entitlement, + gcp_jwt_token: str, +): + response = await affiliate_client.post( + f"/entitlements/{entitlement_gcp.id}/terminate", + headers={"Authorization": f"Bearer {gcp_jwt_token}"}, + ) + + assert response.status_code == 403 + + # ================== # Redeem Entitlement # ================== diff --git a/backend/tests/commands/test_redeem_entitlements.py b/backend/tests/commands/test_redeem_entitlements.py index 9d16700..e040896 100644 --- a/backend/tests/commands/test_redeem_entitlements.py +++ b/backend/tests/commands/test_redeem_entitlements.py @@ -149,8 +149,8 @@ async def test_redeeem_entitlements_error_fetching_datasources( mocker_send_exception.assert_awaited_once_with( "Redeem Entitlements Error", ( - f"Failed to fetch datasources for organization " - f"{apple_inc_organization.id} (ReadTimeout): timed out" + f"1 organizations failed to process: {apple_inc_organization.id}: Failed to fetch " + f"datasources for organization {apple_inc_organization.id} (ReadTimeout): timed out" ), ) diff --git a/backend/tests/conftest.py b/backend/tests/conftest.py index 0e0a453..494f8c5 100644 --- a/backend/tests/conftest.py +++ b/backend/tests/conftest.py @@ -275,6 +275,7 @@ async def _entitlement( redeem_at: datetime | None = None, redeemed_at: datetime | None = None, redeemed_by: Organization | None = None, + terminate_at: datetime | None = None, terminated_at: datetime | None = None, ) -> Entitlement: entitlement = Entitlement( @@ -290,6 +291,7 @@ async def _entitlement( redeem_at=redeem_at, redeemed_at=redeemed_at, redeemed_by=redeemed_by, + terminate_at=terminate_at, terminated_at=terminated_at, ) db_session.add(entitlement)