From 4720e1ee41062f143c66f463555ac51b11288d25 Mon Sep 17 00:00:00 2001
From: Alexey Safonov <safonov.a.r@mail.ru>
Date: Mon, 12 Feb 2024 17:15:57 +0300
Subject: [PATCH] feat: allow to query gateway's height (#5)

* feat: allow to query gateway's height

* refactor: use network singltone for get_network_height
---
 ape_subsquid/__init__.py |  4 +++-
 ape_subsquid/gateway.py  |  3 +++
 ape_subsquid/networks.py |  8 ++++----
 ape_subsquid/query.py    | 18 +++++++++++++-----
 4 files changed, 23 insertions(+), 10 deletions(-)

diff --git a/ape_subsquid/__init__.py b/ape_subsquid/__init__.py
index 40d9d18..24c7262 100644
--- a/ape_subsquid/__init__.py
+++ b/ape_subsquid/__init__.py
@@ -1,6 +1,8 @@
 from ape import plugins
 
-from ape_subsquid.query import SubsquidQueryEngine
+from ape_subsquid.query import SubsquidQueryEngine, get_network_height
+
+__all__ = ["exceptions", "get_network_height"]
 
 
 @plugins.register(plugins.QueryPlugin)
diff --git a/ape_subsquid/gateway.py b/ape_subsquid/gateway.py
index bedfd9b..61a0222 100644
--- a/ape_subsquid/gateway.py
+++ b/ape_subsquid/gateway.py
@@ -294,3 +294,6 @@ def _raise_error(self, error: HTTPError) -> ApeSubsquidError:
             raise DataIsNotAvailable(text)
         else:
             raise ApeSubsquidError(text)
+
+
+gateway = SubsquidGateway()
diff --git a/ape_subsquid/networks.py b/ape_subsquid/networks.py
index 000d025..40a72d5 100644
--- a/ape_subsquid/networks.py
+++ b/ape_subsquid/networks.py
@@ -1,9 +1,9 @@
-from ape.api.query import QueryAPI
+from ape.managers.networks import NetworkManager
 
 
-def get_network(engine: QueryAPI) -> str:
-    ecosystem_name = engine.network_manager.ecosystem.name
-    network_name = engine.network_manager.network.name
+def get_network(network_manager: NetworkManager) -> str:
+    ecosystem_name = network_manager.ecosystem.name
+    network_name = network_manager.network.name
 
     if ecosystem_name == "bsc":
         ecosystem_name = "binance"
diff --git a/ape_subsquid/query.py b/ape_subsquid/query.py
index c8e4fad..299c352 100644
--- a/ape_subsquid/query.py
+++ b/ape_subsquid/query.py
@@ -1,5 +1,6 @@
 from typing import Iterator, Optional, Type, TypeVar, cast
 
+from ape import networks
 from ape.api import BlockAPI, ReceiptAPI
 from ape.api.query import (
     AccountTransactionQuery,
@@ -23,13 +24,14 @@
     Query,
     SubsquidGateway,
     TxFieldSelection,
+    gateway,
 )
 from ape_subsquid.mappings import map_header, map_log, map_receipt
 from ape_subsquid.networks import get_network
 
 
 class SubsquidQueryEngine(QueryAPI):
-    _gateway = SubsquidGateway()
+    _gateway = gateway
 
     @singledispatchmethod
     def estimate_query(self, query: QueryType) -> Optional[int]:  # type: ignore[override]
@@ -61,7 +63,7 @@ def perform_query(self, query: QueryType) -> Iterator:  # type: ignore[override]
 
     @perform_query.register
     def perform_block_query(self, query: BlockQuery) -> Iterator[BlockAPI]:
-        network = get_network(self)
+        network = get_network(self.network_manager)
         q: Query = {
             "fromBlock": query.start_block,
             "toBlock": query.stop_block,
@@ -79,7 +81,7 @@ def perform_block_query(self, query: BlockQuery) -> Iterator[BlockAPI]:
     def perform_account_transaction_query(
         self, query: AccountTransactionQuery
     ) -> Iterator[ReceiptAPI]:
-        network = get_network(self)
+        network = get_network(self.network_manager)
         q: Query = {
             "fromBlock": 0,
             "fields": {
@@ -118,7 +120,7 @@ def perform_account_transaction_query(
 
     @perform_query.register
     def perform_contract_creation_query(self, query: ContractCreationQuery) -> Iterator[ReceiptAPI]:
-        network = get_network(self)
+        network = get_network(self.network_manager)
         contract = query.contract.lower()
         q: Query = {
             "fromBlock": query.start_block,
@@ -160,7 +162,7 @@ def perform_contract_creation_query(self, query: ContractCreationQuery) -> Itera
 
     @perform_query.register
     def perform_contract_event_query(self, query: ContractEventQuery) -> Iterator[ContractLog]:
-        network = get_network(self)
+        network = get_network(self.network_manager)
         if isinstance(query.contract, list):
             address = [address.lower() for address in query.contract]
         else:
@@ -212,3 +214,9 @@ def gateway_ingest(gateway: SubsquidGateway, network: str, query: Query) -> Iter
                 break
 
         query["fromBlock"] = last_block + 1
+
+
+def get_network_height() -> int:
+    network = get_network(networks)
+    height = gateway.get_height(network)
+    return height