From c5dc3aa0f9efde79deae2dff7f25c25dc41fb748 Mon Sep 17 00:00:00 2001
From: Luke Van Seters <lukevanseters@gmail.com>
Date: Fri, 6 Aug 2021 13:10:23 -0400
Subject: [PATCH] Add join table for swaps <=> arbitrages. Write to it

---
 ...ae51c5d56_add_swap_arbitrage_join_table.py | 38 +++++++++++++++
 mev_inspect/crud/arbitrages.py                | 47 ++++++++++++++-----
 2 files changed, 73 insertions(+), 12 deletions(-)
 create mode 100644 alembic/versions/9b8ae51c5d56_add_swap_arbitrage_join_table.py

diff --git a/alembic/versions/9b8ae51c5d56_add_swap_arbitrage_join_table.py b/alembic/versions/9b8ae51c5d56_add_swap_arbitrage_join_table.py
new file mode 100644
index 0000000..390c7a7
--- /dev/null
+++ b/alembic/versions/9b8ae51c5d56_add_swap_arbitrage_join_table.py
@@ -0,0 +1,38 @@
+"""Add swap arbitrage join table
+
+Revision ID: 9b8ae51c5d56
+Revises: 7eec417a4f3e
+Create Date: 2021-08-06 17:06:55.364516
+
+"""
+import sqlalchemy as sa
+from alembic import op
+
+
+# revision identifiers, used by Alembic.
+revision = "9b8ae51c5d56"
+down_revision = "7eec417a4f3e"
+branch_labels = None
+depends_on = None
+
+
+def upgrade():
+    op.create_table(
+        "arbitrage_swaps",
+        sa.Column("created_at", sa.TIMESTAMP, server_default=sa.func.now()),
+        sa.Column("arbitrage_id", sa.String(1024), primary_key=True),
+        sa.Column("swap_transaction_hash", sa.String(66), primary_key=True),
+        sa.Column("swap_trace_address", sa.ARRAY(sa.Integer), primary_key=True),
+        sa.ForeignKeyConstraint(
+            ["arbitrage_id"], ["arbitrages.id"], ondelete="CASCADE"
+        ),
+        sa.ForeignKeyConstraint(
+            ["swap_transaction_hash", "swap_trace_address"],
+            ["swaps.transaction_hash", "swaps.trace_address"],
+            ondelete="CASCADE",
+        ),
+    )
+
+
+def downgrade():
+    op.drop_table("arbitrage_swaps")
diff --git a/mev_inspect/crud/arbitrages.py b/mev_inspect/crud/arbitrages.py
index 25b6f50..364166d 100644
--- a/mev_inspect/crud/arbitrages.py
+++ b/mev_inspect/crud/arbitrages.py
@@ -22,19 +22,42 @@ def write_arbitrages(
     db_session,
     arbitrages: List[Arbitrage],
 ) -> None:
-    arbitrage_models = [
-        ArbitrageModel(
-            id=str(uuid4()),
-            block_number=arbitrage.block_number,
-            transaction_hash=arbitrage.transaction_hash,
-            account_address=arbitrage.account_address,
-            profit_token_address=arbitrage.profit_token_address,
-            start_amount=arbitrage.start_amount,
-            end_amount=arbitrage.end_amount,
-            profit_amount=arbitrage.profit_amount,
+    arbitrage_models = []
+    swap_arbitrage_ids = []
+
+    for arbitrage in arbitrages:
+        arbitrage_id = str(uuid4())
+        arbitrage_models.append(
+            ArbitrageModel(
+                id=arbitrage_id,
+                block_number=arbitrage.block_number,
+                transaction_hash=arbitrage.transaction_hash,
+                account_address=arbitrage.account_address,
+                profit_token_address=arbitrage.profit_token_address,
+                start_amount=arbitrage.start_amount,
+                end_amount=arbitrage.end_amount,
+                profit_amount=arbitrage.profit_amount,
+            )
         )
-        for arbitrage in arbitrages
-    ]
+
+        for swap in arbitrage.swaps:
+            swap_arbitrage_ids.append(
+                {
+                    "arbitrage_id": arbitrage_id,
+                    "swap_transaction_hash": swap.transaction_hash,
+                    "swap_trace_address": swap.trace_address,
+                }
+            )
 
     db_session.bulk_save_objects(arbitrage_models)
+    db_session.execute(
+        """
+        INSERT INTO arbitrage_swaps
+        (arbitrage_id, swap_transaction_hash, swap_trace_address)
+        VALUES
+        (:arbitrage_id, :swap_transaction_hash, :swap_trace_address)
+        """,
+        params=swap_arbitrage_ids,
+    )
+
     db_session.commit()