From c8c4f5b7ae77c8d42939733339bf348737d9c65c Mon Sep 17 00:00:00 2001
From: tmcgroul <safonov.a.r@mail.ru>
Date: Tue, 13 Feb 2024 13:42:48 +0300
Subject: [PATCH] fix: skip estimation if block is not available

---
 ape_subsquid/gateway.py | 19 +++++++++++++------
 ape_subsquid/query.py   | 29 ++++++++++++++++++++++++++---
 2 files changed, 39 insertions(+), 9 deletions(-)

diff --git a/ape_subsquid/gateway.py b/ape_subsquid/gateway.py
index 61a0222..7ec1761 100644
--- a/ape_subsquid/gateway.py
+++ b/ape_subsquid/gateway.py
@@ -241,11 +241,11 @@ class SubsquidGateway:
     _retry_schedule = [5, 10, 20, 30, 60]
 
     @ttl_cache(seconds=30)
-    def get_height(self, network: str) -> int:
-        return self._retry(self._get_height, network)
+    def get_height(self, network: str, **kwargs) -> int:
+        return self._retry(self._get_height, network, **kwargs)
 
-    def query(self, network: str, query: Query) -> list[Block]:
-        return self._retry(self._query, network, query)
+    def query(self, network: str, query: Query, **kwargs) -> list[Block]:
+        return self._retry(self._query, network, query, **kwargs)
 
     def _query(self, network: str, query: Query) -> list[Block]:
         worker_url = self._get_worker(network, query["fromBlock"])
@@ -267,12 +267,13 @@ def _get_height(self, network: str) -> int:
 
     def _retry(self, request: Callable[..., T], *args, **kwargs) -> T:
         retries = 0
+        max_retries = kwargs.pop("max_retries", len(self._retry_schedule))
         while True:
             try:
                 response = request(*args, **kwargs)
             except HTTPError as e:
-                if self._is_retryable_error(e) and retries < len(self._retry_schedule):
-                    pause = self._retry_schedule[retries]
+                if self._is_retryable_error(e) and retries < max_retries:
+                    pause = self._get_retry_pause(retries)
                     retries += 1
                     logger.warning(f"Gateway request failed, will retry in {pause} secs")
                     sleep(pause)
@@ -281,6 +282,12 @@ def _retry(self, request: Callable[..., T], *args, **kwargs) -> T:
             else:
                 return response
 
+    def _get_retry_pause(self, retries: int) -> int:
+        if retries < len(self._retry_schedule):
+            return self._retry_schedule[retries]
+        else:
+            return self._retry_schedule[-1]
+
     def _is_retryable_error(self, error: HTTPError) -> bool:
         assert error.response is not None
         return error.response.status_code == 503
diff --git a/ape_subsquid/query.py b/ape_subsquid/query.py
index 175f530..4891c59 100644
--- a/ape_subsquid/query.py
+++ b/ape_subsquid/query.py
@@ -37,7 +37,11 @@ def estimate_query(self, query: QueryType) -> Optional[int]:  # type: ignore[ove
         return None
 
     @estimate_query.register
-    def estimate_block_query(self, query: BlockQuery) -> int:
+    def estimate_block_query(self, query: BlockQuery) -> Optional[int]:
+        network = get_network(self.network_manager)
+        if not block_is_available(self._gateway, network, query.stop_block):
+            return None
+
         return 100 + (query.stop_block - query.start_block) * 4
 
     @estimate_query.register
@@ -47,11 +51,19 @@ def estimate_account_transaction_query(self, query: AccountTransactionQuery) ->
         return 1000 * 60 * 10
 
     @estimate_query.register
-    def estimate_contract_creation_query(self, query: ContractCreationQuery) -> int:
+    def estimate_contract_creation_query(self, query: ContractCreationQuery) -> Optional[int]:
+        network = get_network(self.network_manager)
+        if not block_is_available(self._gateway, network, query.stop_block):
+            return None
+
         return 100 + (query.stop_block - query.start_block) * 5
 
     @estimate_query.register
-    def estimate_contract_event_query(self, query: ContractEventQuery) -> int:
+    def estimate_contract_event_query(self, query: ContractEventQuery) -> Optional[int]:
+        network = get_network(self.network_manager)
+        if not block_is_available(self._gateway, network, query.stop_block):
+            return None
+
         return 400 + (query.stop_block - query.start_block) * 4
 
     @singledispatchmethod
@@ -193,6 +205,17 @@ def all_fields(cls: Type[T]) -> T:
     return cast(T, fields)
 
 
+def block_is_available(gateway: SubsquidGateway, network: str, block_num: int) -> bool:
+    try:
+        height = gateway.get_height(network, max_retries=0)
+    except Exception:
+        return False
+    else:
+        if block_num > height:
+            return False
+    return True
+
+
 def ensure_range_is_available(gateway: SubsquidGateway, network: str, query: Query):
     height = gateway.get_height(network)
     if query["toBlock"] > height: