import multiprocessing
import glob
import pandas as pd
from sqlalchemy import create_engine

engine = create_engine("postgresql://postgres@localhost:5432/vnpy")


def save_bar(filename):
    vt_symbol = filename.split("_")[0]
    symbol, exchange = vt_symbol.split(".")
    df = pd.read_csv(
        filename,
        compression="gzip",
        parse_dates=True,
        skiprows=1,
        names=[
            "datetime",
            "open_price",
            "high_price",
            "low_price",
            "close_price",
            "volume",
            "open_interest",
        ],
    )
    df["symbol"] = symbol
    df["exchange"] = exchange
    df["interval"] = "1m"
    print(df.head(1))
    df.to_sql("dbbardata", engine, if_exists="append", index=False)


if __name__ == "__main__":
    all_files = glob.glob("*csv.gz")
    pool = multiprocessing.Pool(
        min(len(all_files), multiprocessing.cpu_count()), maxtasksperchild=1
    )
    results = []
    for file_name in all_files:
        res = pool.apply_async(save_bar, args=(file_name,))
        results.append(res)
    [res.get() for r in results]

open	high	low	close	volume	open_interest
2010/4/16 9:16	3450	3488	3450	3470	490	359
2010/4/16 9:17	3468	3473	3467	3467	306	521
2010/4/16 9:18	3467	3471	3466	3467	198	630
2010/4/16 9:19	3467	3468	3448	3448	282	764
2010/4/16 9:20	3448	3460	3448	3454	248	795
2010/4/16 9:21	3454	3456	3454	3456	109	841
2010/4/16 9:22	3457	3458	3456	3458	165	908
2010/4/16 9:23	3459	3460	3455	3455	206	980
2010/4/16 9:24	3455	3455	3450	3454	196	1020
2010/4/16 9:25	3454	3457	3453	3456	137	1076