734 lines
32 KiB
Python
734 lines
32 KiB
Python
# sync_blackbox_flow.py — BlackBox API sync to PostgreSQL
|
|
# Fetches options flow data from BlackBox API and syncs to PostgreSQL database
|
|
|
|
import os, time, hashlib, re, argparse
|
|
from pathlib import Path
|
|
from typing import Dict, Optional, List
|
|
from datetime import datetime, date
|
|
import pandas as pd
|
|
import numpy as np
|
|
import requests
|
|
import json
|
|
import psycopg2
|
|
from psycopg2.extras import execute_values
|
|
from psycopg2 import sql
|
|
from dotenv import load_dotenv
|
|
|
|
load_dotenv()
|
|
|
|
# ─────────────────────────────────────────────
|
|
# Config
|
|
# ─────────────────────────────────────────────
|
|
WRITE_POSTGRES = True
|
|
|
|
# BlackBox API Configuration
|
|
BLACKBOX_API_URL = "https://api.blackboxstocks.com/api/v2/options/getFlowMobile"
|
|
BLACKBOX_API_TOKEN = os.getenv("BLACKBOX_API_TOKEN", "eyJhbGciOiJodHRwOi8vd3d3LnczLm9yZy8yMDAxLzA0L3htbGRzaWctbW9yZSNobWFjLXNoYTI1NiIsInR5cCI6IkpXVCJ9.eyJodHRwOi8vc2NoZW1hcy54bWxzb2FwLm9yZy93cy8yMDA1LzA1L2lkZW50aXR5L2NsYWltcy9uYW1lIjoic3JlZGR5MDgxOUBnbWFpbC5jb20iLCJodHRwOi8vc2NoZW1hcy54bWxzb2FwLm9yZy93cy8yMDA1LzA1L2lkZW50aXR5L2NsYWltcy9uYW1laWRlbnRpZmllciI6Ijk3ZDA5NThjLTBhOTktNDA5OC05OWQwLWNkYTg3Njg5N2Q3ZiIsImh0dHA6Ly9zY2hlbWFzLnhtbHNvYXAub3JnL3dzLzIwMDUvMDUvaWRlbnRpdHkvY2xhaW1zL3NpZCI6InBpS1plMktZN2d2NnErbDZyM3g3YkpHYU9ESlUrZVdOZmMwc1NicWZZNU09IiwiaHR0cDovL3NjaGVtYXMubWljcm9zb2Z0LmNvbS93cy8yMDA4LzA2L2lkZW50aXR5L2NsYWltcy9yb2xlIjpbIlN1YnNjcmliZXIiLCJPcHRpb25zIiwiRXF1aXRpZXMiXSwiZXhwIjoxNzY1NzQ4NjI1fQ.Evq__DugD7s1kAytUqtsknQFIeOPjKM3iwp6cDI0hJI")
|
|
|
|
# PostgreSQL connection parameters
|
|
POSTGRES_HOST = os.getenv("POSTGRES_HOST", "localhost")
|
|
POSTGRES_PORT = int(os.getenv("POSTGRES_PORT", "5432"))
|
|
POSTGRES_DB = os.getenv("POSTGRES_DB", "institutional_trader")
|
|
POSTGRES_USER = os.getenv("POSTGRES_USER", "postgres")
|
|
POSTGRES_PASSWORD = os.getenv("POSTGRES_PASSWORD", "postgres")
|
|
POSTGRES_SCHEMA = os.getenv("POSTGRES_SCHEMA", "public")
|
|
|
|
# ⚠️ Per-table schema style: 'snake' or 'camel'
|
|
TABLE_STYLE = {
|
|
"AlertStream": "snake",
|
|
"AlertStream_monthly": "snake",
|
|
"OptionsFlow": "camel", # ← your Supabase columns are CamelCase here
|
|
"OptionsFlow_monthly": "camel", # ← same
|
|
"OptionsVolume": "camel", # adjust if needed
|
|
"Short_Long": "camel", # adjust if needed
|
|
}
|
|
|
|
_TARGET_STEMS = set(TABLE_STYLE.keys())
|
|
|
|
# Upsert behavior / logging
|
|
INCREMENTAL = True
|
|
SINGLE_CHUNK = False # chunked to avoid 57014
|
|
UPSERT_CHUNK = 3000 # drop to 1000 if still timeouts
|
|
MAX_RETRIES = 3
|
|
PRINT_PROGRESS = False
|
|
PRINT_PRUNE_LOGS = False
|
|
USE_RETURNING_MINIMAL= True
|
|
|
|
print_flush = lambda *a, **k: print(*a, **k, flush=True)
|
|
def _ts(): return time.strftime("%Y-%m-%d %H:%M:%S")
|
|
|
|
# ─────────────────────────────────────────────
|
|
# PostgreSQL connection
|
|
# ─────────────────────────────────────────────
|
|
def _pg_conn():
|
|
"""Create a PostgreSQL connection."""
|
|
return psycopg2.connect(
|
|
host=POSTGRES_HOST,
|
|
port=POSTGRES_PORT,
|
|
database=POSTGRES_DB,
|
|
user=POSTGRES_USER,
|
|
password=POSTGRES_PASSWORD,
|
|
options=f"-c search_path={POSTGRES_SCHEMA}"
|
|
)
|
|
|
|
# ─────────────────────────────────────────────
|
|
# Expected headers (both styles) + mappings
|
|
# ─────────────────────────────────────────────
|
|
# OptionsFlow (CamelCase)
|
|
EXPECTED_OPTIONSFLOW_CAMEL = [
|
|
"CreatedDate","CreatedTime","Symbol","Type","Volume","Price","Side",
|
|
"CallPut","Strike","Spot","Premium","ExpirationDate","Color",
|
|
"ImpliedVolatility","Dte","ER","StockEtf","Sector","Uoa",
|
|
"Weekly","MktCap","OI"
|
|
]
|
|
# OptionsFlow (snake_case)
|
|
EXPECTED_OPTIONSFLOW_SNAKE = [
|
|
"created_date","created_time","symbol","type","volume","price","side",
|
|
"callput","strike","spot","premium","expiration_date","color",
|
|
"implied_volatility","dte","er","stock_etf","sector","uoa",
|
|
"weekly","mktcap","oi"
|
|
]
|
|
# AlertStream (CamelCase in files → snake in DB)
|
|
EXPECTED_ALERTSTREAM_SNAKE = [
|
|
"date","timestamp","ticker","volume","price","pct_of_avg30",
|
|
"notional","message","type","securitytype","industry","sector",
|
|
"avg30day","float","earningsdate"
|
|
]
|
|
|
|
# Camel→snake renames (if file arrives CamelCase but DB expects snake)
|
|
MAP_OPTIONSFLOW_SNAKE = {
|
|
"CreatedDate":"created_date","CreatedTime":"created_time","Symbol":"symbol","Type":"type",
|
|
"Volume":"volume","Price":"price","Side":"side","CallPut":"callput","Strike":"strike",
|
|
"Spot":"spot","Premium":"premium","ExpirationDate":"expiration_date","Color":"color",
|
|
"ImpliedVolatility":"implied_volatility","Dte":"dte","ER":"er","StockEtf":"stock_etf",
|
|
"Sector":"sector","Uoa":"uoa","Weekly":"weekly","MktCap":"mktcap","OI":"oi"
|
|
}
|
|
MAP_ALERTSTREAM_SNAKE = {
|
|
"Date":"date","Timestamp":"timestamp","Ticker":"ticker","Volume":"volume","Price":"price",
|
|
"Pct_of_Avg30Day":"pct_of_avg30","Notional":"notional","Message":"message","Type":"type",
|
|
"SecurityType":"securitytype","Industry":"industry","Sector":"sector",
|
|
"Avg30Day":"avg30day","Float":"float","EarningsDate":"earningsdate"
|
|
}
|
|
|
|
# snake→Camel renames (if file is snake but DB expects Camel)
|
|
MAP_OPTIONSFLOW_CAMEL = {v:k for k,v in MAP_OPTIONSFLOW_SNAKE.items()}
|
|
|
|
# ─────────────────────────────────────────────
|
|
# BlackBox API Integration
|
|
# ─────────────────────────────────────────────
|
|
def build_filter_bitmask(_filter_options: Optional[Dict] = None) -> int:
|
|
"""Build the filter bitmask based on filter options."""
|
|
# Default filter value that enables common filters
|
|
return 2198487171967
|
|
|
|
def build_filters(start_date: datetime, end_date: datetime, custom_filters: Optional[Dict] = None) -> Dict:
|
|
"""Build default filters object."""
|
|
date_str = start_date.isoformat()
|
|
|
|
filters = {
|
|
"optionsDate": {
|
|
"start": date_str,
|
|
"end": end_date.isoformat()
|
|
},
|
|
"expireOptionsDate": {
|
|
"start": date_str,
|
|
"end": end_date.isoformat()
|
|
},
|
|
"optionsFlowPuts": True,
|
|
"optionsFlowCalls": True,
|
|
"optionsFlowYellow": True,
|
|
"optionsFlowWhite": True,
|
|
"optionsFlowMagenta": True,
|
|
"optionsFlowAboveAskOnly": True,
|
|
"optionsFlowBelowBidOnly": True,
|
|
"optionsFlowAtOrAboveAsk": True,
|
|
"optionsFlowAtOrBelowBid": True,
|
|
"optionsFlowMultileg": False,
|
|
"optionsFlowOnlyMultiLeg": False,
|
|
"optionsFlowBelowPoint5": False,
|
|
"optionsFlowBelow5": False,
|
|
"optionsFlow100Contracts": False,
|
|
"optionsFlow500Contracts": False,
|
|
"optionsFlow5000Contracts": False,
|
|
"optionsFlowStock": True,
|
|
"optionsFlowEtf": True,
|
|
"optionsFlowAbove50k": False,
|
|
"optionsFlowAbove100k": False,
|
|
"optionsFlowAbove200k": False,
|
|
"optionsFlowAbove500k": False,
|
|
"optionsFlowAbove1m": False,
|
|
"marketCapAbove750B": False,
|
|
"optionsFlowInTheMoney": False,
|
|
"optionsFlowOutOfTheMoney": False,
|
|
"optionsFlowSweepOnly": False,
|
|
"optionsFlowWeeklyOnly": False,
|
|
"optionsFlowEarningsReportOnly": False,
|
|
"optionsFlowUnusualOnly": False,
|
|
"optionsFlowExDiv": False,
|
|
"optionsFlowConsumerDiscretionary": True,
|
|
"optionsFlowIndustrials": True,
|
|
"optionsFlowInformationTechnology": True,
|
|
"optionsFlowRealEstate": True,
|
|
"optionsFlowHealthCare": True,
|
|
"optionsFlowEnergy": True,
|
|
"optionsFlowFinancials": True,
|
|
"optionsFlowMaterials": True,
|
|
"optionsFlowConsumerStaples": True,
|
|
"optionsFlowCommunicationServices": True,
|
|
"optionsFlowUtilities": True,
|
|
"optionsExpirationRange": False,
|
|
"optionsFlowSectorNone": True,
|
|
}
|
|
|
|
if custom_filters:
|
|
filters.update(custom_filters)
|
|
|
|
return filters
|
|
|
|
# Constants
|
|
TIMEZONE_SUFFIX = "+00:00"
|
|
|
|
def fetch_blackbox_flow(options: Optional[Dict] = None) -> List[Dict]:
|
|
"""Fetch options flow data from BlackBox Stocks API."""
|
|
if options is None:
|
|
options = {}
|
|
|
|
if not BLACKBOX_API_TOKEN:
|
|
raise ValueError(
|
|
"BLACKBOX_API_TOKEN not found in environment variables.\n"
|
|
"Please add BLACKBOX_API_TOKEN to your .env file or set it as an environment variable."
|
|
)
|
|
|
|
# Parse dates - default to today if not provided
|
|
if options.get("startDate"):
|
|
if isinstance(options["startDate"], str):
|
|
start_date = datetime.fromisoformat(options["startDate"].replace("Z", TIMEZONE_SUFFIX))
|
|
elif isinstance(options["startDate"], date):
|
|
start_date = datetime.combine(options["startDate"], datetime.min.time())
|
|
else:
|
|
start_date = options["startDate"]
|
|
else:
|
|
start_date = datetime.now()
|
|
|
|
if options.get("endDate"):
|
|
if isinstance(options["endDate"], str):
|
|
end_date = datetime.fromisoformat(options["endDate"].replace("Z", TIMEZONE_SUFFIX))
|
|
elif isinstance(options["endDate"], date):
|
|
end_date = datetime.combine(options["endDate"], datetime.max.time())
|
|
else:
|
|
end_date = options["endDate"]
|
|
else:
|
|
end_date = start_date
|
|
|
|
# Build request body
|
|
body = {
|
|
"historical": options.get("historical", False),
|
|
"symbol": options.get("symbol", ""),
|
|
"strike": options.get("strike", 0),
|
|
"count": options.get("count") or options.get("limit", 300),
|
|
"filter": build_filter_bitmask(options.get("filters")),
|
|
"filters": build_filters(start_date, end_date, options.get("filters")),
|
|
"fromDate": start_date.isoformat(),
|
|
"toDate": end_date.isoformat()
|
|
}
|
|
|
|
try:
|
|
response = requests.post(
|
|
BLACKBOX_API_URL,
|
|
headers={
|
|
"Content-Type": "application/json",
|
|
"Accept": "application/json",
|
|
"Authorization": f"Bearer {BLACKBOX_API_TOKEN}"
|
|
},
|
|
json=body,
|
|
timeout=30
|
|
)
|
|
|
|
if not response.ok:
|
|
error_text = response.text
|
|
raise RuntimeError(
|
|
f"BlackBox API error: {response.status_code} {response.reason}\n"
|
|
f"Response: {error_text}"
|
|
)
|
|
|
|
data = response.json()
|
|
|
|
# Handle different response structures
|
|
if isinstance(data, list):
|
|
return data
|
|
elif isinstance(data, dict):
|
|
if "data" in data and isinstance(data["data"], list):
|
|
return data["data"]
|
|
elif "flows" in data and isinstance(data["flows"], list):
|
|
return data["flows"]
|
|
elif "results" in data and isinstance(data["results"], list):
|
|
return data["results"]
|
|
else:
|
|
print_flush(f"{_ts()} | ⚠️ Unexpected API response structure: {list(data.keys())}")
|
|
return []
|
|
else:
|
|
return []
|
|
except Exception as e:
|
|
print_flush(f"{_ts()} | ❌ Error fetching BlackBox flow data: {e}")
|
|
raise
|
|
|
|
def map_blackbox_to_database(api_record: Dict) -> Dict:
|
|
"""Map BlackBox API response to database schema."""
|
|
def get_value(obj: Dict, *keys: str) -> Optional[str]:
|
|
"""Safely extract values from object."""
|
|
for key in keys:
|
|
if key in obj and obj[key] is not None:
|
|
return str(obj[key])
|
|
return None
|
|
|
|
def format_date(date_str: Optional[str]) -> Optional[str]:
|
|
"""Format date as YYYY-MM-DD."""
|
|
if not date_str:
|
|
return None
|
|
try:
|
|
dt = datetime.fromisoformat(str(date_str).replace("Z", TIMEZONE_SUFFIX))
|
|
return dt.strftime("%Y-%m-%d")
|
|
except (ValueError, AttributeError):
|
|
try:
|
|
dt = datetime.strptime(str(date_str), "%Y-%m-%d")
|
|
return dt.strftime("%Y-%m-%d")
|
|
except (ValueError, AttributeError):
|
|
return str(date_str)
|
|
|
|
def format_time(time_str: Optional[str]) -> Optional[str]:
|
|
"""Format time string."""
|
|
if not time_str:
|
|
return None
|
|
return str(time_str)
|
|
|
|
# Map fields - try multiple possible field names from API
|
|
mapped = {
|
|
"CreatedDate": format_date(
|
|
get_value(api_record, "createdDate", "CreatedDate", "date", "Date", "timestamp", "Timestamp")
|
|
),
|
|
"CreatedTime": format_time(
|
|
get_value(api_record, "createdTime", "CreatedTime", "time", "Time", "timestamp", "Timestamp")
|
|
),
|
|
"Symbol": get_value(api_record, "symbol", "Symbol", "ticker", "Ticker", "underlying", "Underlying"),
|
|
"Type": get_value(api_record, "type", "Type", "tradeType", "TradeType"),
|
|
"Volume": get_value(api_record, "volume", "Volume", "vol", "Vol", "contracts", "Contracts"),
|
|
"Price": get_value(api_record, "price", "Price", "lastPrice", "LastPrice", "tradePrice", "TradePrice"),
|
|
"Side": get_value(api_record, "side", "Side", "tradeSide", "TradeSide", "direction", "Direction"),
|
|
"CallPut": get_value(api_record, "callPut", "CallPut", "optionType", "OptionType", "putCall", "PutCall", "type", "Type"),
|
|
"Strike": get_value(api_record, "strike", "Strike", "strikePrice", "StrikePrice"),
|
|
"Spot": get_value(api_record, "spot", "Spot", "underlyingPrice", "UnderlyingPrice", "stockPrice", "StockPrice"),
|
|
"Premium": get_value(api_record, "premium", "Premium", "totalPremium", "TotalPremium", "notional", "Notional"),
|
|
"ExpirationDate": format_date(
|
|
get_value(api_record, "expirationDate", "ExpirationDate", "expiry", "Expiry", "expiration", "Expiration")
|
|
),
|
|
"Color": get_value(api_record, "color", "Color", "tradeColor", "TradeColor"),
|
|
"ImpliedVolatility": get_value(api_record, "impliedVolatility", "ImpliedVolatility", "iv", "IV", "volatility", "Volatility"),
|
|
"Dte": get_value(api_record, "dte", "Dte", "DTE", "daysToExpiration", "DaysToExpiration", "daysToExpiry", "DaysToExpiry"),
|
|
"ER": get_value(api_record, "er", "ER", "earnings", "Earnings", "earningsReport", "EarningsReport"),
|
|
"StockEtf": get_value(api_record, "stockEtf", "StockEtf", "assetType", "AssetType", "securityType", "SecurityType"),
|
|
"Sector": get_value(api_record, "sector", "Sector", "industry", "Industry"),
|
|
"Uoa": get_value(api_record, "uoa", "Uoa", "UOA", "underlyingOfAsset", "UnderlyingOfAsset"),
|
|
"Weekly": get_value(api_record, "weekly", "Weekly", "isWeekly", "IsWeekly", "weeklies", "Weeklies"),
|
|
"MktCap": get_value(api_record, "mktCap", "MktCap", "marketCap", "MarketCap", "marketCapitalization", "MarketCapitalization"),
|
|
"OI": get_value(api_record, "oi", "OI", "openInterest", "OpenInterest", "openInt", "OpenInt")
|
|
}
|
|
|
|
return mapped
|
|
|
|
# ─────────────────────────────────────────────
|
|
# Normalization + JSON safety
|
|
# ─────────────────────────────────────────────
|
|
WEIRD_STR = {"inf","+inf","-inf","infinity","+infinity","-infinity","∞","+∞","-∞",
|
|
"nan","-nan","NaN","N/A","NA","NULL","null",""}
|
|
|
|
def _coerce_weird_numbers(df: pd.DataFrame) -> pd.DataFrame:
|
|
df = df.copy()
|
|
for c in df.columns:
|
|
if df[c].dtype == object:
|
|
df[c] = df[c].replace(list(WEIRD_STR), np.nan)
|
|
return df
|
|
|
|
def _normalize_for_table(df: pd.DataFrame, table: str) -> pd.DataFrame:
|
|
"""Rename/select columns to match the DB style of this table."""
|
|
style = TABLE_STYLE.get(table, "snake").lower()
|
|
df = df.copy()
|
|
df.columns = [c.strip() for c in df.columns]
|
|
|
|
if table in ("OptionsFlow","OptionsFlow_monthly"):
|
|
if style == "camel":
|
|
# Ensure CamelCase headers, no renaming needed if file already CamelCase
|
|
# If file is snake, map to Camel
|
|
lower_cols = {c for c in df.columns if c.islower()}
|
|
if lower_cols:
|
|
df = df.rename(columns=MAP_OPTIONSFLOW_CAMEL)
|
|
for c in EXPECTED_OPTIONSFLOW_CAMEL:
|
|
if c not in df.columns: df[c] = None
|
|
df = df[EXPECTED_OPTIONSFLOW_CAMEL]
|
|
else:
|
|
# snake_case target
|
|
# If file CamelCase, map to snake
|
|
has_upper = any(any(ch.isupper() for ch in c) for c in df.columns)
|
|
if has_upper:
|
|
df = df.rename(columns=MAP_OPTIONSFLOW_SNAKE)
|
|
for c in EXPECTED_OPTIONSFLOW_SNAKE:
|
|
if c not in df.columns: df[c] = None
|
|
df = df[EXPECTED_OPTIONSFLOW_SNAKE]
|
|
|
|
elif table in ("AlertStream","AlertStream_monthly"):
|
|
# DB is snake_case per your screenshot
|
|
# If file CamelCase, map to snake
|
|
has_upper = any(any(ch.isupper() for ch in c) for c in df.columns)
|
|
if has_upper:
|
|
df = df.rename(columns=MAP_ALERTSTREAM_SNAKE)
|
|
for c in EXPECTED_ALERTSTREAM_SNAKE:
|
|
if c not in df.columns: df[c] = None
|
|
df = df[EXPECTED_ALERTSTREAM_SNAKE]
|
|
|
|
# Standardize any *date columns → YYYY-MM-DD* strings
|
|
for col in [c for c in df.columns if c.lower().endswith("date")]:
|
|
df[col] = pd.to_datetime(df[col], errors="coerce").dt.strftime("%Y-%m-%d")
|
|
|
|
return df
|
|
|
|
def _json_safe(df: pd.DataFrame) -> pd.DataFrame:
|
|
df = df.copy()
|
|
# numeric: drop non-finite
|
|
for c in df.columns:
|
|
if pd.api.types.is_numeric_dtype(df[c]):
|
|
s = pd.to_numeric(df[c], errors="coerce")
|
|
s[~np.isfinite(s)] = np.nan
|
|
df[c] = s
|
|
# datetimes -> strings
|
|
for c in df.columns:
|
|
if pd.api.types.is_datetime64_any_dtype(df[c]):
|
|
df[c] = df[c].dt.strftime("%Y-%m-%d %H:%M:%S")
|
|
# NA -> None
|
|
return df.astype(object).where(pd.notnull(df), None)
|
|
|
|
def _row_hash_from_series(s: pd.Series) -> str:
|
|
vals=[]
|
|
for _,v in s.items():
|
|
if v is None or (isinstance(v,float) and pd.isna(v)): vals.append("NULL")
|
|
elif isinstance(v,str): vals.append(v.strip())
|
|
else: vals.append(str(v))
|
|
return hashlib.sha1("\x1f".join(vals).encode("utf-8","ignore")).hexdigest()
|
|
|
|
def _df_prepare_for_postgres(df: pd.DataFrame) -> pd.DataFrame:
|
|
if df.empty: return df
|
|
df = _json_safe(df)
|
|
df["row_hash"] = df.apply(_row_hash_from_series, axis=1)
|
|
return df.drop_duplicates(subset=["row_hash"], keep="first").reset_index(drop=True)
|
|
|
|
def _extract_missing_col_from_error(msg: str) -> Optional[str]:
|
|
"""Extract missing column name from PostgreSQL error messages."""
|
|
patterns = [
|
|
r"column \"([^\"]+)\" does not exist",
|
|
r"Could not find the '([^']+)' column",
|
|
]
|
|
for pattern in patterns:
|
|
m = re.search(pattern, msg, re.IGNORECASE)
|
|
if m:
|
|
return m.group(1)
|
|
return None
|
|
|
|
# ─────────────────────────────────────────────
|
|
# PostgreSQL upload (quiet, chunked)
|
|
# ─────────────────────────────────────────────
|
|
def _ensure_table_exists(conn, table_name: str, columns: list):
|
|
"""Ensure table exists with row_hash column and unique constraint."""
|
|
cur = conn.cursor()
|
|
try:
|
|
# Check if table exists
|
|
cur.execute("""
|
|
SELECT EXISTS (
|
|
SELECT FROM information_schema.tables
|
|
WHERE table_schema = current_schema()
|
|
AND table_name = %s
|
|
)
|
|
""", (table_name,))
|
|
exists = cur.fetchone()[0]
|
|
|
|
if not exists:
|
|
# Create table with all columns as text initially (we'll let PostgreSQL infer types)
|
|
# For now, we'll create a basic structure - actual schema should match your data
|
|
col_defs = ", ".join([f'"{col}" TEXT' for col in columns if col != "row_hash"])
|
|
col_defs += ', "row_hash" TEXT UNIQUE'
|
|
|
|
cur.execute(f'CREATE TABLE IF NOT EXISTS "{table_name}" ({col_defs})')
|
|
conn.commit()
|
|
else:
|
|
# Ensure row_hash column and unique constraint exist
|
|
cur.execute("""
|
|
SELECT column_name FROM information_schema.columns
|
|
WHERE table_schema = current_schema()
|
|
AND table_name = %s AND column_name = 'row_hash'
|
|
""", (table_name,))
|
|
if not cur.fetchone():
|
|
cur.execute(f'ALTER TABLE "{table_name}" ADD COLUMN IF NOT EXISTS "row_hash" TEXT')
|
|
conn.commit()
|
|
|
|
# Check for unique constraint on row_hash
|
|
cur.execute("""
|
|
SELECT constraint_name FROM information_schema.table_constraints
|
|
WHERE table_schema = current_schema()
|
|
AND table_name = %s
|
|
AND constraint_type = 'UNIQUE'
|
|
AND constraint_name LIKE %s
|
|
""", (table_name, f'%{table_name}%row_hash%'))
|
|
if not cur.fetchone():
|
|
try:
|
|
cur.execute(f'CREATE UNIQUE INDEX IF NOT EXISTS "{table_name}_row_hash_idx" ON "{table_name}" ("row_hash")')
|
|
conn.commit()
|
|
except Exception:
|
|
conn.rollback()
|
|
finally:
|
|
cur.close()
|
|
|
|
def _upsert_slice(conn, tname: str, rows: list, columns: list):
|
|
"""Upsert a slice of rows using PostgreSQL INSERT ... ON CONFLICT."""
|
|
if not rows:
|
|
return
|
|
|
|
cur = conn.cursor()
|
|
try:
|
|
# Ensure table exists
|
|
_ensure_table_exists(conn, tname, columns)
|
|
|
|
# Build INSERT ... ON CONFLICT statement
|
|
cols_quoted = ", ".join([f'"{col}"' for col in columns])
|
|
updates = ", ".join([f'"{col}" = EXCLUDED."{col}"' for col in columns if col != "row_hash"])
|
|
|
|
# Build the base query template for execute_values
|
|
base_query = f'INSERT INTO "{tname}" ({cols_quoted}) VALUES %s'
|
|
if updates:
|
|
conflict_query = f'{base_query} ON CONFLICT ("row_hash") DO UPDATE SET {updates}'
|
|
else:
|
|
conflict_query = f'{base_query} ON CONFLICT ("row_hash") DO NOTHING'
|
|
|
|
# Prepare values as list of tuples
|
|
values = [tuple(row.get(col) for col in columns) for row in rows]
|
|
|
|
execute_values(cur, conflict_query, values, template=None, page_size=len(rows))
|
|
conn.commit()
|
|
except Exception:
|
|
conn.rollback()
|
|
raise
|
|
finally:
|
|
cur.close()
|
|
|
|
def _postgres_upsert(table: str, df: pd.DataFrame):
|
|
if df.empty:
|
|
return
|
|
|
|
tname = table
|
|
total = len(df)
|
|
columns = [col for col in df.columns if col != "row_hash"] + ["row_hash"]
|
|
|
|
ranges = [(0, min(UPSERT_CHUNK,total))] if SINGLE_CHUNK else \
|
|
[(i, min(i+UPSERT_CHUNK,total)) for i in range(0,total,UPSERT_CHUNK)]
|
|
|
|
if not PRINT_PROGRESS:
|
|
lbl = "(single-chunk)" if SINGLE_CHUNK else f"chunk_size={UPSERT_CHUNK}"
|
|
print_flush(f"{_ts()} | ☁️ Starting upsert → [{tname}] rows={total:,} {lbl}")
|
|
|
|
sent = 0
|
|
conn = None
|
|
|
|
try:
|
|
for start, end in ranges:
|
|
part = df.iloc[start:end].copy()
|
|
if "row_hash" in part.columns:
|
|
part = part.drop_duplicates(subset=["row_hash"], keep="first")
|
|
|
|
# Ensure all columns are present
|
|
for col in columns:
|
|
if col not in part.columns:
|
|
part[col] = None
|
|
|
|
rows = part[columns].to_dict(orient="records")
|
|
|
|
tried_prune = False
|
|
for attempt in range(1, MAX_RETRIES + 1):
|
|
try:
|
|
if conn is None or conn.closed:
|
|
conn = _pg_conn()
|
|
_upsert_slice(conn, tname, rows, columns)
|
|
break
|
|
except Exception as e:
|
|
missing = _extract_missing_col_from_error(str(e))
|
|
if (missing is not None) and (missing in part.columns) and (not tried_prune):
|
|
if PRINT_PRUNE_LOGS:
|
|
print_flush(f"{_ts()} | ⚠️ [{tname}] pruning missing column '{missing}'")
|
|
part = part.drop(columns=[missing])
|
|
columns = [c for c in columns if c != missing]
|
|
rows = part[columns].to_dict(orient="records")
|
|
tried_prune = True
|
|
continue
|
|
if attempt == MAX_RETRIES:
|
|
raise
|
|
time.sleep(1.0 * attempt)
|
|
if conn and not conn.closed:
|
|
conn.close()
|
|
conn = None
|
|
|
|
sent += len(part)
|
|
if PRINT_PROGRESS:
|
|
print_flush(f"{_ts()} | ☁️ [{tname}] {sent:,}/{total:,}")
|
|
|
|
if not PRINT_PROGRESS:
|
|
print_flush(f"{_ts()} | ☁️ [{tname}] done {sent:,}/{total:,}")
|
|
finally:
|
|
if conn and not conn.closed:
|
|
conn.close()
|
|
|
|
def _postgres_replace(table: str, df: pd.DataFrame):
|
|
"""Replace all data in table (delete then insert)."""
|
|
conn = _pg_conn()
|
|
try:
|
|
cur = conn.cursor()
|
|
cur.execute(f'DELETE FROM "{table}"')
|
|
conn.commit()
|
|
cur.close()
|
|
except Exception:
|
|
conn.rollback()
|
|
# Table might not exist, that's okay
|
|
finally:
|
|
conn.close()
|
|
|
|
_postgres_upsert(table, df)
|
|
|
|
def _load_to_postgres(df: pd.DataFrame, table_name: str, _source_path: str = ""):
|
|
ndf = _normalize_for_table(df, table_name)
|
|
if ndf.empty:
|
|
print_flush(f"{_ts()} | ☁️ Empty after normalize. Skipping PostgreSQL [{table_name}]")
|
|
return
|
|
|
|
ndf = _df_prepare_for_postgres(ndf)
|
|
|
|
if INCREMENTAL:
|
|
_postgres_upsert(table_name, ndf)
|
|
print_flush(f"{_ts()} | ☁️✅ Upserted {len(ndf):,} rows → PostgreSQL [{table_name}]")
|
|
else:
|
|
_postgres_replace(table_name, ndf)
|
|
print_flush(f"{_ts()} | ☁️✅ Replaced table with {len(ndf):,} rows → PostgreSQL [{table_name}]")
|
|
|
|
# ─────────────────────────────────────────────
|
|
# Orchestrator
|
|
# ─────────────────────────────────────────────
|
|
def _load_records_to_databases(records: List[Dict], table_name: str):
|
|
"""Load records from API into both SQLite and PostgreSQL."""
|
|
if not records:
|
|
print_flush(f"{_ts()} | ⚠️ No records to process")
|
|
return
|
|
|
|
# Convert records to DataFrame
|
|
df = pd.DataFrame(records)
|
|
df = _coerce_weird_numbers(df)
|
|
|
|
if WRITE_POSTGRES:
|
|
try:
|
|
_load_to_postgres(df, table_name)
|
|
except Exception as e:
|
|
print_flush(f"{_ts()} | ❌ (PostgreSQL) Error: {e}")
|
|
|
|
# ─────────────────────────────────────────────
|
|
# Main
|
|
# ─────────────────────────────────────────────
|
|
def sync_blackbox_flow():
|
|
"""Main sync function."""
|
|
print_flush(f"{_ts()} | 🚀 Starting BlackBox Stocks flow data sync...\n")
|
|
|
|
try:
|
|
# Parse command line arguments
|
|
parser = argparse.ArgumentParser(description="Sync BlackBox Stocks options flow data to databases")
|
|
parser.add_argument("--start-date", type=str, help="Start date (YYYY-MM-DD)")
|
|
parser.add_argument("--end-date", type=str, help="End date (YYYY-MM-DD)")
|
|
parser.add_argument("--limit", type=int, help="Maximum number of records to fetch")
|
|
parser.add_argument("--count", type=int, help="Maximum number of records to fetch (alias for --limit)")
|
|
parser.add_argument("--symbol", type=str, help="Filter by specific symbol")
|
|
parser.add_argument("--table", type=str, default="OptionsFlow_monthly", help="Target table name (default: OptionsFlow_monthly)")
|
|
|
|
args = parser.parse_args()
|
|
|
|
options = {}
|
|
if args.start_date:
|
|
options["startDate"] = args.start_date
|
|
if args.end_date:
|
|
options["endDate"] = args.end_date
|
|
if args.limit:
|
|
options["count"] = args.limit
|
|
elif args.count:
|
|
options["count"] = args.count
|
|
if args.symbol:
|
|
options["symbol"] = args.symbol
|
|
|
|
table_name = args.table
|
|
|
|
# Default to today if no dates provided
|
|
if not options.get("startDate") and not options.get("endDate"):
|
|
today = date.today().isoformat()
|
|
options["startDate"] = today
|
|
options["endDate"] = today
|
|
print_flush(f"{_ts()} | 📅 No date range specified, using today: {today}")
|
|
|
|
print_flush(f"{_ts()} | 📥 Fetching flow data from BlackBox API...")
|
|
print_flush(f"{_ts()} | Options: {options}")
|
|
|
|
# Fetch data from API
|
|
api_records = fetch_blackbox_flow(options)
|
|
print_flush(f"{_ts()} | ✅ Fetched {len(api_records)} records from API")
|
|
|
|
if len(api_records) == 0:
|
|
print_flush(f"{_ts()} | ⚠️ No records returned from API")
|
|
return
|
|
|
|
# Log sample record
|
|
if len(api_records) > 0:
|
|
print_flush(f"\n{_ts()} | 📋 Sample API record structure:")
|
|
print_flush(json.dumps(api_records[0], indent=2, default=str))
|
|
|
|
# Map API records to database schema
|
|
print_flush(f"\n{_ts()} | 🔄 Mapping records to database schema...")
|
|
mapped_records = [map_blackbox_to_database(record) for record in api_records]
|
|
print_flush(f"{_ts()} | ✅ Mapped {len(mapped_records)} records")
|
|
|
|
# Log sample mapped record
|
|
if len(mapped_records) > 0:
|
|
print_flush(f"\n{_ts()} | 📋 Sample mapped record:")
|
|
print_flush(json.dumps(mapped_records[0], indent=2, default=str))
|
|
|
|
# Insert into databases
|
|
print_flush(f"\n{_ts()} | 💾 Inserting records into databases...")
|
|
_load_records_to_databases(mapped_records, table_name)
|
|
|
|
print_flush(f"\n{_ts()} | ✅ Successfully synced {len(mapped_records)} records")
|
|
|
|
# Summary
|
|
print_flush(f"\n{_ts()} | 📊 Sync Summary:")
|
|
print_flush(f"{_ts()} | Fetched from API: {len(api_records)} records")
|
|
print_flush(f"{_ts()} | Inserted into DB: {len(mapped_records)} records")
|
|
|
|
# Get total count in PostgreSQL database
|
|
if WRITE_POSTGRES:
|
|
try:
|
|
conn = _pg_conn()
|
|
cur = conn.cursor()
|
|
cur.execute(f'SELECT COUNT(*) FROM "{table_name}"')
|
|
total_count = cur.fetchone()[0]
|
|
print_flush(f"{_ts()} | Total records in PostgreSQL DB: {total_count}")
|
|
conn.close()
|
|
except Exception as e:
|
|
print_flush(f"{_ts()} | ⚠️ Could not get PostgreSQL count: {e}")
|
|
|
|
except Exception as e:
|
|
print_flush(f"\n{_ts()} | ❌ Sync failed: {e}")
|
|
import traceback
|
|
traceback.print_exc()
|
|
raise
|
|
|
|
if __name__ == "__main__":
|
|
if WRITE_POSTGRES:
|
|
print_flush(f"{_ts()} | ☁️ PostgreSQL: {POSTGRES_HOST}:{POSTGRES_PORT}/{POSTGRES_DB} "
|
|
f"(schema={POSTGRES_SCHEMA}) | INCREMENTAL={INCREMENTAL} | chunk={UPSERT_CHUNK}")
|
|
sync_blackbox_flow()
|