From 19a4140a882283e1e087c2e94bf8b99739eb0e75 Mon Sep 17 00:00:00 2001
From: thePR0M3TH3AN <53631862+PR0M3TH3AN@users.noreply.github.com>
Date: Sat, 31 May 2025 18:02:34 -0400
Subject: [PATCH] Enhance incremental data refresh

---
 scripts/refresh_data.py    | 30 +++++++++++++++++++++++++++---
 tests/test_refresh_data.py | 30 ++++++++++++++++++++++++++++++
 2 files changed, 57 insertions(+), 3 deletions(-)

diff --git a/scripts/refresh_data.py b/scripts/refresh_data.py
index 04ec3a4..f0e46f9 100644
--- a/scripts/refresh_data.py
+++ b/scripts/refresh_data.py
@@ -12,7 +12,7 @@ Example:
 
 from __future__ import annotations
 
-from datetime import date
+from datetime import date, timedelta
 import argparse
 import pathlib
 import sqlite3
@@ -49,9 +49,33 @@ def main(argv: list[str] | None = None) -> None:
     DATA_PATH.mkdir(exist_ok=True)
     with sqlite3.connect(DB_FILE) as conn:
         for s in args.series:
+            start = pd.to_datetime(args.start)
+            end = pd.to_datetime(args.end)
+
+            last = None
+            table_exists = bool(
+                conn.execute(
+                    "SELECT name FROM sqlite_master WHERE type='table' AND name=?",
+                    (s,),
+                ).fetchone()
+            )
+            if table_exists:
+                last_row = conn.execute(f"SELECT MAX(date) FROM {s}").fetchone()
+                if last_row and last_row[0] is not None:
+                    last = pd.to_datetime(last_row[0])
+
+            new_start = start
+            if last is not None:
+                new_start = max(start, last + timedelta(days=1))
+
+            if new_start > end:
+                logger.info("%s up to date; last date %s", s, last.date() if last else "n/a")
+                continue
+
             logger.info("\u21af downloading %s …", s)
-            df = DataReader(s, "fred", args.start, args.end)
-            df.to_sql(s, conn, if_exists="replace", index_label="date")
+            df = DataReader(s, "fred", new_start.date().isoformat(), end.date().isoformat())
+            if_exists = "append" if table_exists else "replace"
+            df.to_sql(s, conn, if_exists=if_exists, index_label="date")
             logger.info("\u2713 wrote %s: %s rows", s, len(df))
     logger.info(
         "\nDone.  SQLite DB at %s (%s KB)",
diff --git a/tests/test_refresh_data.py b/tests/test_refresh_data.py
index 688776f..4c176b1 100644
--- a/tests/test_refresh_data.py
+++ b/tests/test_refresh_data.py
@@ -42,3 +42,33 @@ def test_refresh_data_writes_series(tmp_path, monkeypatch):
             for row in conn.execute("SELECT name FROM sqlite_master WHERE type='table'")
         }
     assert {"A", "B"} <= tables
+
+
+def test_refresh_data_appends_only_new_rows(tmp_path, monkeypatch):
+    calls = []
+
+    def fake_reader(name, src, start, end):
+        calls.append((name, start, end))
+        dates = pd.date_range(start, end)
+        return pd.DataFrame({name: range(len(dates))}, index=dates)
+
+    monkeypatch.setattr(refresh_data, "DataReader", fake_reader)
+    monkeypatch.setattr(refresh_data, "DATA_PATH", tmp_path)
+    db_file = tmp_path / "fred.db"
+    monkeypatch.setattr(refresh_data, "DB_FILE", db_file)
+
+    # Pre-populate DB with two days of data
+    pre_df = pd.DataFrame({"A": [1, 2]}, index=pd.date_range("2000-01-01", periods=2))
+    with sqlite3.connect(db_file) as conn:
+        pre_df.to_sql("A", conn, if_exists="replace", index_label="date")
+
+    refresh_data.main(
+        ["--series", "A", "--start", "2000-01-01", "--end", "2000-01-04"]
+    )
+
+    # Should only fetch after 2000-01-02
+    assert calls == [("A", "2000-01-03", "2000-01-04")]
+
+    with sqlite3.connect(db_file) as conn:
+        df = pd.read_sql("SELECT * FROM A", conn, parse_dates=["date"], index_col="date")
+    assert len(df) == 4