mirror of
https://github.com/PR0M3TH3AN/python_charts.git
synced 2026-10-27 08:50:39 +00:00
Enhance incremental data refresh
This commit is contained in:
1 parent
ed58cd2f43
commit
19a4140a88
2 files changed
+57
-3
No files matched your search
+27
-3
@@ -12,7 +12,7 @@ Example:
|
|||||||
|
|
||||||
from __future__ import annotations
|
from __future__ import annotations
|
||||||
|
|
||||||
from datetime import date
|
from datetime import date, timedelta
|
||||||
import argparse
|
import argparse
|
||||||
import pathlib
|
import pathlib
|
||||||
import sqlite3
|
import sqlite3
|
||||||
@@ -49,9 +49,33 @@ def main(argv: list[str] | None = None) -> None:
|
|||||||
DATA_PATH.mkdir(exist_ok=True)
|
DATA_PATH.mkdir(exist_ok=True)
|
||||||
with sqlite3.connect(DB_FILE) as conn:
|
with sqlite3.connect(DB_FILE) as conn:
|
||||||
for s in args.series:
|
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)
|
logger.info("\u21af downloading %s …", s)
|
||||||
df = DataReader(s, "fred", args.start, args.end)
|
df = DataReader(s, "fred", new_start.date().isoformat(), end.date().isoformat())
|
||||||
df.to_sql(s, conn, if_exists="replace", index_label="date")
|
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("\u2713 wrote %s: %s rows", s, len(df))
|
||||||
logger.info(
|
logger.info(
|
||||||
"\nDone. SQLite DB at %s (%s KB)",
|
"\nDone. SQLite DB at %s (%s KB)",
|
||||||
|
|||||||
@@ -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'")
|
for row in conn.execute("SELECT name FROM sqlite_master WHERE type='table'")
|
||||||
}
|
}
|
||||||
assert {"A", "B"} <= tables
|
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
|
||||||
Reference in new issue
Block a user