Skip to content

Commit

Permalink
improbemtns
Browse files Browse the repository at this point in the history
  • Loading branch information
matusdrobuliak66 committed Jan 14, 2025
1 parent e18149e commit 53a31cb
Show file tree
Hide file tree
Showing 9 changed files with 102 additions and 30 deletions.
20 changes: 18 additions & 2 deletions packages/models-library/src/models_library/resource_tracker.py
Original file line number Diff line number Diff line change
Expand Up @@ -61,11 +61,27 @@ class CreditTransactionStatus(StrAutoEnum):


class CreditClassification(StrAutoEnum):
ADD_WALLET_TOP_UP = auto() # user top up credits
DEDUCT_SERVICE_RUN = auto() # computational/dynamic service run costs)
# Represents the different types of credit classifications.

ADD_WALLET_TOP_UP = auto()
# Indicates that credits have been added to the user's wallet through a top-up.
# Example: The user adds funds to their wallet to increase their available credits.

DEDUCT_SERVICE_RUN = auto()
# Represents a deduction from the user's wallet due to the costs of running a computational or dynamic service.
# Example: Credits are deducted when the user runs a simulation.

DEDUCT_LICENSE_PURCHASE = auto()
# Represents a deduction from the user's wallet for purchasing a license.
# Example: The user purchases a license to access premium features such as VIP models.

ADD_WALLET_EXCHANGE = auto()
# Represents the addition of credits to the user's wallet through an exchange.
# Example: Credits are added due to credit exchange between wallets.

DEDUCT_WALLET_EXCHANGE = auto()
# Represents a deduction of credits from the user's wallet through an exchange.
# Example: Credits are deducted due to credit exchange between wallets.


class PricingPlanClassification(StrAutoEnum):
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -18,6 +18,7 @@
class CreditTransactionStatus(str, enum.Enum):
PENDING = "PENDING"
BILLED = "BILLED"
IN_DEBT = "IN_DEBT"
NOT_BILLED = "NOT_BILLED"
REQUIRES_MANUAL_REVIEW = "REQUIRES_MANUAL_REVIEW"

Expand All @@ -28,6 +29,8 @@ class CreditTransactionClassification(str, enum.Enum):
"DEDUCT_SERVICE_RUN" # computational/dynamic service run costs)
)
DEDUCT_LICENSE_PURCHASE = "DEDUCT_LICENSE_PURCHASE"
ADD_WALLET_EXCHANGE = "ADD_WALLET_EXCHANGE"
DEDUCT_WALLET_EXCHANGE = "DEDUCT_WALLET_EXCHANGE"


resource_tracker_credit_transactions = sa.Table(
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -9,8 +9,10 @@
ServiceRunPage,
)
from models_library.products import ProductName
from models_library.projects import ProjectID
from models_library.rabbitmq_basic_types import RPCMethodName
from models_library.resource_tracker import (
CreditTransactionStatus,
ServiceResourceUsagesFilters,
ServicesAggregatedUsagesTimePeriod,
ServicesAggregatedUsagesType,
Expand All @@ -37,24 +39,30 @@ async def get_service_run_page(
*,
user_id: UserID,
product_name: ProductName,
limit: int = 20,
offset: int = 0,
wallet_id: WalletID | None = None,
access_all_wallet_usage: bool = False,
order_by: OrderBy | None = None,
filters: ServiceResourceUsagesFilters | None = None,
transaction_status: CreditTransactionStatus | None = None,
project_id: ProjectID | None = None,
# pagination
offset: int = 0,
limit: int = 20,
# ordering
order_by: OrderBy | None = None,
) -> ServiceRunPage:
result = await rabbitmq_rpc_client.request(
RESOURCE_USAGE_TRACKER_RPC_NAMESPACE,
_RPC_METHOD_NAME_ADAPTER.validate_python("get_service_run_page"),
user_id=user_id,
product_name=product_name,
limit=limit,
offset=offset,
wallet_id=wallet_id,
access_all_wallet_usage=access_all_wallet_usage,
order_by=order_by,
filters=filters,
transaction_status=transaction_status,
project_id=project_id,
offset=offset,
limit=limit,
order_by=order_by,
timeout_s=_DEFAULT_TIMEOUT_S,
)
assert isinstance(result, ServiceRunPage) # nosec
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -4,7 +4,9 @@
ServiceRunPage,
)
from models_library.products import ProductName
from models_library.projects import ProjectID
from models_library.resource_tracker import (
CreditTransactionStatus,
ServiceResourceUsagesFilters,
ServicesAggregatedUsagesTimePeriod,
ServicesAggregatedUsagesType,
Expand Down Expand Up @@ -34,9 +36,11 @@ async def get_service_run_page(
wallet_id: WalletID | None = None,
access_all_wallet_usage: bool = False,
filters: ServiceResourceUsagesFilters | None = None,
transaction_status: CreditTransactionStatus | None = None,
project_id: ProjectID | None = None,
# pagination
limit: int = 20,
offset: int = 0,
limit: int = 20,
# ordering
order_by: OrderBy | None = None,
) -> ServiceRunPage:
Expand All @@ -47,8 +51,10 @@ async def get_service_run_page(
wallet_id=wallet_id,
access_all_wallet_usage=access_all_wallet_usage,
filters=filters,
limit=limit,
transaction_status=transaction_status,
project_id=project_id,
offset=offset,
limit=limit,
order_by=order_by,
)

Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -116,12 +116,14 @@ async def pay_project_debt(
total_project_debt_amount.available_osparc_credits
!= new_wallet_transaction.osparc_credits
):
raise ValueError("wrong")
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}"
raise ValueError(msg)
if (
make_negative(total_project_debt_amount.available_osparc_credits)
!= current_wallet_transaction.osparc_credits
):
raise ValueError("wrong")
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}"
raise ValueError(msg)

new_wallet_transaction_create = CreditTransactionCreate(
product_name=new_wallet_transaction.product_name,
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -6,6 +6,7 @@
import sqlalchemy as sa
from models_library.api_schemas_storage import S3BucketName
from models_library.products import ProductName
from models_library.projects import ProjectID
from models_library.resource_tracker import (
CreditClassification,
CreditTransactionStatus,
Expand Down Expand Up @@ -197,6 +198,8 @@ async def list_service_runs_by_product_and_user_and_wallet(
service_run_status: ServiceRunStatus | None = None,
started_from: datetime | None = None,
started_until: datetime | None = None,
transaction_status: CreditTransactionStatus | None = None,
project_id: ProjectID | None = None,
# pagination
offset: int,
limit: int,
Expand Down Expand Up @@ -287,6 +290,15 @@ async def list_service_runs_by_product_and_user_and_wallet(
sa.func.DATE(resource_tracker_service_runs.c.started_at)
<= started_until.date()
)
if project_id:
base_query = base_query.where(
resource_tracker_service_runs.c.project_id == project_id
)
if transaction_status:
base_query = base_query.where(
resource_tracker_credit_transactions.c.transaction_status
== transaction_status
)

# Select total count from base_query
subquery = base_query.subquery()
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -10,7 +10,9 @@
)
from models_library.api_schemas_storage import S3BucketName
from models_library.products import ProductName
from models_library.projects import ProjectID
from models_library.resource_tracker import (
CreditTransactionStatus,
ServiceResourceUsagesFilters,
ServicesAggregatedUsagesTimePeriod,
ServicesAggregatedUsagesType,
Expand All @@ -34,8 +36,10 @@ async def list_service_runs(
wallet_id: WalletID | None = None,
access_all_wallet_usage: bool = False,
filters: ServiceResourceUsagesFilters | None = None,
limit: int = 20,
transaction_status: CreditTransactionStatus | None = None,
project_id: ProjectID | None = None,
offset: int = 0,
limit: int = 20,
order_by: OrderBy | None = None,
) -> ServiceRunPage:
started_from = None
Expand All @@ -56,6 +60,8 @@ async def list_service_runs(
wallet_id=None,
started_from=started_from,
started_until=started_until,
transaction_status=transaction_status,
project_id=project_id,
offset=offset,
limit=limit,
order_by=order_by,
Expand All @@ -72,6 +78,8 @@ async def list_service_runs(
wallet_id=wallet_id,
started_from=started_from,
started_until=started_until,
transaction_status=transaction_status,
project_id=project_id,
offset=offset,
limit=limit,
order_by=order_by,
Expand All @@ -88,6 +96,8 @@ async def list_service_runs(
wallet_id=wallet_id,
started_from=started_from,
started_until=started_until,
transaction_status=transaction_status,
project_id=project_id,
offset=offset,
limit=limit,
order_by=order_by,
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -12,6 +12,7 @@
from models_library.wallets import WalletDB, WalletID
from servicelib.rabbitmq.rpc_interfaces.resource_usage_tracker import (
credit_transactions,
service_runs,
)

from ..rabbitmq import get_rabbitmq_rpc_client
Expand Down Expand Up @@ -40,32 +41,47 @@ async def connect_wallet_to_project(
) -> WalletGet:
db: ProjectDBAPI = ProjectDBAPI.get_from_app_context(app)

project_wallet = await db.get_project_wallet(project_uuid=project_id)
# ensure the wallet can be used by the user
wallet: WalletGet = await wallets_api.get_wallet_by_user(
app,
user_id=user_id,
wallet_id=wallet_id,
product_name=product_name,
)

current_project_wallet = await db.get_project_wallet(project_uuid=project_id)
rpc_client = get_rabbitmq_rpc_client(app)

if project_wallet:
# NOTE: Do not allow to change wallet if the project is in DEBT!
rpc_client = get_rabbitmq_rpc_client(app)
if current_project_wallet:
# Do not allow to change wallet if the project connected wallet is in DEBT!
project_wallet_credits = await credit_transactions.get_wallet_total_credits(
rpc_client,
product_name=product_name,
wallet_id=project_wallet.wallet_id,
wallet_id=current_project_wallet.wallet_id,
project_id=project_id,
transaction_status=CreditTransactionStatus.IN_DEBT,
)
if project_wallet_credits.available_osparc_credits > 0:
msg = f"Current Project Wallet {project_wallet.wallet_id} is in DEBT"
msg = (
f"Current Project Wallet {current_project_wallet.wallet_id} is in DEBT"
)
raise ValueError(msg)

# ensure the wallet can be used by the user
wallet: WalletGet = await wallets_api.get_wallet_by_user(
app,
user_id=user_id,
wallet_id=wallet_id,
product_name=product_name,
)

# Allow changing the wallet only if there are no pending transactions within the project.
# TODO: MATUS: check pending transactions
# Do not allow to change wallet if the project has transaction in PENDING!
project_service_runs_in_progress = await service_runs.get_service_run_page(
rpc_client,
user_id=user_id,
product_name=product_name,
wallet_id=wallet_id,
access_all_wallet_usage=True,
transaction_status=CreditTransactionStatus.PENDING,
project_id=project_id,
offset=0,
limit=1,
)
if project_service_runs_in_progress.total > 0:
msg = "Can not change the wallet, as project has currently pending transaction"
raise ValueError(msg)

await db.connect_wallet_to_project(project_uuid=project_id, wallet_id=wallet_id)
return wallet
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -140,7 +140,6 @@ async def pay_project_debt(request: web.Request):
request.app, path_params.project_id
)
if not current_wallet:
_logger.warning("This should not happen?")
raise web.HTTPNotFound(
reason="Project doesn't have any wallet associated to the project"
)
Expand Down

0 comments on commit 53a31cb

Please sign in to comment.