adding work scripts

work/recovery-ui/recovery/lib/function/pgfunctions.py

@@ -0,0 +1,190 @@
+import os
+from io import StringIO
+
+import numpy as np
+import pandas as pd
+import psycopg2
+import psycopg2.extras as extras
+
+
+def execute_many(conn, df, table):
+    """
+    Using cursor.executemany() to insert the dataframe
+    """
+    # Create a list of tupples from the dataframe values
+    tuples = [tuple(x) for x in df.to_numpy()]
+    # Comma-separated dataframe columns
+    cols = ','.join(list(df.columns))
+    # SQL quert to execute
+    query  = "INSERT INTO %s(%s) VALUES(%%s,%%s,%%s)" % (table, cols)
+    cursor = conn.cursor()
+    try:
+        cursor.executemany(query, tuples)
+        conn.commit()
+    except (Exception, psycopg2.DatabaseError) as error:
+        print("Error: %s" % error)
+        conn.rollback()
+        cursor.close()
+        return 1
+    print("execute_many() done")
+    cursor.close()
+
+
+def execute_batch(conn, table, df, page_size):
+    # csvdata = open(csvfile, 'r')
+    # newcvsfile = csvfile.replace('data', 'data_filtered')
+    # newcsvdata = open(newcvsfile, 'w')
+    # for line in csvdata:
+    #     newline = line.replace("bytearray(b'\\xff\\xff\\xff\\x80')","").replace('"','')
+    #     newcsvdata.write(newline.replace('""',''))
+    # csvdata.close()
+    # newcsvdata.close()
+    # df = pd.read_csv(newcvsfile, dtype=object, low_memory=False)
+    """
+    Using psycopg2.extras.execute_batch() to insert the dataframe
+    """
+    # Create a list of tupples from the dataframe values
+    tuples = [tuple(x) for x in df.to_numpy()]
+    # Comma-separated dataframe columns
+    cols = ','.join(list(df.columns))
+    data = ''
+    for col in df.columns:
+        data = data + ',%%s'
+    rowdata = data[1:]
+    dynamicquery = f"INSERT INTO %s(%s) VALUES({rowdata})"
+    # SQL quert to execute
+    query  = dynamicquery % (table, cols)
+    cursor = conn.cursor()
+    try:
+        extras.execute_batch(cursor, query, tuples, page_size)
+        conn.commit()
+    except (Exception, psycopg2.DatabaseError) as error:
+        print("Error: %s" % error)
+        conn.rollback()
+        cursor.close()
+        return 1
+    print("execute_batch() done")
+    cursor.close()
+
+
+def execute_values(conn, df, table):
+    """
+    Using psycopg2.extras.execute_values() to insert the dataframe
+    """
+    # Create a list of tupples from the dataframe values
+    tuples = [tuple(x) for x in df.to_numpy()]
+    # Comma-separated dataframe columns
+    cols = ','.join(list(df.columns))
+    # SQL quert to execute
+    query  = "INSERT INTO %s(%s) VALUES %%s" % (table, cols)
+    cursor = conn.cursor()
+    try:
+        extras.execute_values(cursor, query, tuples)
+        conn.commit()
+    except (Exception, psycopg2.DatabaseError) as error:
+        print("Error: %s" % error)
+        conn.rollback()
+        cursor.close()
+        return 1
+    print("execute_values() done")
+    cursor.close()
+
+
+def execute_mogrify(conn, df, table):
+    """
+    Using cursor.mogrify() to build the bulk insert query
+    then cursor.execute() to execute the query
+    """
+    # Create a list of tupples from the dataframe values
+    tuples = [tuple(x) for x in df.to_numpy()]
+    # Comma-separated dataframe columns
+    cols = ','.join(list(df.columns))
+    # SQL quert to execute
+    cursor = conn.cursor()
+    values = [cursor.mogrify("(%s,%s,%s)", tup).decode('utf8') for tup in tuples]
+    query  = "INSERT INTO %s(%s) VALUES " % (table, cols) + ",".join(values)
+    
+    try:
+        cursor.execute(query, tuples)
+        conn.commit()
+    except (Exception, psycopg2.DatabaseError) as error:
+        print("Error: %s" % error)
+        conn.rollback()
+        cursor.close()
+        return 1
+    print("execute_mogrify() done")
+    cursor.close()
+
+
+def copy_from_file(conn, df, table):
+    """
+    Here we are going save the dataframe on disk as 
+    a csv file, load the csv file  
+    and use copy_from() to copy it to the table
+    """
+    # Save the dataframe to disk
+    tmp_df = "./tmp_dataframe.csv"
+    df.to_csv(tmp_df, index_label='id', header=False)
+    f = open(tmp_df, 'r')
+    cursor = conn.cursor()
+    try:
+        cursor.copy_from(f, table, sep=",")
+        conn.commit()
+    except (Exception, psycopg2.DatabaseError) as error:
+        os.remove(tmp_df)
+        print("Error: %s" % error)
+        conn.rollback()
+        cursor.close()
+        return 1
+    print("copy_from_file() done")
+    cursor.close()
+    os.remove(tmp_df)
+
+
+def copy_from_stringio(conn, df, table):
+    """
+    Here we are going save the dataframe in memory 
+    and use copy_from() to copy it to the table
+    """
+    # save dataframe to an in memory buffer
+    buffer = StringIO()
+    df.to_csv(buffer, index_label='id', header=False)
+    buffer.seek(0)
+    
+    cursor = conn.cursor()
+    try:
+        cursor.copy_from(buffer, table, sep=",")
+        conn.commit()
+    except (Exception, psycopg2.DatabaseError) as error:
+        print("Error: %s" % error)
+        conn.rollback()
+        cursor.close()
+        return 1
+    print("copy_from_stringio() done")
+    cursor.close()
+
+
+#----------------------------------------------------------------
+# SqlAlchemy Only
+#----------------------------------------------------------------
+# from sqlalchemy import create_engine
+
+# connect = "postgresql+psycopg2://%s:%s@%s:5432/%s" % (
+#     param_dic['user'],
+#     param_dic['password'],
+#     param_dic['host'],
+#     param_dic['database']
+# )
+
+# def to_alchemy(df):
+#     """
+#     Using a dummy table to test this call library
+#     """
+#     engine = create_engine(connect)
+#     df.to_sql(
+#         'test_table', 
+#         con=engine, 
+#         index=False, 
+#         if_exists='replace'
+#     )
+#     print("to_sql() done (sqlalchemy)")