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