From c34485c49307b053ba23d10f2335f92c4c12129d Mon Sep 17 00:00:00 2001
From: carlomazzaferro <carlo.mazzaferro@olx.com>
Date: Thu, 25 Nov 2021 15:49:36 +0100
Subject: [PATCH] fix listener

---
 listener.py              |  6 ++++--
 mev_inspect/db.py        |  9 ++++++++-
 mev_inspect/inspector.py | 17 ++++-------------
 3 files changed, 16 insertions(+), 16 deletions(-)

diff --git a/listener.py b/listener.py
index bfa02d3..89fe1d1 100644
--- a/listener.py
+++ b/listener.py
@@ -10,6 +10,7 @@ from mev_inspect.crud.latest_block_update import (
     find_latest_block_update,
     update_latest_block,
 )
+from mev_inspect.db import get_sessions
 from mev_inspect.inspector import MEVInspector
 from mev_inspect.provider import get_base_provider
 from mev_inspect.signal_handler import GracefulKiller
@@ -33,13 +34,14 @@ async def run():
 
     killer = GracefulKiller()
 
-    inspector = MEVInspector(rpc, None, None)
+    inspector = MEVInspector(rpc)
     base_provider = get_base_provider(rpc)
+    _, inspect_session = get_sessions()
 
     while not killer.kill_now:
         await inspect_next_block(
             inspector,
-            None,
+            inspect_session,
             base_provider,
             healthcheck_url,
         )
diff --git a/mev_inspect/db.py b/mev_inspect/db.py
index a01cb52..c44516d 100644
--- a/mev_inspect/db.py
+++ b/mev_inspect/db.py
@@ -1,5 +1,5 @@
 import os
-from typing import Optional
+from typing import Optional, Tuple
 from asyncio import current_task
 
 from sqlalchemy.orm import sessionmaker
@@ -51,3 +51,10 @@ def get_trace_session() -> Optional[async_scoped_session]:
         return _get_session(uri)
 
     return None
+
+
+def get_sessions() -> Tuple[async_scoped_session, Optional[async_scoped_session]]:
+    inspect_db_session = get_inspect_session()
+    trace_db_session = get_trace_session()
+    trace_db_session = trace_db_session() if trace_db_session is not None else None
+    return inspect_db_session, trace_db_session
diff --git a/mev_inspect/inspector.py b/mev_inspect/inspector.py
index dea26b1..8e3cb26 100644
--- a/mev_inspect/inspector.py
+++ b/mev_inspect/inspector.py
@@ -2,15 +2,13 @@ import asyncio
 import logging
 import traceback
 from asyncio import CancelledError
-from typing import Tuple, Optional
 
-from sqlalchemy.ext.asyncio import async_scoped_session
 from web3 import Web3
 from web3.eth import AsyncEth
 
 from mev_inspect.block import create_from_block_number
 from mev_inspect.classifiers.trace import TraceClassifier
-from mev_inspect.db import get_inspect_session, get_trace_session
+from mev_inspect.db import get_sessions
 from mev_inspect.inspect_block import inspect_block
 from mev_inspect.provider import get_base_provider
 
@@ -30,7 +28,7 @@ class MEVInspector:
         self.max_concurrency = asyncio.Semaphore(max_concurrency)
 
     async def create_from_block(self, block_number: int):
-        _, trace_session = _get_sessions()
+        _, trace_session = get_sessions()
         return await create_from_block_number(
             base_provider=self.base_provider,
             w3=self.w3,
@@ -39,7 +37,7 @@ class MEVInspector:
         )
 
     async def inspect_single_block(self, block: int):
-        inspect_session, trace_session = _get_sessions()
+        inspect_session, trace_session = get_sessions()
         return await inspect_block(
             inspect_session,
             trace_session,
@@ -67,7 +65,7 @@ class MEVInspector:
             traceback.print_exc()
 
     async def safe_inspect_block(self, block_number: int):
-        inspect_session, trace_session = _get_sessions()
+        inspect_session, trace_session = get_sessions()
         async with self.max_concurrency:
             return await inspect_block(
                 inspect_session,
@@ -77,10 +75,3 @@ class MEVInspector:
                 self.trace_classifier,
                 block_number,
             )
-
-
-def _get_sessions() -> Tuple[async_scoped_session, Optional[async_scoped_session]]:
-    inspect_db_session = get_inspect_session()
-    trace_db_session = get_trace_session()
-    trace_db_session = trace_db_session() if trace_db_session is not None else None
-    return inspect_db_session, trace_db_session