SQLWorker.py - master - KuzminVD/Dumpy

SQLWorker.py
91 строка3 KB
vladimir
Load code
03 июн 2024, 00:08
03 июн 2024, 00:08f404d0f
import datetime
import os
import sqlite3
import pandas as pd
from datetime import datetime
import fnmatch

class SQLWorker:
    def __init__(self, plugs_dir, db_dir, db_filename):
        self.plugs_dir = plugs_dir
        self.db_dir = db_dir
        self.db_filename = db_filename

    def init_db(self):
        if (os.path.isfile(self.db_dir + self.db_filename)):
            os.remove(self.db_dir + self.db_filename)

        conn = sqlite3.connect(self.db_dir + self.db_filename)
        cur = conn.cursor()
        cur.execute("CREATE TABLE Logs ([Message], [Time])")
        conn.commit()
        conn.close()

    def write_log_to_db(self, message):
        conn = sqlite3.connect(self.db_dir + self.db_filename)
        cur = conn.cursor()
        values = [message, f'{datetime.now():%Y-%m-%d %H:%M:%S%z}']
        q = "INSERT INTO Logs VALUES (?,?)"
        cur.execute(q, values)
        conn.commit()
        conn.close()

    def write_csv_to_db(self, pattern):
        conn = sqlite3.connect(self.db_dir + self.db_filename)
        cur = conn.cursor()
        for root, directories, files in os.walk(self.plugs_dir):
            for filename in files:
                if (not fnmatch.fnmatch(filename, pattern)):
                    continue
                file = self.plugs_dir + filename
                table_name = filename.split('.')[-2]
                df = pd.read_csv(file, sep=',')
                headers = []
                for c in df.columns:
                    headers.append("[" + c + "]")
                cur.execute("CREATE TABLE " + table_name + " (" + ", ".join(headers) + ")")
                for i in range(df.shape[0]):
                    row = df.values[i]
                    values = []
                    for val in row:
                        if (str(val) == 'nan'):
                            values.append("''")
                        elif (type(val) == int):
                            values.append(str(val))
                        elif (type(val) == float):
                            values.append(str(val))
                        elif (type(val) == bool):
                            values.append(str(val))
                        else:
                            values.append("'" + str(val).replace("'", '"') + "'")
                    q = "INSERT INTO " + table_name + " (" + ", ".join(headers) + ") VALUES (" + ", ".join(values) + ")"
                    cur.execute(q)
                conn.commit()
        conn.close()

    def write_results_to_db(self, df):
        conn = sqlite3.connect(self.db_dir + self.db_filename)
        cur = conn.cursor()
        table_name = 'Results'
        headers = []
        for c in df.columns:
            headers.append("[" + c + "]")
        cur.execute("CREATE TABLE " + table_name + " (" + ", ".join(headers) + ")")
        for i in range(df.shape[0]):
            row = df.values[i]
            values = []
            for val in row:
                if (str(val) == 'nan'):
                    values.append("''")
                elif (type(val) == int):
                    values.append(str(val))
                elif (type(val) == float):
                    values.append(str(val))
                elif (type(val) == bool):
                    values.append(str(val))
                else:
                    values.append("'" + str(val).replace("'", '"') + "'")
            q = "INSERT INTO " + table_name + " (" + ", ".join(headers) + ") VALUES (" + ", ".join(values) + ")"
            cur.execute(q)
        conn.commit()
        conn.close()