diff options
| -rw-r--r-- | generalresearch/managers/thl/product.py | 83 | ||||
| -rw-r--r-- | generalresearch/models/thl/finance.py | 51 | ||||
| -rw-r--r-- | generalresearch/models/thl/payout.py | 7 | ||||
| -rw-r--r-- | generalresearch/models/thl/product.py | 99 | ||||
| -rw-r--r-- | generalresearch/thl_django/userprofile/models.py | 16 | ||||
| -rw-r--r-- | tests/models/thl/test_product.py | 11 |
6 files changed, 159 insertions, 108 deletions
diff --git a/generalresearch/managers/thl/product.py b/generalresearch/managers/thl/product.py index c10be38..a00719f 100644 --- a/generalresearch/managers/thl/product.py +++ b/generalresearch/managers/thl/product.py @@ -10,7 +10,7 @@ from threading import Lock from typing import TYPE_CHECKING, Any from cachetools import TTLCache, cachedmethod, keys -from psycopg import Connection +from psycopg import Connection, sql from pydantic import NonNegativeInt from sentry_sdk import capture_exception @@ -40,6 +40,22 @@ logger = logging.getLogger() class ProductManager(PostgresManager): + CACHED_FIELDS = { + "balance", + "user_wallet_balance", + "payouts", + "pop_financial", + "users_active_7d", + "task_completes_7d", + "balance_net_7d", + } + CACHED_FIELDS_JSON = { + "balance", + "user_wallet_balance", + "payouts", + "pop_financial", + } + def __init__( self, pg_config: PostgresConfig, @@ -260,9 +276,8 @@ class ProductManager(PostgresManager): bp.session_config, bp.payout_config, bp.user_create_config, - bp.cache_balance, - bp.cache_user_wallet_balance, - bp.cache_updated_at, + bp.balance, + bp.user_wallet_balance, bp.users_active_7d, bp.task_completes_7d, bp.balance_net_7d @@ -275,9 +290,6 @@ class ProductManager(PostgresManager): bp.*, bp.commission AS commission_pct, bp.grs_domain AS harmonizer_domain, - bp.cache_balance AS balance, - bp.cache_balance AS private_balance, - bp.cache_user_wallet_balance AS user_wallet_balance, COALESCE(t.tags, ARRAY[]::varchar[]) AS tags, sources.value -> 'sources_config' AS sources_config, wallet.value -> 'user_wallet' AS user_wallet_config @@ -636,3 +648,60 @@ class ProductManager(PostgresManager): conn.commit() self.cache_clear(product_uuid) + + def update_cached_fields(self, products: list[Product]) -> None: + """ + Use this to update any of the cache_* keys on a product + or any of the calculated fields such as users_active_7d. + All supported fields are updated at once + """ + updates_by_field = {} + for product in products: + data = product.model_dump(mode="json", include=self.CACHED_FIELDS) + for k, v in data.items(): + if k in self.CACHED_FIELDS_JSON and v is not None: + v = json.dumps(v) + updates_by_field.setdefault(k, {})[product.id] = v + for field, data in updates_by_field.items(): + self.update_field_bulk(field=field, data=data) + + def update_field_bulk(self, field: str, data: dict[str, Any]) -> None: + """ + Update one field for multiple products. + + data maps product UUID -> new value. + """ + if not data: + return + + assert field in self.CACHED_FIELDS, f"Unsupported field: {field!r}" + + value_rows = sql.SQL(", ").join(sql.SQL("(%s, %s)") for _ in data) + value_expression = ( + sql.SQL("updates.value::jsonb") + if field in self.CACHED_FIELDS_JSON + else sql.SQL("updates.value::integer") + ) + + query = sql.SQL(""" + UPDATE userprofile_brokerageproduct AS bp + SET {field} = {value_expression} + FROM (VALUES {value_rows}) AS updates(product_id, value) + WHERE bp.id = updates.product_id::uuid + """).format( + field=sql.Identifier(field), + value_expression=value_expression, + value_rows=value_rows, + ) + + params = [ + item for product_id, value in data.items() for item in (product_id, value) + ] + + with self.pg_config.make_connection() as conn: + with conn.cursor() as cursor: + cursor.execute(query, params) + conn.commit() + + for product_id in data: + self.cache_clear(product_id) diff --git a/generalresearch/models/thl/finance.py b/generalresearch/models/thl/finance.py index 0eb4f8f..f5e8745 100644 --- a/generalresearch/models/thl/finance.py +++ b/generalresearch/models/thl/finance.py @@ -2,7 +2,7 @@ from __future__ import annotations import random from collections.abc import Iterable -from datetime import UTC +from datetime import UTC, datetime from typing import TYPE_CHECKING from uuid import uuid4 @@ -239,6 +239,13 @@ class POPFinancial(BaseModel): return res +class ProductPOPFinancials(BaseModel): + periods: list[POPFinancial] = Field(default_factory=list) + updated_at: AwareDatetimeISO = Field( + default_factory=lambda: datetime.now(tz=UTC), + ) + + class UserWalletBalances(BaseModel): """Cumulative ledger activity and balance for one user's USD wallet. @@ -408,9 +415,7 @@ class UserWalletBalances(BaseModel): data.update( product_id=product_id, product_user_id=product_user_id, - last_event=( - None if wallet_rows.empty else wallet_rows["time_idx"].max() - ), + last_event=(None if wallet_rows.empty else wallet_rows["time_idx"].max()), ) return cls.model_validate(data) @@ -937,6 +942,15 @@ class ProductBalances(BaseModel): "balances elsewhere.", ) + commission: NonNegativeInt | None = Field( + default=None, + description=( + "Net commission revenue from this Brokerage Product. " + "Positive adjustments increase this value and reversals decrease it." + ), + examples=[5_038], + ) + # --- Validate --- @model_validator(mode="after") def check_unknown_fields(self) -> ProductBalances: @@ -1001,7 +1015,11 @@ class ProductBalances(BaseModel): ) @property def expense(self) -> int: - return self.user_bonus_credit - self.user_bonus_debit - self.user_payout_complete_debit + return ( + self.user_bonus_credit + - self.user_bonus_debit + - self.user_payout_complete_debit + ) # --- Properties: account related --- @computed_field( @@ -1226,6 +1244,11 @@ class ProductBalances(BaseModel): "issued_payment", ), ( + "commission_usd", + "Net commission revenue from the Product.", + "commission", + ), + ( "payout_usd", "Total task payouts earned by the Product.", "payout", @@ -1278,9 +1301,12 @@ class ProductBalances(BaseModel): labels=["product_id"], ) for product in products: + value = getattr(product, attribute) + if value is None: + continue metric.add_metric( [str(product.product_id)], - int(getattr(product, attribute)) / 100, + int(value) / 100, ) yield metric @@ -1350,19 +1376,6 @@ class ProductBalances(BaseModel): ).replace("$-", "-$") -class PrivateProductBalances(ProductBalances): - """Product balances with internal revenue visible to administrative APIs.""" - - commission: int = Field( - default=0, - description=( - "Net commission revenue earned by GRL from this Brokerage Product. " - "Positive adjustments increase this value and reversals decrease it." - ), - examples=[5_038], - ) - - class BusinessBalances(BaseModel): product_balances: list[ProductBalances] = Field(default_factory=list) diff --git a/generalresearch/models/thl/payout.py b/generalresearch/models/thl/payout.py index ab231bc..ecc9b0c 100644 --- a/generalresearch/models/thl/payout.py +++ b/generalresearch/models/thl/payout.py @@ -238,6 +238,13 @@ class BrokerageProductPayoutEvent(PayoutEvent): return self.amount_usd.to_usd_str() +class ProductPayouts(BaseModel): + events: list[BrokerageProductPayoutEvent] = Field(default_factory=list) + updated_at: AwareDatetimeISO = Field( + default_factory=lambda: datetime.now(tz=UTC), + ) + + class BusinessPayoutEventCreate(BaseModel): """A single payout event to a supplier Business.""" diff --git a/generalresearch/models/thl/product.py b/generalresearch/models/thl/product.py index 7188390..84230fb 100644 --- a/generalresearch/models/thl/product.py +++ b/generalresearch/models/thl/product.py @@ -48,13 +48,11 @@ from generalresearch.models.custom_types import ( from generalresearch.models.definitions import Source from generalresearch.models.thl.finance import ( POPFinancial, - PrivateProductBalances, ProductBalances, + ProductPOPFinancials, ProductUserWalletBalances, ) -from generalresearch.models.thl.payout import ( - BrokerageProductPayoutEvent, -) +from generalresearch.models.thl.payout import ProductPayouts from generalresearch.models.thl.payout_format import ( PayoutFormatType, format_payout_format, @@ -85,7 +83,6 @@ if TYPE_CHECKING: PRODUCT_BALANCES_METRICS_CACHE_KEY = "metrics:product_balances" -PRODUCT_PRIVATE_BALANCES_METRICS_CACHE_KEY = "metrics:private_product_balances" PRODUCT_USER_WALLET_BALANCES_METRICS_CACHE_KEY = "metrics:product_user_wallet_balances" @@ -969,22 +966,16 @@ class Product(BaseModel, validate_assignment=True): # Initialization is deferred until unless it's called # (see .prebuild_***()) + bp_account: LedgerAccount | None = Field(default=None) + balance: ProductBalances | None = Field(default=None, description="Product Balance") - private_balance: PrivateProductBalances | None = Field( - default=None, description="Product Balance including private keys" - ) user_wallet_balance: ProductUserWalletBalances | None = Field(default=None) - - payouts_total_str: str | None = Field(default=None) - payouts_total: USDCent | None = Field(default=None) - payouts: list[BrokerageProductPayoutEvent] | None = Field( + payouts: ProductPayouts | None = Field( default=None, description="Product Payouts. These are the ACH or Wire payments that were sent to the" "Business on behalf of this specific Product", ) - - pop_financial: list[POPFinancial] | None = Field(default=None) - bp_account: LedgerAccount | None = Field(default=None) + pop_financial: ProductPOPFinancials | None = Field(default=None) users_active_7d: int | None = Field( default=None, description="Count of active users in the past 7 days" @@ -1230,42 +1221,26 @@ class Product(BaseModel, validate_assignment=True): ).set_index("time_idx") balance = ProductBalances.from_pandas(balance_df) balance.product_id = self.uuid - self.balance = balance - - def prebuild_private_balance( - self, - thl_lm: ThlLedgerManager, - pop_ledger_df: pd.DataFrame | None = None, - ) -> None: - from generalresearch.models.thl.finance import PrivateProductBalances commission_account = thl_lm.get_account_or_create_bp_commission_by_uuid( self.uuid ) - - df = pop_ledger_df.loc[pop_ledger_df["account_id"].eq(commission_account.uuid)] - if df.empty: - LOG.warning( - f"Product({self.uuid=}).prebuild_private_balance empty dataframe" - ) - return - - s = df[ + commission_rows = pop_ledger_df.loc[ + pop_ledger_df["account_id"].eq(commission_account.uuid) + ] + commission_totals = commission_rows[ [ "bp_payment.CREDIT", "bp_adjustment.CREDIT", "bp_adjustment.DEBIT", ] ].sum() - - commission = ( - s["bp_payment.CREDIT"] - + s["bp_adjustment.CREDIT"] - - s["bp_adjustment.DEBIT"] - ) - self.private_balance = PrivateProductBalances.model_validate( - self.balance.model_dump() | {"commission": commission} + balance.commission = int( + commission_totals["bp_payment.CREDIT"] + + commission_totals["bp_adjustment.CREDIT"] + - commission_totals["bp_adjustment.DEBIT"] ) + self.balance = balance def prebuild_user_wallet_balances( self, @@ -1331,20 +1306,18 @@ class Product(BaseModel, validate_assignment=True): ] if df.empty: - self.pop_financial = [] + self.pop_financial = ProductPOPFinancials() return df = df.groupby( [pd.Grouper(key="time_idx", freq=rr.interval), "account_id"] ).sum() - from generalresearch.models.thl.finance import POPFinancial - - self.pop_financial = POPFinancial.list_from_pandas( + periods = POPFinancial.list_from_pandas( input_data=df, accounts=[self.bp_account] ) - return + self.pop_financial = ProductPOPFinancials(periods=periods) def prebuild_payouts( self, @@ -1353,18 +1326,25 @@ class Product(BaseModel, validate_assignment=True): LOG.debug(f"Product.prebuild_payouts({self.uuid=})") from generalresearch.models.thl.ledger import OrderBy - self.payouts = bp_pem.get_bp_bp_payout_events_for_products( + events = bp_pem.get_bp_bp_payout_events_for_products( product_uuids=[self.uuid], order_by=OrderBy.DESC, ) + self.payouts = ProductPayouts(events=events) - self.prebuild_payouts_total() - - def prebuild_payouts_total(self) -> None: - assert self.payouts is not None + @computed_field + @property + def payouts_total(self) -> USDCent | None: + if self.payouts is None: + return None + return USDCent(sum(event.amount for event in self.payouts.events)) - self.payouts_total = USDCent(sum([po.amount for po in self.payouts])) - self.payouts_total_str = self.payouts_total.to_usd_str() + @computed_field + @property + def payouts_total_str(self) -> str | None: + if self.payouts_total is None: + return None + return self.payouts_total.to_usd_str() def set_cache( self, @@ -1411,15 +1391,10 @@ class Product(BaseModel, validate_assignment=True): thl_lm=thl_lm, pop_ledger_df=pop_ledger_df, ) - if self.balance: - self.prebuild_private_balance( - thl_lm=thl_lm, + if self.balance and self.user_wallet_enabled: + self.prebuild_user_wallet_balances( pop_ledger_df=pop_ledger_df, ) - if self.user_wallet_enabled: - self.prebuild_user_wallet_balances( - pop_ledger_df=pop_ledger_df, - ) self.prebuild_payouts(bp_pem=bp_pem) self.prebuild_pop_financial( thl_lm=thl_lm, @@ -1439,12 +1414,6 @@ class Product(BaseModel, validate_assignment=True): key=self.uuid, value=self.balance.model_dump_json(), ) - if self.private_balance is not None: - pipe.hset( - name=PRODUCT_PRIVATE_BALANCES_METRICS_CACHE_KEY, - key=self.uuid, - value=self.private_balance.model_dump_json(), - ) if self.user_wallet_balance is not None: pipe.hset( name=PRODUCT_USER_WALLET_BALANCES_METRICS_CACHE_KEY, diff --git a/generalresearch/thl_django/userprofile/models.py b/generalresearch/thl_django/userprofile/models.py index 813dac0..445a53c 100644 --- a/generalresearch/thl_django/userprofile/models.py +++ b/generalresearch/thl_django/userprofile/models.py @@ -79,16 +79,14 @@ class BrokerageProduct(models.Model): # Store configuration regarding user creation. See: models/thl/product.py:UserCreateConfig user_create_config = models.JSONField(default=dict) - # PrivateProductBalances model (can be converted to ProductBalances) - cache_balance = models.JSONField(default=None, null=True) + # ProductBalances model + balance = models.JSONField(default=None, null=True) # ProductUserWalletBalances model - cache_user_wallet_balance = models.JSONField(default=None, null=True) - # list[BrokerageProductPayoutEvent] - cache_payouts = models.JSONField(default=None, null=True) - # list[POPFinancial] - cache_pop_financial = models.JSONField(default=None, null=True) - # We should update all the cache_* fields at once - cache_updated_at = models.DateTimeField(default=None, null=True) + user_wallet_balance = models.JSONField(default=None, null=True) + # ProductPayouts model + payouts = models.JSONField(default=None, null=True) + # ProductPOPFinancials model + pop_financial = models.JSONField(default=None, null=True) users_active_7d = models.IntegerField(default=None, null=True) task_completes_7d = models.IntegerField(default=None, null=True) diff --git a/tests/models/thl/test_product.py b/tests/models/thl/test_product.py index ce3242e..f8f3da3 100644 --- a/tests/models/thl/test_product.py +++ b/tests/models/thl/test_product.py @@ -873,8 +873,7 @@ class TestProductFinancials: body, content_type = p1.balance.to_prometheus() - p1.prebuild_private_balance(thl_lm=thl_ledger_manager, pop_ledger_df=df) - assert p1.private_balance.commission == 5 * 2 + assert p1.balance.commission == 5 * 2 p1.prebuild_user_wallet_balances(pop_ledger_df=df) assert p1.user_wallet_balance.outstanding_liability == 38 * 2 @@ -952,15 +951,11 @@ class TestProductFinancials: p1.prebuild_balance(thl_lm=thl_ledger_manager, pop_ledger_df=df) assert p1.balance.payout == 95 - - p1.prebuild_private_balance(thl_lm=thl_ledger_manager, pop_ledger_df=df) - assert p1.private_balance.commission == 5 + assert p1.balance.commission == 5 p2.prebuild_balance(thl_lm=thl_ledger_manager, pop_ledger_df=df) assert p2.balance.payout == 57 - - p2.prebuild_private_balance(thl_lm=thl_ledger_manager, pop_ledger_df=df) - assert p2.private_balance.commission == 5 + assert p2.balance.commission == 5 p2.prebuild_user_wallet_balances(pop_ledger_df=df) assert p2.user_wallet_balance.outstanding_liability == 38 |