TPM/data_init/data_init_tw_v0.py

import pandas as pd
import json
import yfinance as yf
import numpy as np
import psycopg2
from psycopg2.extras import execute_values
from tqdm import tqdm
from config import SQL_CONFIG
import os

with open('data_init/assets_tw.json') as f:
    data_tw = json.load(f)

# TW Stocks
conn = psycopg2.connect(**SQL_CONFIG)
cursor = conn.cursor()

# Create schema if not exists
try:
    cursor.execute("SELECT 1 FROM stock_price_tw LIMIT 1;")
except psycopg2.errors.UndefinedTable:
    print("📋 Tables not found, creating schema...")
    conn.rollback()
    sql_path = '/flask/sql_script/create_all.sql' if os.path.exists('/flask/sql_script/create_all.sql') else 'sql_script/create_all.sql'
    with open(sql_path, 'r') as f:
        cursor.execute(f.read())
    conn.commit()
    print("✅ Schema created successfully")

cursor.execute("SELECT COUNT(*) FROM stock_price_tw;")
row_count = cursor.fetchone()[0]
if row_count <= 0 :
    for ticker in tqdm(data_tw):
        df = yf.download(ticker, start="2007-01-01", progress=False, threads=False)
        if df is None or df.empty:
            continue
        price_col = 'Close' if 'Close' in df.columns else ('Adj Close' if 'Adj Close' in df.columns else None)
        if price_col is None:
            continue
        value = [(ticker, row.name.strftime('%Y-%m-%d'), float(row[price_col])) for _, row in df.iterrows()]
        with conn:
            with conn.cursor(cursor_factory=psycopg2.extras.RealDictCursor) as curs:
                sql = "insert into stock_price_tw (ticker, date, price) values %s"
                execute_values(curs, sql, value)
    print("Finish initialize")
else:
    print("DB already have data")