postgresql 数据快速转移脚本

# -*- coding:utf-8 -*-
import os
import pandas as pd
import psycopg2
from tqdm import tqdm
from io import BytesIO as StringIO
import time
from sqlalchemy import MetaData
from sqlalchemy import create_engine

homepath = os.environ['HOME']
BufferPath = homepath + "/bufferPath/"
if not os.path.exists(BufferPath):
    os.mkdir(BufferPath)


def get_tables(engine):
    cmd = "select  relname as table from pg_stat_user_tables where schemaname = 'public'"
    df = pd.read_sql(cmd, engine)
    return df['table'].tolist()


template_dict = {
    "database": "testdb",
    "user": "postgres",
    "password": "cohondob",
    "host": "127.0.0.1",
    "port": "5432"
}


class CopyTicksDB(object):
    '''
    账号密码地址数据库
    '''
    timescaleDBMode = False

    bufferPath = BufferPath

    from_database = ""
    from_user = ""
    from_password = ""
    from_host = ""
    from_port = ""

    target_database = ""
    target_user = ""
    target_password = ""
    target_host = ""
    target_port = ""

    from_conn = None
    target_conn = None

    columns = []

    createDDL = '''
    CREATE TABLE IF NOT EXISTS "%s" (
    datetime timestamp NOT NULL,
    instrument_id text NOT NULL,
    exchange text NULL,
    lastprice numeric(20,4) NOT NULL,
    volume numeric(20,4) NOT NULL,
    bidprice1 numeric(20,4) NOT NULL,
    askprice1 numeric(20,4) NOT NULL,
    bidvolume1 numeric(20,4) NOT NULL,
    askvolume1 numeric(20,4) NOT NULL,
    bidvolume5 numeric(20,4) NOT NULL,
    bidvolume4 numeric(20,4) NOT NULL,
    bidvolume3 numeric(20,4) NOT NULL,
    bidvolume2 numeric(20,4) NOT NULL,
    bidprice5 numeric(20,4) NOT NULL,
    bidprice4 numeric(20,4) NOT NULL,
    bidprice3 numeric(20,4) NOT NULL,
    bidprice2 numeric(20,4) NOT NULL,
    askvolume5 numeric(20,4) NOT NULL,
    askvolume4 numeric(20,4) NOT NULL,
    askvolume3 numeric(20,4) NOT NULL,
    askvolume2 numeric(20,4) NOT NULL,
    askprice5 numeric(20,4) NOT NULL,
    askprice4 numeric(20,4) NOT NULL,
    askprice3 numeric(20,4) NOT NULL,
    askprice2 numeric(20,4) NOT NULL,
    openinterest numeric(20,4) NOT NULL,
    turnover numeric(20,4) NOT NULL,
    lowerlimit numeric(20,4) NOT NULL,
    upperlimit numeric(20,4) NOT NULL
    );
    CREATE INDEX  IF NOT EXISTS "index_dt_%s"  ON "%s" USING btree (datetime, instrument_id);
    CREATE INDEX  IF NOT EXISTS "index_dtv_%s" ON "%s" USING btree (datetime,volume);
    '''

    querycmd = ""
    delcmd = ""

    def __init__(self, from_dict, target_dict, timescaleDBmode=False):
        self.timescaleDBMode = timescaleDBmode
        if not os.path.exists(self.bufferPath):
            os.mkdir(self.bufferPath)

        self.from_database = from_dict["database"]
        self.from_host = from_dict["host"]
        self.from_port = from_dict["port"]
        self.from_user = from_dict["user"]
        self.from_password = from_dict["password"]

        self.target_database = target_dict["database"]
        self.target_host = target_dict["host"]
        self.target_port = target_dict["port"]
        self.target_user = target_dict["user"]
        self.target_password = target_dict["password"]

        self.from_conn = psycopg2.connect(database=self.from_database, user=self.from_user, password=self.from_password,
                                          host=self.from_host, port=self.from_port)
        self.target_conn = psycopg2.connect(database=self.target_database, user=self.target_user,
                                            password=self.target_password,
                                            host=self.target_host, port=self.target_port)

    def __del__(self):
        self.from_conn.close()
        self.target_conn.close()

    def __downloadData(self, tableName):
        print("%s download data ....." % tableName)
        file_name = self.bufferPath + tableName + ".txt"
        cur = self.from_conn.cursor()
        cur.copy_expert("COPY (%s) TO STDOUT DELIMITER '|'" % (self.querycmd), open(file_name, "w"))
        self.from_conn.commit()
        cur.execute("Select * FROM \"%s\" limit 1" % tableName)
        self.from_conn.commit()
        cur.close()
        fields_names = [i[0] for i in cur.description]  # 参数0是元组的第一个参数  Column(name='datetime', type_code=1700)
        self.columns = fields_names
        print("download finish!")

    def __uploadData(self, tableName):
        print("try create table %s   ....." % tableName)
        DDL = self.createDDL % (tableName, tableName, tableName, tableName, tableName)
        cur = self.target_conn.cursor()
        cur.execute(DDL)
        self.target_conn.commit()
        if self.timescaleDBMode:
            self.createHyperTable(tableName)

        print("delete table data %s   ....." % tableName)
        cur.execute(self.delcmd)
        self.target_conn.commit()

        if (len(self.columns) == 0):
            cur.execute("Select * FROM \"%s\" limit 1" % tableName)
            self.from_conn.commit()
            fields_names = [i[0] for i in cur.description]  # 参数0是元组的第一个参数  Column(name='datetime', type_code=1700)
            self.columns = fields_names

        print("read file data %s   ....." % tableName)

        file_name = self.bufferPath + tableName + ".txt"
        totalNum = 0
        with open(file_name, 'r') as f:
            k = 0
            content = []
            for line in tqdm(f):
                content.append(line)
                k = k + 1
                totalNum = totalNum + 1
                if k >= 10000000:
                    if (len(content) > 0):
                        s = StringIO()
                        for text in content:
                            s.write(text.encode())
                        s.seek(0)
                        cur.copy_from(s, "\"" + tableName + "\"", columns=tuple(self.columns), sep='|')
                        self.target_conn.commit()

                    k = 0
                    content = []

            if (len(content) > 0):
                s = StringIO()
                for text in content:
                    s.write(text.encode())
                s.seek(0)
                cur.copy_from(s, "\"" + tableName + "\"", columns=tuple(self.columns), sep='|')
                self.target_conn.commit()

        cur.close()
        print("tableName:%s totalDataSize:%s ...." % (tableName, totalNum))
        print("upload finish")

    def download(self, beginDate, endDate, tableName):
        if (beginDate == "" and endDate == ""):
            self.querycmd = "SELECT * FROM \"%s\"   order by datetime asc,volume asc" % (tableName)
        else:
            self.querycmd = "SELECT * FROM \"%s\" where datetime >='%s 20:00:00' and datetime<='%s 17:00:00' order by datetime asc,volume asc" % (
                tableName, beginDate, endDate)
            print("cmd:%s   " % self.querycmd)
        self.__downloadData(tableName)

    def upload(self, beginDate, endDate, tableName):
        if (beginDate == "" and endDate == ""):
            print("---------------------------------------------------------")
            print("warnning you are deleteing all datas")
            print("---------------------------------------------------------")
            time.sleep(10)
            self.delcmd = "delete FROM \"%s\"" % (tableName)
        else:
            self.delcmd = "delete FROM \"%s\" where datetime >='%s 20:00:00' and datetime<='%s 17:00:00'" % (
                tableName, beginDate, endDate)
            print("cmd:%s   " % self.delcmd)
        self.__uploadData(tableName)

    def createHyperTable(self, tableName):
        cmd = "SELECT create_hypertable('\"%s\"', 'datetime');" % (tableName)
        cur = self.target_conn.cursor()
        try:
            print(cmd)
            cur.execute(cmd)
        except Exception as e:
            print(e)
        self.target_conn.commit()


class copyKbars2DB(object):
    '''
    账号密码地址数据库
    '''
    bufferPath = BufferPath

    from_database = ""
    from_user = ""
    from_password = ""
    from_host = ""
    from_port = ""

    target_database = ""
    target_user = ""
    target_password = ""
    target_host = ""
    target_port = ""

    from_conn = None
    target_conn = None

    columns = []

    createDDL = '''
        -- Drop table
        
        -- DROP TABLE "%s_1minute_bars";
        
        CREATE TABLE IF NOT EXISTS "%s_1minute_bars" (
            trading_day bpchar(8) NOT NULL,
            update_time bpchar(8) NOT NULL,
            open_price float8 NULL,
            highest_price float8 NULL,
            lowest_price float8 NULL,
            close_price float8 NULL,
            status int4 NULL,
            instrument_id varchar(15) NOT NULL,
            volume float8 NULL,
            ammount float8 NULL,
            "position" float8 NULL,
            avg_price float8 NULL,
            position_diff float8 NULL,
            CONSTRAINT "%s_1minute_bars_pk" PRIMARY KEY (trading_day, update_time, instrument_id)
        );
    '''

    querycmd = ""
    delcmd = ""

    def __init__(self, from_dict, target_dict):
        if not os.path.exists(self.bufferPath):
            os.mkdir(self.bufferPath)

        self.from_database = from_dict["database"]
        self.from_host = from_dict["host"]
        self.from_port = from_dict["port"]
        self.from_user = from_dict["user"]
        self.from_password = from_dict["password"]

        self.target_database = target_dict["database"]
        self.target_host = target_dict["host"]
        self.target_port = target_dict["port"]
        self.target_user = target_dict["user"]
        self.target_password = target_dict["password"]

        self.from_conn = psycopg2.connect(database=self.from_database, user=self.from_user, password=self.from_password,
                                          host=self.from_host, port=self.from_port)
        self.target_conn = psycopg2.connect(database=self.target_database, user=self.target_user,
                                            password=self.target_password,
                                            host=self.target_host, port=self.target_port)

    def __del__(self):
        self.from_conn.close()
        self.target_conn.close()

    def __downloadData(self, tableName):
        print("%s download data ....." % tableName)
        file_name = self.bufferPath + tableName + ".txt"
        cur = self.from_conn.cursor()
        cur.copy_expert("COPY (%s) TO STDOUT DELIMITER '|'" % (self.querycmd), open(file_name, "w"))
        self.from_conn.commit()
        cur.execute("Select * FROM \"%s_1minute_bars\" limit 1" % tableName)
        self.from_conn.commit()
        cur.close()
        fields_names = [i[0] for i in cur.description]  # 参数0是元组的第一个参数  Column(name='datetime', type_code=1700)
        self.columns = fields_names
        print("download finish!")

    def __uploadData(self, tableName):
        print("try create table %s   ....." % tableName)
        DDL = self.createDDL % (tableName, tableName, tableName)
        cur = self.target_conn.cursor()
        cur.execute(DDL)
        self.target_conn.commit()

        print("delete table data %s   ....." % tableName)
        cur.execute(self.delcmd)
        self.target_conn.commit()

        print("read file data %s   ....." % tableName)
        s = StringIO()
        file_name = self.bufferPath + tableName + ".txt"
        with open(file_name, 'r') as f:
            content = f.readlines()
            for line in tqdm(content):
                s.write(line.encode())
            s.seek(0)

        print("upload data %s   ....." % tableName)
        cur.copy_from(s, "\"%s_1minute_bars\"" % (tableName), columns=tuple(self.columns), sep='|')
        self.target_conn.commit()
        cur.close()
        print("upload finish")

    def download(self, beginDate, endDate, tableName):
        if (beginDate == "" and endDate == ""):
            self.querycmd = "SELECT * FROM \"%s_1minute_bars\"   order by trading_day asc,update_time asc" % (tableName)
        else:
            self.querycmd = "SELECT * FROM \"%s_1minute_bars\" where trading_day >='%s' and trading_day<='%s' order by trading_day asc,update_time asc" % (
                tableName, beginDate, endDate)
            print("cmd:%s   " % self.querycmd)
        self.__downloadData(tableName)

    def upload(self, beginDate, endDate, tableName):
        if (beginDate == "" and endDate == ""):
            print("---------------------------------------------------------")
            print("warnning you are deleteing all datas")
            print("---------------------------------------------------------")
            time.sleep(10)
            self.delcmd = "delete FROM \"%s\"_1minute_bars" % (tableName)
        else:
            self.delcmd = "delete FROM \"%s_1minute_bars\" where trading_day >='%s' and trading_day<='%s'" % (
                tableName, beginDate, endDate)
            print("cmd:%s   " % self.delcmd)
        self.__uploadData(tableName)


if __name__ == "__main__":
    engine = create_engine('postgresql+psycopg2://xxxxx:xxxxx@%s:xxxx/xxxxx' % ("127.0.0.1"))
    alltables = get_tables(engine)

    fd = {
        "database": "xxxx",
        "user": "xxxx",
        "password": "xxxx",
        "host": "xxxx",
        "port": "xxxx"
    }

    td = {
        "database": "xxxx",
        "user": "xxxx",
        "password": "xxxx",
        "host": "xxxx",
        "port": "xxxx"
    }

    # for symbol in alltables:
    #     copytickdb = CopyTicksDB(fd, td, True)
    #     copytickdb.download("20160101", "20200608", symbol)
    # copytickdb.upload("20200501", "20200608", symbol)

    symbollist = []
    for tableName in alltables:
        symbollist.append(tableName.split("_")[0])
    copyKbars2 = copyKbars2DB(fd, td)
    for symbol in symbollist:
        copyKbars2.download("20200614", "20200615", symbol)
        copyKbars2.upload("20200614", "20200615", symbol)