gctools/database/cet.py

import plac
import pandas as pd
from sqlalchemy import create_engine
import psycopg2
from os import path


@plac.pos("query", "SQL Query", type=str)
@plac.pos("mode", "", choices=["in", "out", "queryout"])
@plac.pos("csv_file", "", type=str)
@plac.opt("Server", "Hostname or DSN", type=str)
@plac.opt("database", "", type=str)
@plac.opt("User", "", type=str)
@plac.opt("Password", "", type=str)
@plac.flg("charset", "")
@plac.opt("Codepage", "", type=str)
@plac.opt("errorlog", "", type=str)
def run(
    query,
    mode,
    csv_file,
    Server="localhost\\GLOBALCUBE",
    database="master",
    User="sa",
    Password="Mffu3011#",
    charset=False,
    Codepage="65001",
    errorlog="error.log",
):
    # dsn = f"dsn={Server};uid={User};pwd={Password}"
    dsn = f"postgresql://{User}:{Password}@{Server}/{database}"
    if Codepage.isnumeric():
        Codepage = "cp" + Codepage
    if mode == "queryout":
        queryout(dsn, query, csv_file, Codepage, errorlog)
        return
    print("This is madness")


def convert_data(element):
    txt = str(element)
    txt = txt.replace("None", "")
    txt = txt.replace("False", "0").replace("True", "1")
    txt = txt.replace("\t", "").replace("\r", "").replace("\n", "")
    txt = txt.replace("\x81", "").replace("\x90", "")

    txt = "" if txt in ["nan", "NaT"] else txt
    return txt


def queryout(dsn, query, csv_file, codepage, errorlog):
    if path.exists(query):
        with open(query, "r", encoding=codepage) as frh:
            query = frh.read()

    try:
        conn = create_engine(dsn).connect().execution_options(stream_results=True)
        df = pd.read_sql(query, conn, chunksize=1000)
    except Exception as e:
        print(e.args[1])

    with open(csv_file, "w", encoding=codepage, errors="replace") as fwh:
        print("Kopiervorgang wird gestartet...")
        i = 0
        for chunk in df:
            chunk_dict = chunk.to_dict(orient="records")
            for row in chunk_dict:
                try:
                    fwh.write(("\t".join(map(convert_data, row.values())) + "\n"))
                except Exception as e:
                    print(e.args[1])
            i += len(chunk_dict)
            if (len(chunk_dict)) == 1000:
                print(f"1000 Zeilen zum SQL Server gesendet. Insgesamt gesendet: {i}")
        print("")
        print(f"{i} Zeilen kopiert.")


if __name__ == "__main__":
    plac.call(run)