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

Write to arbitrage protocols column
2022-01-20 19:15:23 -05:00 · 2022-01-20 19:15:23 -05:00 · 8fd382e4b1
commit 8fd382e4b1
parent ac47974daf f37de76824
10 changed files with 27 additions and 15 deletions
--- a/mev_inspect/abis/uniswap_v2/UniswapV2Pair.json
+++ b/mev_inspect/abis/uniswap_v2/UniswapV2Pair.json
--- a/mev_inspect/abis/uniswap_v3/UniswapV3Pool.json
+++ b/mev_inspect/abis/uniswap_v3/UniswapV3Pool.json
--- a/mev_inspect/classifiers/specs/uniswap.py
+++ b/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,
    },
--- a/mev_inspect/crud/arbitrages.py
+++ b/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},
            )
        )

--- a/mev_inspect/models/arbitrages.py
+++ b/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))
--- a/mev_inspect/schemas/swaps.py
+++ b/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]
--- a/tests/helpers.py
+++ b/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,
 ):
--- a/tests/sandwiches/12775690.json
+++ b/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
    }
-]
+]
--- a/tests/test_arbitrages.py
+++ b/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",
--- a/tests/test_swaps.py
+++ b/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