642 lines
26 KiB
Python
642 lines
26 KiB
Python
import argparse
|
||
import sys, time, re, requests, random, os
|
||
from typing import List, Tuple, Set, Optional, Iterable
|
||
import pandas as pd
|
||
import yfinance as yf
|
||
from bs4 import BeautifulSoup
|
||
from datetime import datetime, timedelta, timezone, date
|
||
from io import StringIO
|
||
import psycopg2
|
||
from psycopg2.extras import execute_values
|
||
from dotenv import load_dotenv
|
||
try:
|
||
from zoneinfo import ZoneInfo
|
||
CT = ZoneInfo("America/Chicago")
|
||
except Exception:
|
||
CT = None
|
||
|
||
load_dotenv()
|
||
|
||
|
||
|
||
# ========= USER CONFIG =======================================================
|
||
# PostgreSQL connection - uses environment variables or defaults to local
|
||
# Set DATABASE_URL or USE_LOCAL_DB=true with LOCAL_DB_* variables
|
||
SEGMENTS = ["large", "mid", "nasdaq"] # auto universes to include
|
||
TAG_AUTO_SEGMENTS = True # tag auto symbols as AUTO:SP500/SP400/NDX100
|
||
EXTRA_SYMBOLS = [] # optional flat list, e.g. ["TSLA","COIN","SPY"]
|
||
|
||
EXTRA_GROUPS = {
|
||
"OIL": ["HUSA","INDO","USEG","MARPS","BRN","PED","GBR","KLXE","CEIN","IMPP","CVX","OXY","GUSH","RCON","APA","VIVK"],
|
||
"FASTFOOD_REST": ["SG","CMG","WEN","QSR","WING","BROS","SHAK","CAVA","SBUX","MCD","YUM"],
|
||
"ROBOTICS": ["SERV","RR","KITT","MBOT","IRBT","ARBE","SYM","TSLA","NVDA","ROBO","BOTZ","TER","ROK"],
|
||
"SEMIS": ["NVDA","MCHP","NVTS","ADI","QCOM","AVGO","ON","TXN","QQQ","OLED","STM","ASML","RMBS","ARM","SMH","INTC","LRCX","DELL","MRVL","TSM","SMCI","MU","SOXL","AMD"],
|
||
"QUANTUM": ["IONQ","RGTI","QUBT","QBTS","ARQQ","QMCO","COHR","MRVL","HON","TSEM","FORM"],
|
||
"CYBER": ["CYBR","FTNT","PANW","ZS","CRWD","OKTA","S"],
|
||
"FINTECH": ["SOFI","AFRM","NU","V","MA","PYPL","UPST","FOUR","AXP","PSFE","PAYO","ADYEY","QFIN"],
|
||
"CATHY": ["TXG","TWST","CRSP","COIN","PACB","DNA","SOFI","PD","HOOD","ROKU","ZM","EXAS","U","NTLA","SPY","TER","PATH","PINS","META","ARKK","CERS","BEAM","VCYT","RBLX","TTD","TSLA","TDOC","ACHR"],
|
||
"CHINA": ["FUTU","KC","LI","EDU","BILI","KWEB","PDD","JD","BABA","BIDU","NIO","XPEV","GDS"],
|
||
"SOLAR": ["DQ","SEDG","JKS","ENPH","TAN","CSIQ","ARRY","FSLR","RUN"],
|
||
"BANKS": ["MS","C","GS","KRE","WFC","JPM","BAC","SCHW","XLF"],
|
||
"TRAVEL": ["EXPE","BKNG","WYNN","LVS","TRIP","ABNB","JETS","RCL","LUV","DAL","UAL","NCLH","CCL","MAR","HLT"],
|
||
"WEED": ["ACB","MSOS","CGC","CURLF","GTBIF","TLRY","SNDL","MSOX","MJ"],
|
||
"DOW": ["WMT","CRM","PG","HD","MSFT","V","MCD","IBM","AMGN","VZ","CSCO","AAPL","MRK","DIS","TRV","HON","JNJ","JPM","NKE","AXP","MMM","GS","CVX","CAT","UNH","BA","BE"],
|
||
"NUCLEAR": ["SMR","OKLO","NNE","ASPI","LTBR","LEU","TLN","CEG","PEG"],
|
||
"URANIUM": ["URA","NLR","URNM","NXE","URAN","URNJ","CCJ","UEC"],
|
||
"SAAS-CLOUD": ["NOW","CRM","WDAY","OKTA","MDB","DDOG","TWLO","DOCU","WIX","TEAM","SNOW","NET","U","PLTR","BILL","DOCN","VEEV","CRWV"],
|
||
"SPACE-ROCKET": ["RKLB","SPCE","SIDU","MNTS","RDW","LUNR","ASTS","SATS","ARKX","VSAT","IRDM","SES"],
|
||
"MEME": ["NEGG","AMC","BBAI","BYND","OPEN","UPST","NVAX","GME","DJT","BB","TRUP","KOSS","KODK","SPWR","RDDT","HOOD"],
|
||
"DRONES": ["NOC","LMT","KTOS","BA","AIRO","UMAC","DPRO","ONDS","RTX","RCAT","EVTL","AMBA","ACHR","JOBY","GPRO","HON","GOGO","EH"]
|
||
}
|
||
SYMBOLS_ONLY = False # True = ignore SEGMENTS and use only EXTRA_SYMBOLS + EXTRA_GROUPS
|
||
LOOKBACK_DAYS_INTRADAY = 30 # 30 days (1 month) - uses 5m interval for >7 days
|
||
DAILY_PERIOD = "2y"
|
||
BATCH_SIZE = 110 # 50–120 good
|
||
SLEEP_BETWEEN_BATCHES = 0.25 # seconds
|
||
LOOP = False # Set to False for one-shot update, True for continuous loop
|
||
INTRADAY_INTERVAL_MIN = 1
|
||
DAILY_INTERVAL_MIN = 1
|
||
# ============================================================================
|
||
|
||
USER_AGENT = "Mozilla/5.0 (Windows NT 10.0; Win64; x64)"
|
||
HEADERS = {"User-Agent": USER_AGENT}
|
||
URL_SP500 = "https://en.wikipedia.org/wiki/List_of_S%26P_500_companies"
|
||
URL_SP400 = "https://en.wikipedia.org/wiki/List_of_S%26P_400_companies"
|
||
URL_NASDAQ = "https://en.wikipedia.org/wiki/Nasdaq-100"
|
||
VALID_TICKER_RE = re.compile(r"^[A-Z][A-Z0-9.\-]{0,6}$")
|
||
BAD_WORDS = {"CLOSING","INTRADAY","TICKER","SYMBOL","INDEX","COMPANY","WEIGHT","WEIGHTS"}
|
||
|
||
|
||
# ---------- time helpers ----------
|
||
def now_utc() -> str:
|
||
return datetime.now(timezone.utc).strftime("%Y-%m-%d %H:%M:%S UTC")
|
||
def now_ct() -> str:
|
||
if CT:
|
||
return datetime.now(CT).strftime("%Y-%m-%d %I:%M:%S %p CT")
|
||
return time.strftime("%Y-%m-%d %I:%M:%S %p (local)")
|
||
def stamp(label: str) -> str:
|
||
return f"[{label}] {now_ct()} | {now_utc()}"
|
||
|
||
|
||
# ---------- db ----------
|
||
def connect_rw() -> psycopg2.extensions.connection:
|
||
"""Create a PostgreSQL connection using environment variables."""
|
||
use_local_db = os.getenv('USE_LOCAL_DB', 'false').lower() == 'true'
|
||
|
||
if use_local_db:
|
||
config = {
|
||
'host': os.getenv('LOCAL_DB_HOST', '192.168.8.151'),
|
||
'port': int(os.getenv('LOCAL_DB_PORT', '5432')),
|
||
'user': os.getenv('LOCAL_DB_USER', 'postgres'),
|
||
'password': os.getenv('LOCAL_DB_PASSWORD', 'postgres'),
|
||
'database': os.getenv('LOCAL_DB_NAME', 'institutional_trader')
|
||
}
|
||
elif os.getenv('DATABASE_URL'):
|
||
# Parse connection string
|
||
import urllib.parse
|
||
url = urllib.parse.urlparse(os.getenv('DATABASE_URL'))
|
||
config = {
|
||
'host': url.hostname,
|
||
'port': url.port or 5432,
|
||
'user': url.username,
|
||
'password': url.password,
|
||
'database': url.path.lstrip('/')
|
||
}
|
||
else:
|
||
raise ValueError(
|
||
'No database configuration found. Set either:\n'
|
||
' - USE_LOCAL_DB=true with LOCAL_DB_* variables\n'
|
||
' - OR DATABASE_URL'
|
||
)
|
||
|
||
return psycopg2.connect(**config)
|
||
|
||
|
||
def ensure_metadata_tables(con: psycopg2.extensions.connection):
|
||
with con.cursor() as cur:
|
||
cur.execute("""
|
||
CREATE TABLE IF NOT EXISTS universe_symbols (
|
||
symbol VARCHAR(20) PRIMARY KEY,
|
||
source VARCHAR(50),
|
||
added_at TIMESTAMP WITH TIME ZONE
|
||
)""")
|
||
cur.execute("""
|
||
CREATE TABLE IF NOT EXISTS symbol_sectors (
|
||
symbol VARCHAR(20) NOT NULL,
|
||
sector VARCHAR(100) NOT NULL,
|
||
PRIMARY KEY (symbol, sector)
|
||
)""")
|
||
con.commit()
|
||
|
||
def ensure_prices_daily_schema(con: psycopg2.extensions.connection):
|
||
with con.cursor() as cur:
|
||
# Note: PostgreSQL schema uses "Date" (capital D) to match existing schema
|
||
cur.execute("""
|
||
CREATE TABLE IF NOT EXISTS prices_daily (
|
||
symbol VARCHAR(10) NOT NULL,
|
||
"Date" DATE NOT NULL,
|
||
open NUMERIC(12, 4),
|
||
high NUMERIC(12, 4),
|
||
low NUMERIC(12, 4),
|
||
close NUMERIC(12, 4),
|
||
volume NUMERIC(15, 2),
|
||
PRIMARY KEY (symbol, "Date")
|
||
)""")
|
||
cur.execute("""
|
||
CREATE INDEX IF NOT EXISTS idx_prices_daily_symbol_date
|
||
ON prices_daily(symbol, "Date" DESC)
|
||
""")
|
||
con.commit()
|
||
def ensure_intraday_schema(con: psycopg2.extensions.connection):
|
||
with con.cursor() as cur:
|
||
# PostgreSQL schema uses TIMESTAMP WITH TIME ZONE instead of INTEGER
|
||
cur.execute("""
|
||
CREATE TABLE IF NOT EXISTS prices_intraday_1m (
|
||
symbol VARCHAR(10) NOT NULL,
|
||
ts TIMESTAMP WITH TIME ZONE NOT NULL,
|
||
open NUMERIC(12, 4),
|
||
high NUMERIC(12, 4),
|
||
low NUMERIC(12, 4),
|
||
close NUMERIC(12, 4),
|
||
volume NUMERIC(15, 2),
|
||
PRIMARY KEY (symbol, ts)
|
||
)""")
|
||
cur.execute("""
|
||
CREATE INDEX IF NOT EXISTS idx_prices_intraday_symbol_ts
|
||
ON prices_intraday_1m(symbol, ts DESC)
|
||
""")
|
||
cur.execute("""
|
||
CREATE INDEX IF NOT EXISTS idx_prices_intraday_ts
|
||
ON prices_intraday_1m(ts DESC)
|
||
""")
|
||
con.commit()
|
||
|
||
|
||
def get_last_daily_date(con: psycopg2.extensions.connection, sym: str) -> Optional[str]:
|
||
"""Get the last date for a symbol from prices_daily table."""
|
||
with con.cursor() as cur:
|
||
cur.execute('SELECT MAX("Date") FROM prices_daily WHERE symbol=%s', (sym,))
|
||
row = cur.fetchone()
|
||
if row and row[0]:
|
||
# Convert date to string format YYYY-MM-DD
|
||
if isinstance(row[0], date):
|
||
return row[0].isoformat()
|
||
return str(row[0])
|
||
return None
|
||
|
||
|
||
# ---------- symbol universe ----------
|
||
def normalize_yahoo_symbol(s: str) -> str:
|
||
return s.strip().upper().replace(".", "-")
|
||
|
||
def _unique(seq: Iterable[str]) -> List[str]:
|
||
seen, out = set(), []
|
||
for x in seq:
|
||
if x not in seen:
|
||
out.append(x); seen.add(x)
|
||
return out
|
||
|
||
def _parse_first_wikitable_with_columns(html: str, preferred_cols: Tuple[str, ...]) -> List[str]:
|
||
soup = BeautifulSoup(html, "html.parser")
|
||
for t in soup.select("table.wikitable"):
|
||
try:
|
||
df = pd.read_html(StringIO(str(t)))[0]
|
||
except Exception:
|
||
continue
|
||
for pc in preferred_cols:
|
||
if pc in df.columns:
|
||
ser = (df[pc].astype(str)
|
||
.str.replace(r"\[.*?\]", "", regex=True)
|
||
.str.strip())
|
||
out = []
|
||
for raw in ser.tolist():
|
||
up = raw.upper().replace(" ", "").replace("\u200b","")
|
||
if up in BAD_WORDS: # safety against header junk
|
||
continue
|
||
if VALID_TICKER_RE.match(up):
|
||
out.append(normalize_yahoo_symbol(up))
|
||
if out:
|
||
return out
|
||
return []
|
||
|
||
def build_ticker_universe(segments: List[str]) -> Tuple[List[str], dict]:
|
||
# manual group tags
|
||
sector_map = {}
|
||
manual_group_syms = []
|
||
for sector, lst in (EXTRA_GROUPS or {}).items():
|
||
for s in lst:
|
||
sym = normalize_yahoo_symbol(s)
|
||
manual_group_syms.append(sym)
|
||
sector_map.setdefault(sym, set()).add(sector)
|
||
manual_group_syms = _unique(manual_group_syms)
|
||
|
||
manual_flat = [normalize_yahoo_symbol(s) for s in EXTRA_SYMBOLS]
|
||
for s in manual_flat:
|
||
sector_map.setdefault(s, set()).add("MANUAL_LIST")
|
||
|
||
if SYMBOLS_ONLY:
|
||
return _unique(manual_group_syms + manual_flat), sector_map
|
||
|
||
auto_syms: List[str] = []
|
||
if "large" in segments:
|
||
print("Fetching S&P 500 tickers…")
|
||
r = requests.get(URL_SP500, headers=HEADERS, timeout=20); r.raise_for_status()
|
||
spx = _parse_first_wikitable_with_columns(r.text, ("Symbol","Ticker","Ticker symbol"))
|
||
auto_syms += spx
|
||
if TAG_AUTO_SEGMENTS:
|
||
for s in spx: sector_map.setdefault(s, set()).add("AUTO:SP500")
|
||
|
||
if "mid" in segments:
|
||
print("Fetching S&P 400 tickers…")
|
||
r = requests.get(URL_SP400, headers=HEADERS, timeout=20); r.raise_for_status()
|
||
mid = _parse_first_wikitable_with_columns(r.text, ("Symbol","Ticker","Ticker symbol"))
|
||
auto_syms += mid
|
||
if TAG_AUTO_SEGMENTS:
|
||
for s in mid: sector_map.setdefault(s, set()).add("AUTO:SP400")
|
||
|
||
if "nasdaq" in segments:
|
||
print("Fetching Nasdaq-100 tickers…")
|
||
r = requests.get(URL_NASDAQ, headers=HEADERS, timeout=20); r.raise_for_status()
|
||
ndx = _parse_first_wikitable_with_columns(r.text, ("Ticker","Ticker symbol","Symbol"))
|
||
auto_syms += ndx
|
||
if TAG_AUTO_SEGMENTS:
|
||
for s in ndx: sector_map.setdefault(s, set()).add("AUTO:NDX100")
|
||
|
||
return _unique(auto_syms + manual_group_syms + manual_flat), sector_map
|
||
|
||
|
||
# ---------- yfinance helpers ----------
|
||
def yf_download_multi(tickers: List[str], period: str, interval: str) -> pd.DataFrame:
|
||
return yf.download(tickers=tickers, period=period, interval=interval,
|
||
auto_adjust=False, progress=False, threads=True, group_by="ticker")
|
||
|
||
def chunks(seq, n):
|
||
for i in range(0, len(seq), n):
|
||
yield seq[i:i+n]
|
||
|
||
|
||
# ---------- metadata upserts ----------
|
||
def ensure_meta(con: psycopg2.extensions.connection):
|
||
ensure_metadata_tables(con)
|
||
|
||
def upsert_universe(con: psycopg2.extensions.connection, symbols: List[str], sector_map: dict):
|
||
now_iso = datetime.now(timezone.utc)
|
||
rows = []
|
||
for s in symbols:
|
||
tags = sector_map.get(s, set())
|
||
is_manual = ("MANUAL_LIST" in tags) or any(not t.startswith("AUTO:") for t in tags)
|
||
rows.append((s, "manual" if is_manual else "auto", now_iso))
|
||
|
||
with con.cursor() as cur:
|
||
execute_values(cur, """
|
||
INSERT INTO universe_symbols(symbol, source, added_at)
|
||
VALUES %s
|
||
ON CONFLICT (symbol) DO NOTHING
|
||
""", rows)
|
||
|
||
pairs = []
|
||
for s, tags in sector_map.items():
|
||
for t in tags:
|
||
pairs.append((s, t))
|
||
if pairs:
|
||
execute_values(cur, """
|
||
INSERT INTO symbol_sectors(symbol, sector)
|
||
VALUES %s
|
||
ON CONFLICT (symbol, sector) DO NOTHING
|
||
""", pairs)
|
||
con.commit()
|
||
|
||
|
||
# ---------- data upserts ----------
|
||
def upsert_daily(con: psycopg2.extensions.connection, sym: str, df: pd.DataFrame, source="yfinance", diag: bool=False) -> int:
|
||
if df is None or df.empty:
|
||
print(f"[daily upsert] {sym}: EMPTY dataframe from yfinance")
|
||
return 0
|
||
|
||
# Handle single or multi-ticker DataFrame
|
||
if isinstance(df.columns, pd.MultiIndex):
|
||
g = None
|
||
for level in (0, 1):
|
||
try:
|
||
g = df.xs(sym, axis=1, level=level)
|
||
break
|
||
except Exception:
|
||
pass
|
||
if g is None:
|
||
print(f"[daily upsert] {sym}: could not slice MultiIndex for ticker")
|
||
return 0
|
||
else:
|
||
g = df
|
||
|
||
g = g.copy()
|
||
idx = pd.to_datetime(g.index, errors="coerce")
|
||
g = g[~idx.isna()]
|
||
try:
|
||
idx = idx.tz_localize(None)
|
||
except Exception:
|
||
pass
|
||
g.index = idx
|
||
|
||
# Normalize column names to our schema
|
||
g = g.rename(columns={
|
||
"Open":"open", "High":"high", "Low":"low", "Close":"close",
|
||
"Adj Close":"adj_close", "Volume":"volume"
|
||
})
|
||
for c in ["open","high","low","close","volume"]:
|
||
if c not in g.columns:
|
||
g[c] = None
|
||
|
||
# Use "Date" (capital D) to match PostgreSQL schema, convert to date object
|
||
g["Date"] = g.index.date
|
||
g["symbol"] = sym
|
||
|
||
if diag:
|
||
dmin = g["Date"].min() if len(g) else "NA"
|
||
dmax = g["Date"].max() if len(g) else "NA"
|
||
print(f"[daily upsert] {sym}: normalized rows={len(g)} range=({dmin} -> {dmax})")
|
||
try:
|
||
g.reset_index(drop=True).to_csv(f"daily_{sym}.csv", index=False)
|
||
print(f"[diag] wrote preview CSV -> daily_{sym}.csv")
|
||
except Exception as e:
|
||
print(f"[diag] CSV error: {e}")
|
||
|
||
rows = g[["symbol","Date","open","high","low","close","volume"]].values.tolist()
|
||
try:
|
||
with con.cursor() as cur:
|
||
execute_values(cur, """
|
||
INSERT INTO prices_daily(symbol,"Date",open,high,low,close,volume)
|
||
VALUES %s
|
||
ON CONFLICT(symbol,"Date") DO UPDATE SET
|
||
open=EXCLUDED.open, high=EXCLUDED.high, low=EXCLUDED.low,
|
||
close=EXCLUDED.close, volume=EXCLUDED.volume
|
||
""", rows)
|
||
con.commit()
|
||
except Exception as e:
|
||
print(f"[daily upsert] {sym}: postgres error -> {e}")
|
||
con.rollback()
|
||
return 0
|
||
return len(rows)
|
||
|
||
# ---------- single-symbol seeder (CLI) ----------
|
||
def seed_single_daily(con: psycopg2.extensions.connection, sym: str, period: str = "2y", diag: bool=False) -> None:
|
||
print(stamp(f"SEED DAILY {sym} (period={period})"))
|
||
df = yf.download(sym, period=period, interval="1d", auto_adjust=False, progress=False)
|
||
if df is None or df.empty:
|
||
print(f"[SEED] {sym}: yfinance returned EMPTY — cannot seed")
|
||
return
|
||
|
||
idx = pd.to_datetime(df.index, errors="coerce")
|
||
try: idx = idx.tz_localize(None)
|
||
except Exception: pass
|
||
dmin = str(pd.to_datetime(idx.min()).date()) if len(idx) else "NA"
|
||
dmax = str(pd.to_datetime(idx.max()).date()) if len(idx) else "NA"
|
||
print(f"[SEED] {sym}: fetched {len(df)} rows ({dmin} -> {dmax})")
|
||
|
||
wrote = upsert_daily(con, sym, df, diag=diag)
|
||
print(f"[SEED] {sym}: inserted/updated {wrote} rows")
|
||
|
||
with con.cursor() as cur:
|
||
cur.execute("""
|
||
SELECT COUNT(*), MIN("Date"), MAX("Date")
|
||
FROM prices_daily
|
||
WHERE symbol=%s
|
||
""", (sym,))
|
||
row = cur.fetchone()
|
||
print(f"[SEED CHECK] {sym}: count={row[0]} range=({row[1]} -> {row[2]})")
|
||
|
||
|
||
def upsert_intraday(con: psycopg2.extensions.connection, sym: str, df: pd.DataFrame, source="yfinance") -> int:
|
||
if df is None or df.empty:
|
||
return 0
|
||
|
||
# Slice the multi-index (group_by="ticker") or pass-through single frame
|
||
g = df.xs(sym, axis=1, level=0) if isinstance(df.columns, pd.MultiIndex) else df
|
||
if g is None or g.empty:
|
||
return 0
|
||
|
||
g = g.rename(columns={"Open": "open", "High": "high", "Low": "low", "Close": "close", "Volume": "volume"}).copy()
|
||
|
||
# Ensure UTC time index
|
||
idx = pd.to_datetime(g.index, errors="coerce")
|
||
if idx.tz is None:
|
||
idx = idx.tz_localize("UTC", nonexistent="shift_forward", ambiguous="NaT")
|
||
else:
|
||
idx = idx.tz_convert("UTC")
|
||
|
||
# Convert to pandas Timestamp for PostgreSQL TIMESTAMP WITH TIME ZONE
|
||
# PostgreSQL will handle the timezone conversion
|
||
g["symbol"] = sym.upper()
|
||
g["ts"] = idx
|
||
|
||
# Order matches table columns (no date_cst, min_cst, source in PostgreSQL schema)
|
||
rows = g[["symbol","ts","open","high","low","close","volume"]].values.tolist()
|
||
|
||
# Use PostgreSQL ON CONFLICT syntax
|
||
try:
|
||
with con.cursor() as cur:
|
||
execute_values(cur, """
|
||
INSERT INTO prices_intraday_1m
|
||
(symbol, ts, open, high, low, close, volume)
|
||
VALUES %s
|
||
ON CONFLICT (symbol, ts) DO UPDATE SET
|
||
open=EXCLUDED.open, high=EXCLUDED.high, low=EXCLUDED.low,
|
||
close=EXCLUDED.close, volume=EXCLUDED.volume
|
||
""", rows)
|
||
con.commit()
|
||
except Exception as e:
|
||
print(f"[intraday upsert] {sym}: postgres error -> {e}")
|
||
con.rollback()
|
||
return 0
|
||
|
||
return len(g)
|
||
|
||
|
||
|
||
# ---------- runners ----------
|
||
def run_intraday_batched(con, symbols, lookback_days, batch_size, sleep_between):
|
||
# Determine interval based on lookback days (Yahoo Finance limitations)
|
||
# 1m: Last 7 days maximum
|
||
# 5m: Last 60 days maximum
|
||
# 15m: Last 60 days maximum
|
||
if lookback_days <= 7:
|
||
interval = "1m"
|
||
period = f"{lookback_days}d"
|
||
elif lookback_days <= 60:
|
||
interval = "5m" # 5-minute bars for longer lookback (30 days)
|
||
period = f"{lookback_days}d"
|
||
else:
|
||
interval = "15m" # 15-minute for very long lookback
|
||
period = "60d" # Max 60 days for 15m
|
||
|
||
total_rows, errors = 0, 0
|
||
print(stamp(f"INTRADAY START last {lookback_days}d using {interval} interval"))
|
||
print(f"Universe size: {len(symbols)}; batch={batch_size}", flush=True)
|
||
|
||
for bi, batch in enumerate(chunks(symbols, batch_size), 1):
|
||
tag = f"[batch {bi}/{(len(symbols)+batch_size-1)//batch_size}]"
|
||
print(stamp(f"{tag} download {len(batch)} tickers"), flush=True)
|
||
try:
|
||
df = yf_download_multi(batch, period=period, interval=interval)
|
||
except KeyboardInterrupt:
|
||
print("\n[interrupt] stopping intraday loop gracefully."); break
|
||
except Exception as e:
|
||
errors += len(batch); print(f"{tag} download error: {e}")
|
||
time.sleep(max(0.2, sleep_between)); continue
|
||
|
||
try:
|
||
# Process batch - each upsert commits individually
|
||
for sym in batch:
|
||
try:
|
||
total_rows += upsert_intraday(con, sym, df)
|
||
except Exception as e:
|
||
errors += 1; print(f"{tag} upsert error {sym}: {e}")
|
||
except KeyboardInterrupt:
|
||
print("\n[interrupt] stopping mid-batch."); break
|
||
|
||
print(stamp(f"{tag} done; cumulative rows={total_rows}, errors={errors}"), flush=True)
|
||
time.sleep(max(0.2, sleep_between))
|
||
|
||
print(stamp("INTRADAY END"))
|
||
print(f"[INTRADAY] rows upserted: {total_rows}, errors: {errors}")
|
||
return total_rows, errors
|
||
|
||
|
||
def _daily_count(con: psycopg2.extensions.connection, sym: str) -> int:
|
||
with con.cursor() as cur:
|
||
cur.execute("SELECT COUNT(*) FROM prices_daily WHERE symbol=%s", (sym,))
|
||
row = cur.fetchone()
|
||
return int(row[0] or 0)
|
||
|
||
def run_daily_incremental(con: psycopg2.extensions.connection, symbols: List[str], period: str) -> Tuple[int,int]:
|
||
print(stamp(f"DAILY START (period={period})"))
|
||
total, errors = 0, 0
|
||
for i, sym in enumerate(symbols, 1):
|
||
try:
|
||
have = _daily_count(con, sym)
|
||
need_seed = have < 10
|
||
last = None if need_seed else get_last_daily_date(con, sym)
|
||
|
||
if last is None:
|
||
mode = "FULL"
|
||
df = yf.download(sym, period=period, interval="1d",
|
||
auto_adjust=False, progress=False)
|
||
else:
|
||
mode = f"INC from {last} -5d"
|
||
start_dt = (datetime.strptime(last, "%Y-%m-%d") - timedelta(days=5)).date().isoformat()
|
||
df = yf.download(sym, start=start_dt, interval="1d",
|
||
auto_adjust=False, progress=False)
|
||
|
||
if df is None or df.empty:
|
||
print(f"[DAILY {mode}] {sym}: yfinance returned EMPTY (have={have})")
|
||
else:
|
||
i0 = pd.to_datetime(df.index, errors="coerce")
|
||
try: i0 = i0.tz_localize(None)
|
||
except Exception: pass
|
||
dmin = str(pd.to_datetime(i0.min()).date()) if len(i0) else "NA"
|
||
dmax = str(pd.to_datetime(i0.max()).date()) if len(i0) else "NA"
|
||
print(f"[DAILY {mode}] {sym}: fetched {len(df)} rows ({dmin} -> {dmax}), have={have}")
|
||
|
||
wrote = upsert_daily(con, sym, df)
|
||
total += wrote
|
||
|
||
if i % 50 == 0:
|
||
print(stamp(f"DAILY progress {i}/{len(symbols)} (rows={total})"), flush=True)
|
||
|
||
except Exception as e:
|
||
errors += 1
|
||
print(f"[daily error] {sym}: {e}")
|
||
time.sleep(0.02)
|
||
# Note: upsert_daily commits individually, no need for final commit
|
||
print(stamp("DAILY END"))
|
||
print(f"[DAILY] rows upserted: {total}, errors: {errors}")
|
||
return total, errors
|
||
|
||
|
||
# ---------- single-symbol seeder (CLI) ----------
|
||
# ---------- main loop ----------
|
||
def main():
|
||
db_info = os.getenv('DATABASE_URL', 'local PostgreSQL')
|
||
print(stamp(f"BOOT — PostgreSQL: {db_info}"))
|
||
con = connect_rw()
|
||
try:
|
||
ensure_prices_daily_schema(con)
|
||
ensure_intraday_schema(con)
|
||
|
||
symbols, sector_map = build_ticker_universe(SEGMENTS)
|
||
ensure_meta(con)
|
||
upsert_universe(con, symbols, sector_map)
|
||
print(f"Total tickers in universe: {len(symbols)}")
|
||
|
||
if not LOOP:
|
||
run_intraday_batched(con, symbols, LOOKBACK_DAYS_INTRADAY, BATCH_SIZE, SLEEP_BETWEEN_BATCHES)
|
||
run_daily_incremental(con, symbols, DAILY_PERIOD)
|
||
print(stamp("ONE-SHOT DONE"))
|
||
return 0
|
||
|
||
last_daily = 0.0
|
||
cycle = 0
|
||
while True:
|
||
cycle += 1
|
||
print("\n" + "="*7 + f" LOOP cycle {cycle} " + "="*7)
|
||
print(stamp("LOOP START"))
|
||
|
||
run_intraday_batched(con, symbols, LOOKBACK_DAYS_INTRADAY, BATCH_SIZE, SLEEP_BETWEEN_BATCHES)
|
||
|
||
now = time.time()
|
||
if DAILY_INTERVAL_MIN == 0 or now - last_daily >= DAILY_INTERVAL_MIN * 60:
|
||
run_daily_incremental(con, symbols, DAILY_PERIOD)
|
||
last_daily = time.time()
|
||
|
||
if INTRADAY_INTERVAL_MIN <= 0:
|
||
time.sleep(1.0)
|
||
else:
|
||
interval = INTRADAY_INTERVAL_MIN * 60
|
||
now = time.time()
|
||
next_ts = (int(now // interval) + 1) * interval
|
||
jitter = random.uniform(0, 4)
|
||
delay = max(0.0, (next_ts - now) + jitter)
|
||
print(stamp(f"SLEEP {delay:.1f}s until next intraday run"))
|
||
time.sleep(delay)
|
||
|
||
finally:
|
||
con.close()
|
||
print(stamp("CLOSED DB"))
|
||
|
||
|
||
# ---------- CLI entry ----------
|
||
def cli_entry():
|
||
import argparse
|
||
parser = argparse.ArgumentParser(description="Load prices to PostgreSQL (daily + intraday)")
|
||
parser.add_argument("--only", help="Run only for one symbol (e.g. --only ASTS)")
|
||
parser.add_argument("--daily", action="store_true", help="Run only daily for the symbol (with --only)")
|
||
parser.add_argument("--intraday", action="store_true", help="Run only intraday for the symbol (with --only)")
|
||
parser.add_argument("--diag", action="store_true", help="Diagnostic: print ranges and write CSV previews")
|
||
args = parser.parse_args()
|
||
|
||
db_info = os.getenv('DATABASE_URL', 'local PostgreSQL')
|
||
print(stamp(f"BOOT — PostgreSQL: {db_info}"))
|
||
con = connect_rw()
|
||
ensure_prices_daily_schema(con)
|
||
ensure_intraday_schema(con)
|
||
|
||
if args.only:
|
||
sym = args.only.strip().upper()
|
||
if args.daily:
|
||
seed_single_daily(con, sym, DAILY_PERIOD, diag=args.diag)
|
||
elif args.intraday:
|
||
# diag not needed for intraday path
|
||
run_intraday_batched(con, [sym], LOOKBACK_DAYS_INTRADAY, 1, 0.1)
|
||
else:
|
||
seed_single_daily(con, sym, DAILY_PERIOD, diag=args.diag)
|
||
run_intraday_batched(con, [sym], LOOKBACK_DAYS_INTRADAY, 1, 0.1)
|
||
con.close()
|
||
print(stamp("DONE (CLI one-shot)"))
|
||
sys.exit(0)
|
||
|
||
# Default = full loop
|
||
rc = main()
|
||
sys.exit(rc if isinstance(rc, int) else 0)
|
||
|
||
|
||
if __name__ == "__main__":
|
||
cli_entry() |