improvements and tests

Author: matusdrobuliak66

packages/postgres-database/src/simcore_postgres_database/migration/versions/a3a58471b0f1_add_credit_transaction_classification_.py

@@ -1,7 +1,7 @@
 """add credit transaction classification enums
 
 Revision ID: a3a58471b0f1
-Revises: 307017ee1a49
+Revises: f19905923355
 Create Date: 2025-01-14 13:44:05.025647+00:00
 
 """
@@ -10,7 +10,7 @@
 
 # revision identifiers, used by Alembic.
 revision = "a3a58471b0f1"
-down_revision = "307017ee1a49"
+down_revision = "f19905923355"
 branch_labels = None
 depends_on = None
 

packages/service-library/src/servicelib/rabbitmq/rpc_interfaces/resource_usage_tracker/credit_transactions.py

@@ -32,16 +32,34 @@ async def get_wallet_total_credits(
     *,
     product_name: ProductName,
     wallet_id: WalletID,
-    transaction_status: CreditTransactionStatus | None = None,
-    project_id: ProjectID | None = None,
 ) -> WalletTotalCredits:
     result = await rabbitmq_rpc_client.request(
         RESOURCE_USAGE_TRACKER_RPC_NAMESPACE,
         _RPC_METHOD_NAME_ADAPTER.validate_python("get_wallet_total_credits"),
         product_name=product_name,
         wallet_id=wallet_id,
-        transaction_status=transaction_status,
+        timeout_s=_DEFAULT_TIMEOUT_S,
+    )
+    assert isinstance(result, WalletTotalCredits)  # nosec
+    return result
+
+
+@log_decorator(_logger, level=logging.DEBUG)
+async def get_project_wallet_total_credits(
+    rabbitmq_rpc_client: RabbitMQRPCClient,
+    *,
+    product_name: ProductName,
+    wallet_id: WalletID,
+    project_id: ProjectID,
+    transaction_status: CreditTransactionStatus | None = None,
+) -> WalletTotalCredits:
+    result = await rabbitmq_rpc_client.request(
+        RESOURCE_USAGE_TRACKER_RPC_NAMESPACE,
+        _RPC_METHOD_NAME_ADAPTER.validate_python("get_project_wallet_total_credits"),
+        product_name=product_name,
+        wallet_id=wallet_id,
         project_id=project_id,
+        transaction_status=transaction_status,
         timeout_s=_DEFAULT_TIMEOUT_S,
     )
     assert isinstance(result, WalletTotalCredits)  # nosec

services/resource-usage-tracker/src/simcore_service_resource_usage_tracker/api/rest/_resource_tracker.py

@@ -34,7 +34,7 @@
 async def get_credit_transactions_sum(
     wallet_total_credits: Annotated[
         WalletTotalCredits,
-        Depends(credit_transactions.sum_credit_transactions_by_product_and_wallet),
+        Depends(credit_transactions.sum_wallet_credits),
     ],
 ):
     return wallet_total_credits

services/resource-usage-tracker/src/simcore_service_resource_usage_tracker/api/rpc/_credit_transactions.py

@@ -9,7 +9,7 @@
 from models_library.wallets import WalletID
 from servicelib.rabbitmq import RPCRouter
 
-from ...services import credit_transactions
+from ...services import credit_transactions, service_runs
 
 router = RPCRouter()
 
@@ -20,20 +20,33 @@ async def get_wallet_total_credits(
     *,
     product_name: ProductName,
     wallet_id: WalletID,
-    # internal filters
+) -> WalletTotalCredits:
+    return await credit_transactions.sum_wallet_credits(
+        db_engine=app.state.engine,
+        product_name=product_name,
+        wallet_id=wallet_id,
+    )
+
+
+@router.expose(reraise_if_error_type=())
+async def get_project_wallet_total_credits(
+    app: FastAPI,
+    *,
+    product_name: ProductName,
+    wallet_id: WalletID,
+    project_id: ProjectID,
     transaction_status: CreditTransactionStatus | None = None,
-    project_id: ProjectID | None = None,
 ) -> WalletTotalCredits:
-    return await credit_transactions.sum_credit_transactions_by_product_and_wallet(
+    return await service_runs.sum_project_wallet_total_credits(
         db_engine=app.state.engine,
         product_name=product_name,
         wallet_id=wallet_id,
-        transaction_status=transaction_status,
         project_id=project_id,
+        transaction_status=transaction_status,
     )
 
 
-@router.expose(reraise_if_error_type=())
+@router.expose(reraise_if_error_type=(ValueError,))
 async def pay_project_debt(
     app: FastAPI,
     *,

services/resource-usage-tracker/src/simcore_service_resource_usage_tracker/services/credit_transactions.py

@@ -20,9 +20,10 @@
 
 from ..api.rest.dependencies import get_resource_tracker_db_engine
 from ..models.credit_transactions import CreditTransactionCreate
+from ..services.modules.db import service_runs_db
 from .modules.db import credit_transactions_db
 from .modules.rabbitmq import get_rabbitmq_client_from_request
-from .utils import make_negative, sum_credit_transactions_and_publish_to_rabbitmq
+from .utils import sum_credit_transactions_and_publish_to_rabbitmq
 
 
 async def create_credit_transaction(
@@ -74,21 +75,16 @@ async def create_credit_transaction(
         return transaction_id
 
 
-async def sum_credit_transactions_by_product_and_wallet(
+async def sum_wallet_credits(
     db_engine: Annotated[AsyncEngine, Depends(get_resource_tracker_db_engine)],
     *,
     product_name: ProductName,
     wallet_id: WalletID,
-    # attribute filters
-    transaction_status: CreditTransactionStatus | None = None,
-    project_id: ProjectID | None = None,
 ) -> WalletTotalCredits:
-    return await credit_transactions_db.sum_credit_transactions_by_product_and_wallet(
+    return await credit_transactions_db.sum_wallet_credits(
         db_engine,
         product_name=product_name,
         wallet_id=wallet_id,
-        transaction_status=transaction_status,
-        project_id=project_id,
     )
 
 
@@ -102,27 +98,42 @@ async def pay_project_debt(
 ):
     # NOTE: `current_wallet_transaction` is the Wallet in DEBT
 
-    total_project_debt_amount = (
-        await credit_transactions_db.sum_credit_transactions_by_product_and_wallet(
-            db_engine,
-            product_name=current_wallet_transaction.product_name,
-            wallet_id=current_wallet_transaction.wallet_id,
-            transaction_status=CreditTransactionStatus.IN_DEBT,
-            project_id=project_id,
-        )
+    total_project_debt_amount = await service_runs_db.sum_project_wallet_total_credits(
+        db_engine,
+        product_name=current_wallet_transaction.product_name,
+        wallet_id=current_wallet_transaction.wallet_id,
+        project_id=project_id,
+        transaction_status=CreditTransactionStatus.IN_DEBT,
     )
 
     if (
         total_project_debt_amount.available_osparc_credits
         != new_wallet_transaction.osparc_credits
     ):
-        msg = f"Project DEBT of {total_project_debt_amount.available_osparc_credits} does not equal to payment: new_wallet {new_wallet_transaction.osparc_credits}, current wallet {current_wallet_transaction.osparc_credits}"
+        msg = f"Project DEBT of {total_project_debt_amount.available_osparc_credits} does not equal to payment: new_wallet {new_wallet_transaction.wallet_id} credits {new_wallet_transaction.osparc_credits}, current wallet {current_wallet_transaction.wallet_id} credits {current_wallet_transaction.osparc_credits}"
         raise ValueError(msg)
     if (
-        make_negative(total_project_debt_amount.available_osparc_credits)
+        -total_project_debt_amount.available_osparc_credits
         != current_wallet_transaction.osparc_credits
     ):
-        msg = f"Project DEBT of {total_project_debt_amount.available_osparc_credits} does not equal to payment: new_wallet {new_wallet_transaction.osparc_credits}, current wallet {current_wallet_transaction.osparc_credits}"
+        msg = f"Project DEBT of {total_project_debt_amount.available_osparc_credits} does not equal to payment: new_wallet {new_wallet_transaction.wallet_id} credits {new_wallet_transaction.osparc_credits}, current wallet {current_wallet_transaction.wallet_id} credits {current_wallet_transaction.osparc_credits}"
+        raise ValueError(msg)
+    if current_wallet_transaction.product_name != new_wallet_transaction.product_name:
+        msg = f"Currently we do not support credit exchange between different products. New wallet {new_wallet_transaction.wallet_id}, current wallet {current_wallet_transaction.wallet_id}"
+        raise ValueError(msg)
+
+    # Does the new wallet has enough credits to pay the debt?
+    new_wallet_total_credit_amount = await credit_transactions_db.sum_wallet_credits(
+        db_engine,
+        product_name=new_wallet_transaction.product_name,
+        wallet_id=new_wallet_transaction.wallet_id,
+    )
+    if (
+        new_wallet_total_credit_amount.available_osparc_credits
+        + total_project_debt_amount.available_osparc_credits
+        < 0
+    ):
+        msg = f"New wallet {new_wallet_transaction.wallet_id} doesn't have enough credits {new_wallet_total_credit_amount.available_osparc_credits} to pay the debt {total_project_debt_amount.available_osparc_credits} of current wallet {current_wallet_transaction.wallet_id}"
         raise ValueError(msg)
 
     new_wallet_transaction_create = CreditTransactionCreate(

services/resource-usage-tracker/src/simcore_service_resource_usage_tracker/services/modules/db/credit_transactions_db.py

@@ -13,6 +13,9 @@
 from simcore_postgres_database.models.resource_tracker_credit_transactions import (
     resource_tracker_credit_transactions,
 )
+from simcore_postgres_database.models.resource_tracker_service_runs import (
+    resource_tracker_service_runs,
+)
 from simcore_postgres_database.utils_repos import transaction_context
 from sqlalchemy.ext.asyncio import AsyncConnection, AsyncEngine
 
@@ -158,22 +161,19 @@ async def batch_update_credit_transaction_status_for_in_debt_transactions(
 
     if project_id:
         update_stmt = update_stmt.where(
-            resource_tracker_credit_transactions.c.project_id == f"{project_id}"
+            resource_tracker_service_runs.c.project_id == f"{project_id}"
         )
     async with transaction_context(engine, connection) as conn:
         result = await conn.execute(update_stmt)
         print(result)
 
 
-async def sum_credit_transactions_by_product_and_wallet(
+async def sum_wallet_credits(
     engine: AsyncEngine,
     connection: AsyncConnection | None = None,
     *,
     product_name: ProductName,
     wallet_id: WalletID,
-    # attribute filters
-    transaction_status: CreditTransactionStatus | None = None,
-    project_id: ProjectID | None = None,
 ) -> WalletTotalCredits:
     async with transaction_context(engine, connection) as conn:
         sum_stmt = sa.select(
@@ -192,16 +192,6 @@ async def sum_credit_transactions_by_product_and_wallet(
             )
         )
 
-        if project_id:
-            sum_stmt = sum_stmt.where(
-                resource_tracker_credit_transactions.c.project_id == f"{project_id}"
-            )
-        if transaction_status:
-            sum_stmt = sum_stmt.where(
-                resource_tracker_credit_transactions.c.transaction_status
-                == transaction_status
-            )
-
         result = await conn.execute(sum_stmt)
     row = result.first()
     if row is None or row[0] is None:

services/resource-usage-tracker/src/simcore_service_resource_usage_tracker/services/modules/db/service_runs_db.py

@@ -1,9 +1,14 @@
-# pylint: disable=too-many-arguments
 import logging
 from datetime import datetime
+
+# pylint: disable=too-many-arguments
+from decimal import Decimal
 from typing import cast
 
 import sqlalchemy as sa
+from models_library.api_schemas_resource_usage_tracker.credit_transactions import (
+    WalletTotalCredits,
+)
 from models_library.api_schemas_storage import S3BucketName
 from models_library.products import ProductName
 from models_library.projects import ProjectID
@@ -434,6 +439,70 @@ async def get_osparc_credits_aggregated_by_service(
     )
 
 
+async def sum_project_wallet_total_credits(
+    engine: AsyncEngine,
+    connection: AsyncConnection | None = None,
+    *,
+    product_name: ProductName,
+    wallet_id: WalletID,
+    project_id: ProjectID,
+    transaction_status: CreditTransactionStatus | None = None,
+) -> WalletTotalCredits:
+    async with pass_or_acquire_connection(engine, connection) as conn:
+        sum_stmt = (
+            sa.select(
+                sa.func.SUM(resource_tracker_credit_transactions.c.osparc_credits),
+            )
+            .select_from(
+                resource_tracker_service_runs.join(
+                    resource_tracker_credit_transactions,
+                    (
+                        resource_tracker_service_runs.c.product_name
+                        == resource_tracker_credit_transactions.c.product_name
+                    )
+                    & (
+                        resource_tracker_service_runs.c.service_run_id
+                        == resource_tracker_credit_transactions.c.service_run_id
+                    ),
+                    isouter=True,
+                )
+            )
+            .where(
+                (resource_tracker_service_runs.c.product_name == product_name)
+                & (resource_tracker_service_runs.c.project_id == f"{project_id}")
+                & (
+                    resource_tracker_credit_transactions.c.transaction_classification
+                    == CreditClassification.DEDUCT_SERVICE_RUN
+                )
+                & (resource_tracker_credit_transactions.c.wallet_id == wallet_id)
+            )
+        )
+
+        if transaction_status:
+            sum_stmt = sum_stmt.where(
+                resource_tracker_credit_transactions.c.transaction_status
+                == transaction_status
+            )
+        else:
+            sum_stmt = sum_stmt.where(
+                resource_tracker_credit_transactions.c.transaction_status.in_(
+                    [
+                        CreditTransactionStatus.BILLED,
+                        CreditTransactionStatus.PENDING,
+                        CreditTransactionStatus.IN_DEBT,
+                    ]
+                )
+            )
+
+        result = await conn.execute(sum_stmt)
+        row = result.first()
+        if row is None or row[0] is None:
+            return WalletTotalCredits(
+                wallet_id=wallet_id, available_osparc_credits=Decimal(0)
+            )
+        return WalletTotalCredits(wallet_id=wallet_id, available_osparc_credits=row[0])
+
+
 async def export_service_runs_table_to_s3(
     engine: AsyncEngine,
     connection: AsyncConnection | None = None,

services/resource-usage-tracker/src/simcore_service_resource_usage_tracker/services/process_message_running_service.py

@@ -293,12 +293,10 @@ async def _process_stop_event(
             running_service.pricing_unit_cost,
         )
 
-        wallet_total_credits = (
-            await credit_transactions_db.sum_credit_transactions_by_product_and_wallet(
-                db_engine,
-                product_name=running_service.product_name,
-                wallet_id=running_service.wallet_id,
-            )
+        wallet_total_credits = await credit_transactions_db.sum_wallet_credits(
+            db_engine,
+            product_name=running_service.product_name,
+            wallet_id=running_service.wallet_id,
         )
         _transaction_status = (
             CreditTransactionStatus.BILLED

services/resource-usage-tracker/src/simcore_service_resource_usage_tracker/services/service_runs.py

@@ -1,7 +1,11 @@
+# pylint: disable=too-many-arguments
 from datetime import UTC, datetime, timedelta
 
 import shortuuid
 from aws_library.s3 import SimcoreS3API
+from models_library.api_schemas_resource_usage_tracker.credit_transactions import (
+    WalletTotalCredits,
+)
 from models_library.api_schemas_resource_usage_tracker.service_runs import (
     OsparcCreditsAggregatedByServiceGet,
     OsparcCreditsAggregatedUsagesPage,
@@ -182,6 +186,23 @@ async def export_service_runs(
     )
 
 
+async def sum_project_wallet_total_credits(
+    db_engine: AsyncEngine,
+    *,
+    product_name: ProductName,
+    wallet_id: WalletID,
+    project_id: ProjectID,
+    transaction_status: CreditTransactionStatus | None = None,
+) -> WalletTotalCredits:
+    return await service_runs_db.sum_project_wallet_total_credits(
+        db_engine,
+        product_name=product_name,
+        wallet_id=wallet_id,
+        project_id=project_id,
+        transaction_status=transaction_status,
+    )
+
+
 async def get_osparc_credits_aggregated_usages_page(
     user_id: UserID,
     product_name: ProductName,