From 6764d526340f41a5d9369d2d2732a47ddd4df8f8 Mon Sep 17 00:00:00 2001
From: Lex Berezhny <lex@damoti.com>
Date: Wed, 21 Oct 2020 16:48:40 -0400
Subject: [PATCH] postgres optimized balance query

---
 lbry/db/queries/base.py |  3 +++
 lbry/db/queries/txio.py | 14 ++++++++++----
 lbry/db/tables.py       |  7 ++++++-
 3 files changed, 19 insertions(+), 5 deletions(-)

diff --git a/lbry/db/queries/base.py b/lbry/db/queries/base.py
index 35e35673f..dade83f90 100644
--- a/lbry/db/queries/base.py
+++ b/lbry/db/queries/base.py
@@ -8,6 +8,7 @@ from ..query_context import context
 from ..tables import (
     SCHEMA_VERSION, metadata, Version,
     Claim, Support, Block, BlockFilter, BlockGroupFilter, TX, TXFilter,
+    pg_add_account_address_constraints_and_indexes
 )
 
 
@@ -105,6 +106,8 @@ def check_version_and_create_tables():
         ctx.execute(Version.insert().values(version=SCHEMA_VERSION))
         for table in metadata.sorted_tables:
             disable_trigger_and_constraints(table.name)
+        if ctx.is_postgres:
+            ctx.execute(text(pg_add_account_address_constraints_and_indexes))
 
 
 def disable_trigger_and_constraints(table_name):
diff --git a/lbry/db/queries/txio.py b/lbry/db/queries/txio.py
index faccaaf88..51ee95844 100644
--- a/lbry/db/queries/txio.py
+++ b/lbry/db/queries/txio.py
@@ -535,7 +535,14 @@ def get_txo_sum(**constraints):
 
 
 def get_balance(account_ids):
+    ctx = context()
     my_addresses = select(AccountAddress.c.address).where(in_account_ids(account_ids))
+    if ctx.is_postgres:
+        txo_address_check = TXO.c.address == func.any(func.array(my_addresses))
+        txi_address_check = TXI.c.address == func.any(func.array(my_addresses))
+    else:
+        txo_address_check = TXO.c.address.in_(my_addresses)
+        txi_address_check = TXI.c.address.in_(my_addresses)
     query = (
         select(
             func.sum(TXO.c.amount).label("total"),
@@ -553,18 +560,17 @@ def get_balance(account_ids):
             )).label("supports"),
             func.sum(case(
                 [(where_txo_type_in(TXO_TYPES["support"]) & (
-                   (TXI.c.address.isnot(None)) &
-                   (TXI.c.address.in_(my_addresses))
+                   (TXI.c.address.isnot(None)) & txi_address_check
                 ), TXO.c.amount)],
                 else_=0
             )).label("my_supports"),
         )
-        .where((TXO.c.spent_height == 0) & (TXO.c.address.in_(my_addresses)))
+        .where((TXO.c.spent_height == 0) & txo_address_check)
         .select_from(
             TXO.join(TXI, (TXI.c.position == 0) & (TXI.c.tx_hash == TXO.c.tx_hash), isouter=True)
         )
     )
-    result = context().fetchone(query)
+    result = ctx.fetchone(query)
     return {
         "total": result["total"],
         "available": result["total"] - result["reserved"],
diff --git a/lbry/db/tables.py b/lbry/db/tables.py
index 4c861dd78..88cdf4e37 100644
--- a/lbry/db/tables.py
+++ b/lbry/db/tables.py
@@ -45,6 +45,11 @@ AccountAddress = Table(
 )
 
 
+pg_add_account_address_constraints_and_indexes = [
+    "CREATE UNIQUE INDEX account_address_idx ON account_address (account, address);"
+]
+
+
 Block = Table(
     'block', metadata,
     Column('height', Integer, primary_key=True),
@@ -160,7 +165,7 @@ pg_add_txo_constraints_and_indexes = [
     f"CREATE INDEX txo_unspent_supports ON txo (claim_hash) INCLUDE (amount) "
     f"WHERE spent_height = 0 AND txo_type={TXO_TYPES['support']};",
     # for calculating balance
-    f"CREATE INDEX txo_unspent_by_address ON txo (address) INCLUDE (amount) "
+    f"CREATE INDEX txo_unspent_by_address ON txo (address) INCLUDE (amount, txo_type, tx_hash) "
     f"WHERE spent_height = 0;",
     # for finding modified claims in a block range
     f"CREATE INDEX txo_claim_changes "