From bab2043575f89d2d7dc8a79b99390016501920d7 Mon Sep 17 00:00:00 2001
From: Luke Van Seters <lukevanseters@gmail.com>
Date: Mon, 3 Jan 2022 13:38:34 -0500
Subject: [PATCH] Abstract out csv writing

---
 mev_inspect/crud/blocks.py | 35 +++++++++--------------------------
 mev_inspect/db.py          | 25 ++++++++++++++++++++++++-
 2 files changed, 33 insertions(+), 27 deletions(-)

diff --git a/mev_inspect/crud/blocks.py b/mev_inspect/crud/blocks.py
index e8d7152..41199a5 100644
--- a/mev_inspect/crud/blocks.py
+++ b/mev_inspect/crud/blocks.py
@@ -1,7 +1,7 @@
-import io
 from datetime import datetime
-from typing import Any, List, Optional
+from typing import List
 
+from mev_inspect.db import write_as_csv
 from mev_inspect.schemas.blocks import Block
 
 
@@ -25,32 +25,15 @@ def delete_blocks(
     db_session.commit()
 
 
-def clean_csv_value(value: Optional[Any]) -> str:
-    if value is None:
-        return r"\N"
-    return str(value).replace("\n", "\\n")
-
-
 def write_blocks(
     db_session,
     blocks: List[Block],
 ) -> None:
-    csv_file_like_object = io.StringIO()
-    for block in blocks:
-        csv_file_like_object.write(
-            "|".join(
-                map(
-                    clean_csv_value,
-                    (
-                        block.block_number,
-                        datetime.fromtimestamp(block.block_timestamp),
-                    ),
-                )
-            )
-            + "\n"
+    items_generator = (
+        (
+            block.block_number,
+            datetime.fromtimestamp(block.block_timestamp),
         )
-
-    csv_file_like_object.seek(0)
-
-    with db_session.connection().connection.cursor() as cursor:
-        cursor.copy_from(csv_file_like_object, "blocks", sep="|")
+        for block in blocks
+    )
+    write_as_csv(db_session, "blocks", items_generator)
diff --git a/mev_inspect/db.py b/mev_inspect/db.py
index 15ccdc3..38f588e 100644
--- a/mev_inspect/db.py
+++ b/mev_inspect/db.py
@@ -1,5 +1,6 @@
+import io
 import os
-from typing import Optional
+from typing import Any, Iterable, Optional
 
 from sqlalchemy import create_engine, orm
 from sqlalchemy.orm import sessionmaker
@@ -63,3 +64,25 @@ def get_trace_session() -> Optional[orm.Session]:
         return Session()
 
     return None
+
+
+def _clean_csv_value(value: Optional[Any]) -> str:
+    if value is None:
+        return r"\N"
+    return str(value).replace("\n", "\\n")
+
+
+def write_as_csv(
+    db_session,
+    table_name: str,
+    items: Iterable[Iterable[Any]],
+) -> None:
+    csv_file_like_object = io.StringIO()
+
+    for item in items:
+        csv_file_like_object.write("|".join(map(_clean_csv_value, item)) + "\n")
+
+    csv_file_like_object.seek(0)
+
+    with db_session.connection().connection.cursor() as cursor:
+        cursor.copy_from(csv_file_like_object, table_name, sep="|")