Update utils/import_savant_csvs.py

utils/import_savant_csvs.py

@@ -1,5 +1,10 @@
+from __future__ import annotations
+
 from pathlib import Path
+from typing import Iterable
+
 import pandas as pd
+from sqlalchemy import text
 
 from database.remote_db import get_connection
 
@@ -7,48 +12,84 @@ from database.remote_db import get_connection
 BATTER_CSV_PATH = Path("data/batter_savant_data.csv")
 PITCHER_CSV_PATH = Path("data/pitcher_savant_data.csv")
 
+CHUNK_SIZE = 500
 
-def import_savant_csvs():
 
-    if not BATTER_CSV_PATH.exists():
-        raise FileNotFoundError("Missing batter CSV")
+def _clean_dataframe(df: pd.DataFrame, source_file: str) -> pd.DataFrame:
+    df = df.copy()
+    df.columns = [str(c).strip().lower() for c in df.columns]
+    df["source_file"] = source_file
+    df = df.where(pd.notnull(df), None)
+    return df
 
-    if not PITCHER_CSV_PATH.exists():
-        raise FileNotFoundError("Missing pitcher CSV")
 
-    batter_df = pd.read_csv(BATTER_CSV_PATH)
-    pitcher_df = pd.read_csv(PITCHER_CSV_PATH)
+def _chunk_dataframe(df: pd.DataFrame, chunk_size: int) -> Iterable[pd.DataFrame]:
+    for start in range(0, len(df), chunk_size):
+        yield df.iloc[start:start + chunk_size].copy()
 
-    batter_df.columns = [c.lower() for c in batter_df.columns]
-    pitcher_df.columns = [c.lower() for c in pitcher_df.columns]
 
-    conn = get_connection()
+def _truncate_if_requested(table_name: str, clear_first: bool) -> None:
+    if not clear_first:
+        return
 
+    conn = get_connection()
     try:
-        batter_df.to_sql(
-            "mlb_batter_statcast_features",
-            conn,
-            if_exists="append",
-            index=False,
-            method="multi",
-            chunksize=1000,
-        )
-
-        pitcher_df.to_sql(
-            "mlb_pitcher_statcast_features",
-            conn,
-            if_exists="append",
-            index=False,
-            method="multi",
-            chunksize=1000,
-        )
-
+        conn.execute(text(f"TRUNCATE TABLE {table_name}"))
         conn.commit()
-
     finally:
         conn.close()
 
+
+def _import_dataframe(df: pd.DataFrame, table_name: str, clear_first: bool = False) -> dict:
+    _truncate_if_requested(table_name, clear_first=clear_first)
+
+    total_rows = len(df)
+    inserted_rows = 0
+
+    for chunk in _chunk_dataframe(df, CHUNK_SIZE):
+        conn = get_connection()
+        try:
+            chunk.to_sql(
+                table_name,
+                conn,
+                if_exists="append",
+                index=False,
+                method="multi",
+                chunksize=250,
+            )
+            conn.commit()
+            inserted_rows += len(chunk)
+        finally:
+            conn.close()
+
     return {
-        "batter_rows": len(batter_df),
-        "pitcher_rows": len(pitcher_df),
-    }
+        "table_name": table_name,
+        "total_rows": total_rows,
+        "inserted_rows": inserted_rows,
+    }
+
+
+def import_batter_savant_csv(clear_first: bool = False) -> dict:
+    if not BATTER_CSV_PATH.exists():
+        raise FileNotFoundError(f"Missing file: {BATTER_CSV_PATH}")
+
+    batter_df = pd.read_csv(BATTER_CSV_PATH)
+    batter_df = _clean_dataframe(batter_df, "batter_savant_data.csv")
+    return _import_dataframe(
+        batter_df,
+        "mlb_batter_statcast_features",
+        clear_first=clear_first,
+    )
+
+
+def import_pitcher_savant_csv(clear_first: bool = False) -> dict:
+    if not PITCHER_CSV_PATH.exists():
+        raise FileNotFoundError(f"Missing file: {PITCHER_CSV_PATH}")
+
+    pitcher_df = pd.read_csv(PITCHER_CSV_PATH)
+    pitcher_df = _clean_dataframe(pitcher_df, "pitcher_savant_data.csv")
+    return _import_dataframe(
+        pitcher_df,
+        "mlb_pitcher_statcast_features",
+        clear_first=clear_first,
+    )