business prebuild_metrics

2 files changed, 83 insertions, 7 deletions

diff --git a/generalresearch/managers/gr/business.py b/generalresearch/managers/gr/business.py
index 9338f03..16e81ba 100644
--- a/generalresearch/managers/gr/business.py
+++ b/generalresearch/managers/gr/business.py
@@ -1,5 +1,6 @@
 from __future__ import annotations
 
+import json
 from typing import TYPE_CHECKING
 from uuid import UUID, uuid4
 
@@ -155,6 +156,21 @@ class BusinessManager(PostgresManagerWithRedis):
         referenced by the gr-api services
 
     """
+    CACHED_FIELDS: frozenset[str] = frozenset(
+        {
+            "balance",
+            "user_wallet_balance",
+            "users_active_7d",
+            "task_completes_7d",
+            "balance_net_7d",
+        }
+    )
+    CACHED_FIELDS_JSON: frozenset[str] = frozenset(
+        {
+            "balance",
+            "user_wallet_balance",
+        }
+    )
 
     def get_or_create(
         self,
@@ -408,3 +424,30 @@ class BusinessManager(PostgresManagerWithRedis):
             return None
 
         return Business.model_validate(res[0])
+
+    def update_cached_fields(self, business: Business) -> None:
+        """
+        Use this to update any of the cached keys on a business
+        or any of the calculated fields such as users_active_7d.
+        All supported fields are updated at once
+        """
+        data = business.model_dump(mode="json", include=set(self.CACHED_FIELDS))
+        for k, v in data.items():
+            if k in self.CACHED_FIELDS_JSON and v is not None:
+                data[k] = json.dumps(v)
+
+        assert business.id is not None, "Business must be persisted before updating"
+        assignments = sql.SQL(", ").join(
+            sql.SQL("{} = %s").format(sql.Identifier(field)) for field in data
+        )
+        query = sql.SQL("""
+            UPDATE common_business
+            SET {assignments}
+            WHERE id = %s
+        """).format(assignments=assignments)
+        params = [*data.values(), business.id]
+
+        with self.pg_config.make_connection() as conn:
+            with conn.cursor() as cursor:
+                cursor.execute(query, params)
+            conn.commit()
diff --git a/generalresearch/models/gr/business.py b/generalresearch/models/gr/business.py
index 1c8e40c..24894cf 100644
--- a/generalresearch/models/gr/business.py
+++ b/generalresearch/models/gr/business.py
@@ -33,7 +33,6 @@ from generalresearch.models.gr.team import Team
 from generalresearch.models.thl.finance import BusinessBalances, POPFinancial
 from generalresearch.models.thl.ledger import LedgerAccount, OrderBy
 from generalresearch.models.thl.payout import BusinessPayoutEvent
-from generalresearch.utils.aggregation import group_by_year
 
 if TYPE_CHECKING:
     from generalresearch.incite.mergers.pop_ledger import PopLedgerMerge
@@ -219,6 +218,17 @@ class Business(BaseModel):
     pop_financial: list[POPFinancial] | None = Field(default=None)
     bp_accounts: list[LedgerAccount] | None = Field(default=None)
 
+    users_active_7d: int | None = Field(
+        default=None, description="Count of active users in the past 7 days"
+    )
+    task_completes_7d: int | None = Field(
+        default=None, description="Count of completes in the past 7 days"
+    )
+    balance_net_7d: int | None = Field(
+        default=None,
+        description="Net Earnings over the last 7 days (in USD Cents, this can be positive or negative)",
+    )
+
     def __str__(self) -> str:
         return (
             f"Name: {self.name} ({self.uuid})\n"
@@ -238,9 +248,10 @@ class Business(BaseModel):
         # --- Prefetch ---
 
     def prefetch_addresses(self, pg_config: PostgresConfig) -> None:
-        with pg_config.make_connection() as conn, conn.cursor(
-            row_factory=dict_row
-        ) as c:
+        with (
+            pg_config.make_connection() as conn,
+            conn.cursor(row_factory=dict_row) as c,
+        ):
             c.execute(
                 query="""
                         SELECT *
@@ -258,9 +269,10 @@ class Business(BaseModel):
         self.addresses = [BusinessAddress.model_validate(i) for i in res]
 
     def prefetch_teams(self, pg_config: PostgresConfig) -> None:
-        with pg_config.make_connection() as conn, conn.cursor(
-            row_factory=dict_row
-        ) as c:
+        with (
+            pg_config.make_connection() as conn,
+            conn.cursor(row_factory=dict_row) as c,
+        ):
             c: Cursor
 
             c.execute(
@@ -434,6 +446,26 @@ class Business(BaseModel):
 
         return
 
+    def prebuild_metrics(self, product_manager: ProductManager):
+        """
+        Simply sum the products' metrics
+        """
+        product_ids = self.product_uuids
+        query = """
+        SELECT
+            SUM(users_active_7d) AS users_active_7d,
+            SUM(task_completes_7d) AS task_completes_7d,
+            SUM(balance_net_7d) AS balance_net_7d
+        FROM userprofile_brokerageproduct
+        WHERE id = ANY(%(product_ids)s)
+        """
+        with product_manager.pg_config.make_connection() as conn, conn.cursor() as c:
+            c.execute(query, {"product_ids": product_ids})
+            res = c.fetchone()
+        self.users_active_7d = res["users_active_7d"]
+        self.task_completes_7d = res["task_completes_7d"]
+        self.balance_net_7d = res["balance_net_7d"]
+
     def prebuild_payouts(
         self,
         bpem: BusinessPayoutEventManager,
@@ -658,6 +690,7 @@ class Business(BaseModel):
             pop_ledger=pop_ledger,
         )
         self.prebuild_payouts(bpem=bpem)
+        self.prebuild_metrics(product_manager=product_manager)
         # self.prebuild_pop_financial(
         #     product_manager=product_manager,
         #     thl_lm=thl_lm,