From 63c692c5a26c6613bd3ffbd5feb46f3497f2353d Mon Sep 17 00:00:00 2001
From: Rafael Cardenas <rafael@rafaelcr.com>
Date: Mon, 27 Jan 2025 13:35:28 -0600
Subject: [PATCH] feat: add estimated inbound and outbound stx balance

---
 src/api/routes/address.ts            | 20 ++++++++++++++++----
 src/api/schemas/entities/balances.ts | 10 ++++++++++
 src/datastore/pg-store.ts            | 20 ++++++++++++--------
 3 files changed, 38 insertions(+), 12 deletions(-)

diff --git a/src/api/routes/address.ts b/src/api/routes/address.ts
index 7cf31bedb..2737754b0 100644
--- a/src/api/routes/address.ts
+++ b/src/api/routes/address.ts
@@ -121,13 +121,19 @@ export const AddressRoutes: FastifyPluginAsync<
           blockHeight
         );
         let mempoolBalance: bigint | undefined = undefined;
+        let mempoolInbound: bigint | undefined = undefined;
+        let mempoolOutbound: bigint | undefined = undefined;
         if (req.query.until_block === undefined) {
-          const delta = await fastify.db.getPrincipalMempoolStxBalanceDelta(sql, stxAddress);
-          mempoolBalance = stxBalanceResult.balance + delta;
+          const pending = await fastify.db.getPrincipalMempoolStxBalanceDelta(sql, stxAddress);
+          mempoolInbound = pending.inbound;
+          mempoolOutbound = pending.outbound;
+          mempoolBalance = stxBalanceResult.balance + pending.delta;
         }
         const result: AddressStxBalance = {
           balance: stxBalanceResult.balance.toString(),
           estimated_balance: mempoolBalance?.toString(),
+          estimated_balance_inbound: mempoolInbound?.toString(),
+          estimated_balance_outbound: mempoolOutbound?.toString(),
           total_sent: stxBalanceResult.totalSent.toString(),
           total_received: stxBalanceResult.totalReceived.toString(),
           total_fees_sent: stxBalanceResult.totalFeesSent.toString(),
@@ -211,15 +217,21 @@ export const AddressRoutes: FastifyPluginAsync<
         });
 
         let mempoolBalance: bigint | undefined = undefined;
+        let mempoolInbound: bigint | undefined = undefined;
+        let mempoolOutbound: bigint | undefined = undefined;
         if (req.query.until_block === undefined) {
-          const delta = await fastify.db.getPrincipalMempoolStxBalanceDelta(sql, stxAddress);
-          mempoolBalance = stxBalanceResult.balance + delta;
+          const pending = await fastify.db.getPrincipalMempoolStxBalanceDelta(sql, stxAddress);
+          mempoolInbound = pending.inbound;
+          mempoolOutbound = pending.outbound;
+          mempoolBalance = stxBalanceResult.balance + pending.delta;
         }
 
         const result: AddressBalance = {
           stx: {
             balance: stxBalanceResult.balance.toString(),
             estimated_balance: mempoolBalance?.toString(),
+            estimated_balance_inbound: mempoolInbound?.toString(),
+            estimated_balance_outbound: mempoolOutbound?.toString(),
             total_sent: stxBalanceResult.totalSent.toString(),
             total_received: stxBalanceResult.totalReceived.toString(),
             total_fees_sent: stxBalanceResult.totalFeesSent.toString(),
diff --git a/src/api/schemas/entities/balances.ts b/src/api/schemas/entities/balances.ts
index 4adb466fc..cec755499 100644
--- a/src/api/schemas/entities/balances.ts
+++ b/src/api/schemas/entities/balances.ts
@@ -26,6 +26,16 @@ export const StxBalanceSchema = Type.Object(
         description: 'Total STX balance considering pending mempool transactions',
       })
     ),
+    estimated_balance_inbound: Type.Optional(
+      Type.String({
+        description: 'Inbound STX balance from pending mempool transactions',
+      })
+    ),
+    estimated_balance_outbound: Type.Optional(
+      Type.String({
+        description: 'Outbound STX balance from pending mempool transactions',
+      })
+    ),
     total_sent: Type.String(),
     total_received: Type.String(),
     total_fees_sent: Type.String(),
diff --git a/src/datastore/pg-store.ts b/src/datastore/pg-store.ts
index 8213150f2..256642601 100644
--- a/src/datastore/pg-store.ts
+++ b/src/datastore/pg-store.ts
@@ -2480,8 +2480,8 @@ export class PgStore extends BasePgStore {
    * Returns the total STX balance delta affecting a principal from transactions currently in the
    * mempool.
    */
-  async getPrincipalMempoolStxBalanceDelta(sql: PgSqlClient, principal: string): Promise<bigint> {
-    const results = await sql<{ delta: string }[]>`
+  async getPrincipalMempoolStxBalanceDelta(sql: PgSqlClient, principal: string) {
+    const results = await sql<{ inbound: string; outbound: string; delta: string }[]>`
       WITH sent AS (
         SELECT SUM(COALESCE(token_transfer_amount, 0) + fee_rate) AS total
         FROM mempool_txs
@@ -2496,14 +2496,18 @@ export class PgStore extends BasePgStore {
         SELECT SUM(COALESCE(token_transfer_amount, 0)) AS total
         FROM mempool_txs
         WHERE pruned = false AND token_transfer_recipient_address = ${principal}
+      ),
+      values AS (
+        COALESCE((SELECT total FROM received), 0) AS inbound,
+        COALESCE((SELECT total FROM sent), 0) + COALESCE((SELECT total FROM sponsored), 0) AS outbound
       )
-      SELECT
-        COALESCE((SELECT total FROM received), 0)
-        - COALESCE((SELECT total FROM sent), 0)
-        - COALESCE((SELECT total FROM sponsored), 0)
-        AS delta
+      SELECT inbound, outbound, (inbound - outbound) AS delta
     `;
-    return BigInt(results[0]?.delta ?? '0');
+    return {
+      inbound: BigInt(results[0]?.inbound ?? '0'),
+      outbound: BigInt(results[0]?.outbound ?? '0'),
+      delta: BigInt(results[0]?.delta ?? '0'),
+    };
   }
 
   async getUnlockedStxSupply(