Merge pull request #38 from BalancerMaxis/feat/query_all_gauges

feat: query all pools and gauges via apiv3

bal_tools/graphql/apiv3/get_gauges.gql

@@ -0,0 +1,14 @@
+query AllGauges($chain: GqlChain!) {
+  poolGetPools(where: { chainIn: [$chain] }) {
+    staking {
+      gauge {
+        gaugeAddress
+        otherGauges {
+          id
+        }
+      }
+    }
+    chain
+    symbol
+  }
+}

bal_tools/graphql/apiv3/get_pools.gql

@@ -0,0 +1,9 @@
+query AllPools($chain: GqlChain!) {
+  poolGetPools(where: { chainIn: [$chain] }) {
+    address
+    symbol
+    dynamicData {
+      swapEnabled
+    }
+  }
+}

bal_tools/models.py

@@ -3,6 +3,7 @@
 from dataclasses import dataclass, fields
 from enum import Enum
 from pydantic import BaseModel, field_validator, model_validator, Field
+import json_fix
 
 
 class GqlChain(Enum):
@@ -43,6 +44,8 @@ def __iter__(self):
     def __len__(self):
         return len(self.pools)
 
+    def __json__(self):
+        return self.pools
 
 
 @dataclass
@@ -154,3 +157,21 @@ def validate_model(cls, values):
             if field in values:
                 values[field] = cls.str_to_decimal(values[field])
         return values
+
+
+class PoolData(BaseModel):
+    address: str
+    symbol: str
+    dynamicData: dict
+
+
+class GaugePoolData(BaseModel):
+    staking: Optional[dict]
+    chain: str
+    symbol: str
+
+
+@dataclass
+class GaugeData:
+    address: str
+    symbol: str

bal_tools/pools_gauges.py

@@ -1,11 +1,11 @@
-from typing import Dict
+from typing import Dict, List
 import requests
-from .utils import to_checksum_address
-from .models import CorePools, PoolId, Symbol
+from .utils import to_checksum_address, flatten_nested_dict
 
 from gql.transport.exceptions import TransportQueryError
 from bal_tools.subgraph import Subgraph
 from bal_tools.errors import NoResultError
+from bal_tools.models import PoolData, GaugePoolData, GaugeData, CorePools, PoolId, Symbol
 
 GITHUB_RAW_OUTPUTS = (
     "https://raw.githubusercontent.com/BalancerMaxis/bal_addresses/main/outputs"
@@ -116,6 +116,41 @@ def query_root_gauges(self, skip=0, step_size=100) -> list:
             result += self.query_root_gauges(skip + step_size, step_size)
         return result
 
+    def query_all_gauges(self, include_other_gauges=True) -> List[GaugeData]:
+        """
+        query all gauges from the apiv3 subgraph
+        """
+        data = self.subgraph.fetch_graphql_data("apiv3", "get_gauges", {"chain": self.chain.upper()})
+        all_gauges = []
+        for pool in data["poolGetPools"]:
+            gauge_pool = GaugePoolData(**flatten_nested_dict(pool))
+            if gauge_pool.staking is not None and gauge_pool.staking.get('gauge') is not None:
+                gauge = gauge_pool.staking['gauge']
+                all_gauges.append(GaugeData(
+                    address=gauge['gaugeAddress'],
+                    symbol=f"{gauge_pool.symbol}-gauge"
+                ))
+                if include_other_gauges:
+                    for other_gauge in gauge.get('otherGauges', []):
+                        all_gauges.append(GaugeData(
+                            address=other_gauge['id'],
+                            symbol=f"{gauge_pool.symbol}-gauge"
+                        ))
+        return all_gauges
+
+    def query_all_pools(self) -> List[PoolData]:
+        """
+        query all pools from the apiv3 subgraph
+        filters out disabled pools
+        """
+        data = self.subgraph.fetch_graphql_data("apiv3", "get_pools", {"chain": self.chain.upper()})
+        all_pools = []
+        for pool in data["poolGetPools"]:
+            pool_data = PoolData(**flatten_nested_dict(pool))
+            if pool_data.dynamicData['swapEnabled']:
+                all_pools.append(pool_data)
+        return all_pools
+
     def get_last_join_exit(self, pool_id: int) -> int:
         """
         Returns a timestamp of the last join/exit for a given pool id
@@ -135,7 +170,6 @@ def get_pool_tvl(self, pool_id: str) -> float:
         """
         Returns the TVL of a pool as per the API V3 subgraph
         """
-        print(pool_id)
         try:
             data = self.subgraph.fetch_graphql_data(
                 "apiv3", "get_pool_tvl", {"chain": self.chain.upper(), "poolId": pool_id}

bal_tools/requirements.txt

@@ -7,5 +7,6 @@ munch==4.0.0
 gql[requests]
 python-dotenv
 pydantic==2.7.4
+json-fix
 
 git+https://github.com/BalancerMaxis/[email protected]

bal_tools/utils.py

@@ -27,8 +27,10 @@ def get_abi(contract_name: str) -> Union[Dict, List[Dict]]:
 
 
 def flatten_nested_dict(d):
-    for key, value in list(d.items()):
+    result = d.copy()
+    for key, value in list(result.items()):
         if isinstance(value, dict):
-            d.pop(key)
-            d.update(value)
-    return d
+            if key not in ['dynamicData', 'staking']:
+                result.pop(key)
+                result.update(value)
+    return result

tests/test_pools_gauges.py

@@ -1,5 +1,6 @@
 import pytest
 from gql.transport.exceptions import TransportQueryError
+from bal_tools.models import PoolData, GaugeData
 from bal_tools.models import CorePools
 import json
 
@@ -130,3 +131,22 @@ def test_get_preferential_gauge(bal_pools_gauges):
         pytest.skip(f"Skipping {bal_pools_gauges.chain}, no example preferential gauge")
 
     assert bal_pools_gauges.get_preferential_gauge(example[0]) == example[1]
+
+
+def test_query_all_pools(bal_pools_gauges):
+    """
+    test return data of v3 AllGauges query
+    """
+    response = bal_pools_gauges.query_all_pools()
+
+    if len(response) > 0:
+        assert isinstance(response[0], PoolData)
+
+def test_query_all_gauges(bal_pools_gauges):
+    """
+    test return data of v3 AllPools query
+    """
+    response = bal_pools_gauges.query_all_gauges()
+
+    if len(response) > 0:
+        assert isinstance(response[0], GaugeData)