Merge pull request #246 from flashbots/write-to-arbs-protocol-column

Write to arbitrage protocols column

mev_inspect/classifiers/specs/uniswap.py

@@ -103,6 +103,7 @@ UNISWAP_V3_CONTRACT_SPECS = [
 UNISWAP_V3_GENERAL_SPECS = [
     ClassifierSpec(
         abi_name=UNISWAP_V3_POOL_ABI_NAME,
+        protocol=Protocol.uniswap_v3,
         classifiers={
             "swap(address,bool,int256,uint160,bytes)": UniswapV3SwapClassifier,
         },
@@ -134,6 +135,7 @@ UNISWAPPY_V2_CONTRACT_SPECS = [
 
 UNISWAPPY_V2_PAIR_SPEC = ClassifierSpec(
     abi_name=UNISWAP_V2_PAIR_ABI_NAME,
+    protocol=Protocol.uniswap_v2,
     classifiers={
         "swap(uint256,uint256,address,bytes)": UniswapV2SwapClassifier,
     },

mev_inspect/crud/arbitrages.py

@@ -41,6 +41,7 @@ def write_arbitrages(
                 end_amount=arbitrage.end_amount,
                 profit_amount=arbitrage.profit_amount,
                 error=arbitrage.error,
+                protocols={swap.protocol.value for swap in arbitrage.swaps},
             )
         )
 

mev_inspect/models/arbitrages.py

@@ -1,4 +1,4 @@
-from sqlalchemy import Column, Numeric, String
+from sqlalchemy import ARRAY, Column, Numeric, String
 
 from .base import Base
 
@@ -15,3 +15,4 @@ class ArbitrageModel(Base):
     end_amount = Column(Numeric, nullable=False)
     profit_amount = Column(Numeric, nullable=False)
     error = Column(String, nullable=True)
+    protocols = Column(ARRAY(String))

mev_inspect/schemas/swaps.py

@@ -18,5 +18,5 @@ class Swap(BaseModel):
     token_in_amount: int
     token_out_address: str
     token_out_amount: int
-    protocol: Optional[Protocol]
+    protocol: Protocol
     error: Optional[str]

tests/helpers.py

@@ -1,4 +1,4 @@
-from typing import List, Optional
+from typing import List
 
 from mev_inspect.schemas.traces import (
     Classification,
@@ -48,7 +48,7 @@ def make_swap_trace(
     contract_address: str,
     abi_name: str,
     function_signature: str,
-    protocol: Optional[Protocol],
+    protocol: Protocol,
     recipient_address: str,
     recipient_input_key: str,
 ):

tests/sandwiches/12775690.json

@@ -18,7 +18,7 @@
             "token_in_amount": 12108789017249529876,
             "token_out_address": "0x9813037ee2218799597d83d4a5b6f3b6778218d9",
             "token_out_amount": 1114969767487478978357,
-            "protocol": null,
+            "protocol": "uniswap_v2",
             "error": null
         },
         "backrun_swap": {
@@ -37,7 +37,7 @@
             "token_in_amount": 1114969767487478978357,
             "token_out_address": "0xc02aaa39b223fe8d0a0e5c4f27ead9083c756cc2",
             "token_out_amount": 12158780499164852150,
-            "protocol": null,
+            "protocol": "uniswap_v2",
             "error": null
         },
         "sandwiched_swaps": [
@@ -56,7 +56,7 @@
                 "token_in_amount": 652974555369106606,
                 "token_out_address": "0x9813037ee2218799597d83d4a5b6f3b6778218d9",
                 "token_out_amount": 60000000000000000000,
-                "protocol": null,
+                "protocol": "uniswap_v2",
                 "error": null
             },
             {
@@ -74,7 +74,7 @@
                 "token_in_amount": 300000000000000000,
                 "token_out_address": "0x9813037ee2218799597d83d4a5b6f3b6778218d9",
                 "token_out_amount": 27561865602394087181,
-                "protocol": null,
+                "protocol": "uniswap_v2",
                 "error": null
             },
             {
@@ -92,7 +92,7 @@
                 "token_in_amount": 125000000000000000,
                 "token_out_address": "0x9813037ee2218799597d83d4a5b6f3b6778218d9",
                 "token_out_amount": 11483313070817976324,
-                "protocol": null,
+                "protocol": "uniswap_v2",
                 "error": null
             },
             {
@@ -110,11 +110,11 @@
                 "token_in_amount": 30000000000000000000,
                 "token_out_address": "0x9813037ee2218799597d83d4a5b6f3b6778218d9",
                 "token_out_amount": 2742522049933966038599,
-                "protocol": null,
+                "protocol": "uniswap_v2",
                 "error": null
             }
         ],
         "profit_token_address": "0xc02aaa39b223fe8d0a0e5c4f27ead9083c756cc2",
         "profit_amount": 49991481915322274
     }
-]
+]

tests/test_arbitrages.py

@@ -6,6 +6,7 @@ from mev_inspect.classifiers.specs.uniswap import (
     UNISWAP_V3_POOL_ABI_NAME,
 )
 from mev_inspect.schemas.swaps import Swap
+from mev_inspect.schemas.traces import Protocol
 
 
 def test_two_pool_arbitrage(get_transaction_hashes, get_addresses):
@@ -35,6 +36,7 @@ def test_two_pool_arbitrage(get_transaction_hashes, get_addresses):
             block_number=block_number,
             trace_address=[0],
             contract_address=first_pool_address,
+            protocol=Protocol.uniswap_v2,
             from_address=account_address,
             to_address=second_pool_address,
             token_in_address=first_token_address,
@@ -48,6 +50,7 @@ def test_two_pool_arbitrage(get_transaction_hashes, get_addresses):
             transaction_position=transaction_position,
             block_number=block_number,
             trace_address=[1],
+            protocol=Protocol.uniswap_v3,
             contract_address=second_pool_address,
             from_address=first_pool_address,
             to_address=account_address,
@@ -62,6 +65,7 @@ def test_two_pool_arbitrage(get_transaction_hashes, get_addresses):
         abi_name=UNISWAP_V3_POOL_ABI_NAME,
         transaction_hash=transaction_hash,
         transaction_position=transaction_position,
+        protocol=Protocol.uniswap_v3,
         block_number=block_number,
         trace_address=[2, 0],
         contract_address=unrelated_pool_address,
@@ -117,6 +121,7 @@ def test_three_pool_arbitrage(get_transaction_hashes, get_addresses):
             abi_name=UNISWAP_V2_PAIR_ABI_NAME,
             transaction_hash=transaction_hash,
             transaction_position=transaction_position,
+            protocol=Protocol.uniswap_v2,
             block_number=block_number,
             trace_address=[0],
             contract_address=first_pool_address,
@@ -131,6 +136,7 @@ def test_three_pool_arbitrage(get_transaction_hashes, get_addresses):
             abi_name=UNISWAP_V3_POOL_ABI_NAME,
             transaction_hash=transaction_hash,
             transaction_position=transaction_position,
+            protocol=Protocol.uniswap_v3,
             block_number=block_number,
             trace_address=[1],
             contract_address=second_pool_address,
@@ -145,6 +151,7 @@ def test_three_pool_arbitrage(get_transaction_hashes, get_addresses):
             abi_name=UNISWAP_V3_POOL_ABI_NAME,
             transaction_hash=transaction_hash,
             transaction_position=transaction_position,
+            protocol=Protocol.uniswap_v3,
             block_number=block_number,
             trace_address=[2],
             contract_address=third_pool_address,
@@ -245,6 +252,7 @@ def create_generic_swap(
         abi_name=UNISWAP_V3_POOL_ABI_NAME,
         transaction_hash="0xfake",
         transaction_position=0,
+        protocol=Protocol.uniswap_v2,
         block_number=0,
         trace_address=trace_address,
         contract_address="0xfake",

tests/test_swaps.py

@@ -72,7 +72,7 @@ def test_swaps(
             from_address=alice_address,
             contract_address=first_pool_address,
             abi_name=UNISWAP_V2_PAIR_ABI_NAME,
-            protocol=None,
+            protocol=Protocol.uniswap_v2,
             function_signature="swap(uint256,uint256,address,bytes)",
             recipient_address=bob_address,
             recipient_input_key="to",
@@ -93,7 +93,7 @@ def test_swaps(
             from_address=bob_address,
             contract_address=second_pool_address,
             abi_name=UNISWAP_V3_POOL_ABI_NAME,
-            protocol=None,
+            protocol=Protocol.uniswap_v3,
             function_signature="swap(address,bool,int256,uint160,bytes)",
             recipient_address=carl_address,
             recipient_input_key="recipient",
@@ -198,7 +198,7 @@ def test_swaps(
     assert uni_v2_swap.transaction_hash == first_transaction_hash
     assert uni_v2_swap.block_number == block_number
     assert uni_v2_swap.trace_address == [1]
-    assert uni_v2_swap.protocol is None
+    assert uni_v2_swap.protocol == Protocol.uniswap_v2
     assert uni_v2_swap.contract_address == first_pool_address
     assert uni_v2_swap.from_address == alice_address
     assert uni_v2_swap.to_address == bob_address
@@ -211,7 +211,7 @@ def test_swaps(
     assert uni_v3_swap.transaction_hash == second_transaction_hash
     assert uni_v3_swap.block_number == block_number
     assert uni_v3_swap.trace_address == []
-    assert uni_v3_swap.protocol is None
+    assert uni_v3_swap.protocol == Protocol.uniswap_v3
     assert uni_v3_swap.contract_address == second_pool_address
     assert uni_v3_swap.from_address == bob_address
     assert uni_v3_swap.to_address == carl_address