From bfc048c61e84a84cc198c6fdbff627acff76d285 Mon Sep 17 00:00:00 2001 From: Armor <2654988228@qq.com> Date: Tue, 9 Jul 2024 11:47:13 +0800 Subject: [PATCH] Mainframe-Pro update --- data/note.txt | 66 ++++++++ logs/log.txt | 0 src/KSP-tools/EngineDataProcess_DB_upload.py | 49 ++++++ src/KSP-tools/SRB_plot.py | 156 +++++++++++++++++++ src/KSP-tools/antenna_patch_generator.py | 5 +- src/KSP-tools/data/thrust_curve.txt | 129 +++++++++++++++ src/KSP-tools/engine_auto_price.py | 134 ++++++++++++++++ src/My_Logger/__init__.py | 0 src/My_Logger/logger.py | 62 ++++++++ src/PostgresSQL/__init__.py | 0 src/PostgresSQL/databaseHandler.py | 77 +++++++++ 11 files changed, 676 insertions(+), 2 deletions(-) create mode 100644 data/note.txt create mode 100644 logs/log.txt create mode 100644 src/KSP-tools/EngineDataProcess_DB_upload.py create mode 100644 src/KSP-tools/SRB_plot.py create mode 100644 src/KSP-tools/data/thrust_curve.txt create mode 100644 src/KSP-tools/engine_auto_price.py create mode 100644 src/My_Logger/__init__.py create mode 100644 src/My_Logger/logger.py create mode 100644 src/PostgresSQL/__init__.py create mode 100644 src/PostgresSQL/databaseHandler.py diff --git a/data/note.txt b/data/note.txt new file mode 100644 index 0000000..64e0c53 --- /dev/null +++ b/data/note.txt @@ -0,0 +1,66 @@ +Soooo, lets think about it. +在银河系，一个离银心不近也不远的地方，诞生了希尔文明。 +他们蹒跚起步，征服了陆地，迈向了海洋。他们很幸运， + +一个路过的黑洞扯碎了希尔文明的母星系，但是希尔文明没有心灰气馁，他们抢在灾难降临之前，倾尽全部的智慧，把文明的种子播撒了出去。 +现在，距离希尔舰队启航，已经过去了15个千年。新生的希尔人已经习惯了在群星中的生活。他们出生在星舰，在星河间穿行，以探索为荣耀。 +舰队偶尔会在某个星系落脚，补充星际旅行所必需的燃料和物质。但他们不会长久的停留。 +为了保证舰队的出现不会干涉到孕育中的文明种子，在舰队泊入一个恒星系之前，需要派出探险队确认目标星系没有生命诞生。 +这是一个重要而有趣的工作，在年轻人中很受欢迎，但只有最优秀的探星者才能参与这项充满荣耀的工作，驾驶着【星火】级星系探索飞船，在舰载AI的帮助下对恒星系展开细致而全面的探索。 +以及为后续舰队补充资源建立前哨基地。 +这是严垣从小就梦寐以求的工作。能够离开【此处需要插入某个形容狭小的形容词】的星舰，前往未知的世界探索，这样的生活实在是令人心潮澎湃。 +在经过了10余年的勤学苦读，12个月的艰苦训练后，严垣如愿以偿，成为了一名光荣的探星者。 + + + + + + + + + + + + +We need to ensure a 1600 arrival in the destination, leaving us with around 30 minutes of redundancy. +The current estimation of travel time is around 90 minutes. +Which means a 1400 departure is optimal. + +We need to check the traffic status 10 minutes before departure. Anything longer than 120 minutes will forbid the driving plan. +In which case we will transition into public traffic. + +We will need to pack the consumables, including water and necessary food. Which I intend to be prepared before lunchtime. +We also need to pack 2 set of personal terminals, power bank, Etc + +1200 Consumables preparation complete +1350 Final travel plan decision point. + In the event of the selection of public transportation, the plan is bike-9-1-5-bus. With a 1400 departure we should catch the 1540 bus. +1355 Donning the Doctor's gear. +1400 Departure. + + + + + + + + + + + + +I am listening to the OST of Lone-trail while I type this down. +Kristen. +I'm not going to deny, that I like her. +I like the idea of exploration. I like the idea of pushing the boundary. And she managed to do just that. +Comparing to her colleague, she is just a mortal. She is not as indestructible as Seria. She does not have the noble blood like Muelsyse. +But her none-stopping spirit of advance is the purest of them all. And under such spirit, Rhine Lab was found and united. +Her sins cannot be neglected, that I admit. But I can not resist my desire to salute her. +I salute to a pioneer, a scientist, a visionary, who gave all, risk all to challenge the unknown. +And now, even sky is not the limit. +Mortals are weak, but in that weakness, greatness is born. I find this phrase very suitable for her. +Ad astra, to the stars. Kristen, good night. + +And about Lone-trail, there's more I enjoy. +Kristen is a great character that I like so much, but the storyline did not just focus on how great her is. +It did not just shape a great character and call her sins necessary evil. \ No newline at end of file diff --git a/logs/log.txt b/logs/log.txt new file mode 100644 index 0000000..e69de29 diff --git a/src/KSP-tools/EngineDataProcess_DB_upload.py b/src/KSP-tools/EngineDataProcess_DB_upload.py new file mode 100644 index 0000000..a2f6922 --- /dev/null +++ b/src/KSP-tools/EngineDataProcess_DB_upload.py @@ -0,0 +1,49 @@ +import pandas as pd +from src.PostgresSQL.databaseHandler import DatabaseHandler + + +def process_excel(): + path = 'C://Users//26549//OneDrive//文档//KSP Engine Tweak Chart.xlsx' + excel_content = pd.read_excel(path) + data_list = [] + for row, content in excel_content.iterrows(): + if isinstance(content['Vac-ISP'], str) and '-' in content['Vac-ISP']: + content['SL-ISP'], content['Vac-ISP'] = float(content['Vac-ISP'].split('-')[0]), float( + content['Vac-ISP'].split('-')[1]) + else: + content['SL-ISP'] = content['Vac-ISP'] + data_list.append(dict(content)) + + df = pd.DataFrame(data_list, columns=[item for item in data_list[0]]) + df.to_excel('data//KSP Engine Tweak Chart_NEW.xlsx') + + +def upload_to_database(database_config): + path = 'C://Users//26549//OneDrive//文档//KSP Engine Tweak Chart.xlsx' + excel_content = pd.read_excel(path, keep_default_na=False) + + table_structure = dict(excel_content.loc[1]) + data_content_list = [dict(values) for _key, values in excel_content.iterrows()] + for key, values in table_structure.items(): + if str(values) in ('nan', 'nat'): + table_structure[key] = 'TEXT' + elif isinstance(values, float): + table_structure[key] = 'FLOAT4' + else: + table_structure[key] = 'TEXT' + handler = DatabaseHandler(database_config) + + handler.create_table(table_structure, 'Engine_Data', force_create=True) + handler.insert_into_table('Engine_Data', data_content_list) + + handler.close() + + +if __name__ == '__main__': + # process_excel() + config = {'database': 'KSP', 'user': 'postgres', 'password': 'Armor17909', 'host': '127.0.0.1', 'port': '5432'} + upload_to_database(config) + + + + print('done') diff --git a/src/KSP-tools/SRB_plot.py b/src/KSP-tools/SRB_plot.py new file mode 100644 index 0000000..4899c71 --- /dev/null +++ b/src/KSP-tools/SRB_plot.py @@ -0,0 +1,156 @@ +import numpy as np +from matplotlib import pyplot as plt + + +def cal_fuel_consumption_ratio(thrust, ISP): + return float(thrust) / (9.80665 * float(ISP)) + + +def get_thrust(fuel_percentage): + global thrust_curve + fuel_level = [i for i in thrust_curve] + for i in range(len(fuel_level) - 1): + if fuel_level[i] >= fuel_percentage > fuel_level[i + 1]: + prev_fuel = fuel_level[i] + nex_fuel = fuel_level[i + 1] + prev_thrust = thrust_curve[prev_fuel] + next_thrust = thrust_curve[nex_fuel] + if prev_thrust == next_thrust: + return prev_thrust + k = (prev_thrust - next_thrust) / (prev_fuel - nex_fuel) + result = prev_thrust - k * abs(prev_fuel - fuel_percentage) + return result + return 0 + + +def output_thrust_curve(): + global thrust_curve + print("thrustCurve") + print("{") + for item in thrust_curve: + print(f' key = {item} {thrust_curve[item]}') + print("}") + +""" +thrust_curve = { + 1.0000: 0.95, + 0.9997: 0.95, + 0.9500: 0.98, + 0.8000: 1.0, + 0.7500: 0.98, + 0.7000: 0.95, + 0.6500: 0.9, + 0.6000: 0.85, + 0.5500: 0.8, + 0.5000: 0.75, + 0.4700: 0.764, + 0.4400: 0.778, + 0.4000: 0.7966666666666666, + 0.3500: 0.82, + 0.3200: 0.826, + 0.3000: 0.83, + 0.2700: 0.826, + 0.2500: 0.84, + 0.2300: 0.824, + 0.2000: 0.80, + 0.1700: 0.77, + 0.1500: 0.75, + 0.1300: 0.67, + 0.1000: 0.55, + 0.0900: 0.5, + 0.0800: 0.45, + 0.0700: 0.40, + 0.0600: 0.35714285714285715, + 0.0500: 0.3142857142857143, + 0.0400: 0.27142857142857146, + 0.0350: 0.25, + 0.0300: 0.23, + 0.0200: 0.19, + 0.0150: 0.17, + 0.0100: 0.15, + 0.0050: 0.12916666666666665, + 0.0028: 0.12, + 0.0018: 0.08, + 0.0005: 0.05, + 0.0000: 0.03} +""" +f = open('data/thrust_curve.txt', 'r') +char = f.readlines() +f.close() + +thrust_curve = {} +for line in char: + temp = line.strip().replace('key = ', '').split(' ') + thrust_curve[float(temp[0])] = float(temp[1]) + # print(temp) + + +print(thrust_curve) + + +dry_mass = 90 +fuel_mass = 500 +thrust = 14750 +ISP = 268 + +plt.subplot(2, 2, 1) +plt.title("Thrust Level to Fuel Comsumption") +x = [1 - i for i in thrust_curve] +y = [thrust_curve[i] for i in thrust_curve] + +plt.plot(x, y) + +plt.ylabel('Thrust Level') + +plt.subplot(2, 2, 2) +plt.title("Acceleraion to Fuel Comsumption") +twr = [thrust * thrust_curve[i] / (dry_mass + i * fuel_mass) for i in thrust_curve] + +plt.plot(x, twr) + +plt.xlabel('Fuel Consumed') +plt.ylabel('Acceleraion') + +remaining_fuel_mass = fuel_mass + +time_list = [] +thrust_list = [] +acc_list = [] + +now_step_time = 0 +now_step_thrust = 1 +now_step_fuel_percentage = 1 +while remaining_fuel_mass >= 0: + time_list.append(now_step_time) + thrust_list.append(get_thrust(now_step_fuel_percentage)) + now_step_acc = thrust_list[-1] * thrust / (dry_mass + fuel_mass * now_step_fuel_percentage) + acc_list.append(now_step_acc) + # print(now_step_time, now_step_thrust, now_step_fuel_percentage, now_step_acc, remaining_fuel_mass) + now_step_time += 0.01 + remaining_fuel_mass -= cal_fuel_consumption_ratio(thrust * now_step_thrust, ISP) * 0.01 + now_step_fuel_percentage = remaining_fuel_mass / fuel_mass + now_step_thrust = get_thrust(now_step_fuel_percentage) + if now_step_acc == 0: + print('error') + print('remaining_fuel_mass', remaining_fuel_mass) + print('now_step_fuel_percentage', now_step_fuel_percentage) + break +plt.subplot(2, 2, 4) +plt.title("Acceleraion to Time") + +plt.plot(time_list, acc_list) + +plt.xlabel('Time in second') +plt.ylabel('Acceleraion') + +plt.subplot(2, 2, 3) +plt.title("Thrust to Time") + +plt.plot(time_list, thrust_list) + +plt.xlabel('Time in second') +plt.ylabel('Thrust level') + +plt.show() +print(time_list[-1]) +output_thrust_curve() diff --git a/src/KSP-tools/antenna_patch_generator.py b/src/KSP-tools/antenna_patch_generator.py index 1f4b0c2..63c2694 100644 --- a/src/KSP-tools/antenna_patch_generator.py +++ b/src/KSP-tools/antenna_patch_generator.py @@ -5,8 +5,8 @@ import datetime def gen_single_patch(antenna): result_str = f'// {antenna["Part Showname"]}\n// Range: ' range_raw = antenna["Range Raw"] - if range_raw >= 299792458*86400*365: - range_str, range_unit = range_raw / (299792458*86400*365), ' light years' + if range_raw >= 299792458 * 86400 * 365: + range_str, range_unit = range_raw / (299792458 * 86400 * 365), ' light years' elif range_raw >= 149597871000: range_str, range_unit = range_raw / 149597871000, 'AU' elif range_raw >= 1000000000: @@ -69,6 +69,7 @@ def gen_single_patch(antenna): # print(result_str) return result_str + if __name__ == '__main__': path = 'D://OneDrive//文档//KSP Engine Tweak Chart.xlsx' excel_content = pd.read_excel(path, sheet_name='Communication') diff --git a/src/KSP-tools/data/thrust_curve.txt b/src/KSP-tools/data/thrust_curve.txt new file mode 100644 index 0000000..38f9a65 --- /dev/null +++ b/src/KSP-tools/data/thrust_curve.txt @@ -0,0 +1,129 @@ +key = 1.00000 0.945 +key = 0.98942 0.945 +key = 0.97888 0.942 +key = 0.96834 0.942 +key = 0.95773 0.948 +key = 0.947 0.959 +key = 0.93618 0.967 +key = 0.92524 0.978 +key = 0.91426 0.981 +key = 0.90323 0.986 +key = 0.89216 0.989 +key = 0.8811 0.989 +key = 0.87 0.992 +key = 0.8589 0.992 +key = 0.84777 0.994 +key = 0.83665 0.994 +key = 0.82552 0.994 +key = 0.81439 0.994 +key = 0.80323 0.997 +key = 0.79207 0.997 +key = 0.78088 1 +key = 0.7697 1 +key = 0.75851 1 +key = 0.74744 0.989 +key = 0.73665 0.964 +key = 0.7261 0.942 +key = 0.71568 0.932 +key = 0.70544 0.915 +key = 0.69526 0.91 +key = 0.68518 0.901 +key = 0.67521 0.89 +key = 0.66537 0.88 +key = 0.65562 0.871 +key = 0.64599 0.86 +key = 0.63649 0.849 +key = 0.62711 0.838 +key = 0.61782 0.83 +key = 0.60862 0.822 +key = 0.59952 0.814 +key = 0.59047 0.808 +key = 0.58152 0.8 +key = 0.57269 0.789 +key = 0.56392 0.784 +key = 0.55527 0.773 +key = 0.54672 0.765 +key = 0.53823 0.759 +key = 0.52976 0.756 +key = 0.52136 0.751 +key = 0.51302 0.745 +key = 0.5048 0.734 +key = 0.49671 0.723 +key = 0.48868 0.718 +key = 0.4807 0.713 +key = 0.47273 0.713 +key = 0.46476 0.713 +key = 0.45675 0.715 +key = 0.44869 0.721 +key = 0.44056 0.726 +key = 0.43241 0.729 +key = 0.42422 0.732 +key = 0.416 0.734 +key = 0.40772 0.74 +key = 0.39941 0.743 +key = 0.39107 0.745 +key = 0.3827 0.748 +key = 0.3743 0.751 +key = 0.36586 0.754 +key = 0.3574 0.756 +key = 0.34891 0.759 +key = 0.34035 0.765 +key = 0.33176 0.767 +key = 0.32315 0.77 +key = 0.31453 0.77 +key = 0.30588 0.773 +key = 0.29723 0.773 +key = 0.28855 0.776 +key = 0.27987 0.776 +key = 0.2712 0.776 +key = 0.26252 0.776 +key = 0.25387 0.773 +key = 0.24531 0.765 +key = 0.23682 0.759 +key = 0.22841 0.751 +key = 0.2201 0.743 +key = 0.21185 0.737 +key = 0.20376 0.724 +key = 0.19575 0.715 +key = 0.18781 0.71 +key = 0.18005 0.694 +key = 0.17241 0.683 +key = 0.16493 0.669 +key = 0.1575 0.663 +key = 0.15011 0.661 +key = 0.14278 0.655 +key = 0.13548 0.652 +key = 0.12824 0.647 +key = 0.12115 0.633 +key = 0.11416 0.625 +key = 0.10726 0.617 +key = 0.10045 0.609 +key = 0.0937 0.603 +key = 0.08704 0.595 +key = 0.08051 0.584 +key = 0.07406 0.576 +key = 0.06771 0.568 +key = 0.06155 0.551 +key = 0.05556 0.535 +key = 0.0497 0.524 +key = 0.04387 0.521 +key = 0.03804 0.521 +key = 0.03231 0.513 +key = 0.0269 0.483 +key = 0.02193 0.444 +key = 0.01751 0.395 +key = 0.0138 0.332 +key = 0.01085 0.264 +key = 0.00863 0.198 +key = 0.00682 0.162 +key = 0.00528 0.138 +key = 0.00396 0.118 +key = 0.00294 0.091 +key = 0.0022 0.066 +key = 0.00164 0.05 +key = 0.00126 0.033 +key = 0.00098 0.025 +key = 0.00073 0.023 +key = 0.00051 0.02 +key = 0.00041 0.009 +key = 0.0 0.005 diff --git a/src/KSP-tools/engine_auto_price.py b/src/KSP-tools/engine_auto_price.py new file mode 100644 index 0000000..7b08eed --- /dev/null +++ b/src/KSP-tools/engine_auto_price.py @@ -0,0 +1,134 @@ +import pandas as pd +import numpy as np +import copy +import time +import json + + +class Queue: + def __init__(self, length): + self.queue = [] + self.length = length + + def in_queue(self, item): + self.queue.append(item) + if len(self.queue) > self.length: + self.queue.pop(0) + + def content(self): + return self.queue + + +def cal_price(engine, factor_dict): + result = 0 + for key, value in factor_dict.items(): + x = value * engine[key] + result += x if str(x) != 'nan' else 0 + return result + + +def generate_full_config(single_config, full_config_list): + new_dimension = list(np.arange(single_config[0], single_config[1], single_config[2]).round(4)) + if not full_config_list: + return [[item] for item in new_dimension] + new_full_config_list = [] + for og_list in full_config_list: + for item in new_dimension: + new_full_config_list.append(copy.deepcopy(og_list + [item])) + return new_full_config_list + + +def generate_full_config_dict_list(single_config, full_config_list, key): + new_dimension = list(np.arange(single_config[0], single_config[1], single_config[2]).round(4)) + if not full_config_list: + return [{key: float(val)} for val in new_dimension] + new_full_config_list = [] + for og_dict in full_config_list: + for item in new_dimension: + new_dict = copy.deepcopy(og_dict) + new_dict[key] = float(item) + new_full_config_list.append(new_dict) + return new_full_config_list + + +def calculate_config_factor(local_config_dict, target_list): + gen_config_start_time = time.time() + full_config = [] + for key, content in local_config_dict.items(): + # full_config = generate_full_config(content, full_config) + full_config = generate_full_config_dict_list(content, full_config, key) + with open('data/full_config.txt', 'w', encoding='utf-8') as file: + file.write(json.dumps(full_config)) + # file = open('data/full_config.txt', 'r', encoding='utf-8') + # full_config = json.loads(file.read()) + print('Config Time:', time.time() - gen_config_start_time) + + calculation_start = time.time() + min_div = 999999999999999999999 + best_config = None + best_config_queue = Queue(10) + for config in full_config: + div = 0 + for engine in target_list: + if engine['SL-ISP'] >= 200: + # auto_price = (engine['SL-ISP'] / 260) * (engine["Vac-ISP"] / 290) * engine['TWR'] * engine['Max Thrust'] * (1 - engine['Throttle Range']) * (1 + 0.01*engine['TVC']) / 4 + auto_price = cal_price(engine, config) + price = engine['Price KD'] + div += abs(auto_price - price) + # print(auto_price, price, auto_price - price, auto_price / price) + # if auto_price / price >= 2: + # raise ValueError + if min_div > div: + min_div = div + best_config = config + best_config_queue.in_queue([min_div, config]) + print('Calculation Time:', time.time() - calculation_start) + return best_config_queue + + +def cal_single_div(target_list, config): + div = 0 + for engine in target_list: + if engine['SL-ISP'] >= 200: + auto_price = cal_price(engine, config) + price = engine['Price KD'] + div += abs(auto_price - price) + return div + + +if __name__ == '__main__': + path = 'D://OneDrive//文档//KSP Engine Tweak Chart.xlsx' + excel_content = pd.read_excel(path, sheet_name='Engine Database') + full_list = [] + for key, content in excel_content.iterrows(): + if content['Fuel Type'] == 'Kerolox' and content['Cycle'] == 'Gas Generator': + full_list.append(dict(content)) + + tmp_config = {'Max Thrust': 1.7, 'SL-ISP': 0.2, 'TR Adjusted': 0.2, 'TVC Adjusted': 0.5, 'TWR': 0.5, 'Vac-ISP': 0.2} + div = cal_single_div(full_list, tmp_config) + + config_dict = { + 'SL-ISP': (0.01, 0.51, 0.05), + 'Vac-ISP': (0.01, 0.51, 0.05), + 'TWR': (0.01, 1.01, 0.1), + 'Max Thrust': (1.5, 2.0, 0.05), + 'TR Adjusted': (0.01, 1.01, 0.1), + 'TVC Adjusted': (0.01, 1.01, 0.1), + } + + factor_result = calculate_config_factor(config_dict, full_list) + + auto_dict = {'Max Thrust': 0.3, 'SL-ISP': 0.0012, 'TR Adjusted': 0.9, 'TVC Adjusted': 0.52, 'TWR': 0.5, 'Vac-ISP': 0.0013} + div = 0 + for engine in full_list: + if engine['SL-ISP'] >= 200: + auto_price = cal_price(engine, auto_dict) + price = engine['Price KD'] + print(auto_price, price, auto_price - price, auto_price / price) + + # if auto_price / price >= 2: + # raise ValueError + print('done') + + + diff --git a/src/My_Logger/__init__.py b/src/My_Logger/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/src/My_Logger/logger.py b/src/My_Logger/logger.py new file mode 100644 index 0000000..60cf30c --- /dev/null +++ b/src/My_Logger/logger.py @@ -0,0 +1,62 @@ +import logging +import datetime + + +class Logger: + def __init__(self, level="DEBUG"): + # 创建日志器对象 + self.logger = logging.getLogger(__name__) + self.logger.setLevel(level) + self.format = logging.Formatter(f'[%(filename)s: %(funcName)s:%(lineno)4d][%(asctime)s][%(levelname)-.5s]: ' + f'%(message)s', + datefmt='%Y-%m-%d %H:%M:%S') + + def console_handler(self, level="DEBUG"): + # 创建控制台的日志处理器 + console_handler = logging.StreamHandler() + console_handler.setLevel(level) + + # 处理器添加输出格式 + console_handler.setFormatter(self.format) + + # 返回控制器 + return console_handler + + def file_handler(self, level="DEBUG"): + # 创建文件的日志处理器 + file_handler = logging.FileHandler(f"../logs/log.txt", mode="a", encoding="utf-8") + file_handler.setLevel(level) + + # 处理器添加输出格式 + file_handler.setFormatter(self.format) + + # 返回控制器 + return file_handler + + def get_log(self): + # 日志器中添加控制台处理器 + self.logger.addHandler(self.console_handler()) + # 日志器中添加文件处理器 + self.logger.addHandler(self.file_handler()) + + # 返回日志实例对象 + return self.logger + + +if __name__ == '__main__': + print('111') + + # class TestLog(): + # def __init__(self): + # log = My_Logger() + # self.logger = log.get_log() + # + # def test_baili_01(self): + # self.logger.info("开始执行") + # self.logger.warning("结束执行") + # + # + # # 实例化 + # test = TestLog() + # # 调用类中的方法 + # test.test_baili_01() diff --git a/src/PostgresSQL/__init__.py b/src/PostgresSQL/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/src/PostgresSQL/databaseHandler.py b/src/PostgresSQL/databaseHandler.py new file mode 100644 index 0000000..86c3816 --- /dev/null +++ b/src/PostgresSQL/databaseHandler.py @@ -0,0 +1,77 @@ +import psycopg2 +from src.My_Logger.logger import Logger + + +class DatabaseHandler: + def __init__(self, config, logger=None): + if not logger: + log = Logger() + self.logger = log.get_log() + self.config = config + # {database="postgres", user="postgres", password="123456", host="localhost", port="5432"} + self.connection = psycopg2.connect(database=config['database'], user=config['user'], + password=config['password'], host=config['host'], port=config['port']) + self.cursor = self.connection.cursor() + self.cursor.execute(f"select tablename from pg_tables where schemaname='public'") + self.table_list = [item[0] for item in self.cursor.fetchall()] + + def create_table(self, table_structure, table_name, force_create=False): + table_name = table_name.lower() + if force_create and self.exists_table(table_name): + self.drop_table(table_name) + sql = f'Create Table {table_name}(' + content_list = [] + for key, values in table_structure.items(): + content_list.append(f'{key.replace(" ", "_").replace("-", "_")} {values}') + sql += ',\n'.join(content_list) + sql += ');' + self.connection.cursor().execute(sql) + self.connection.commit() + + def exists_table(self, table_name): + table_name = table_name.lower() + if table_name in self.table_list: + return True + else: + return False + + def drop_table(self, table_name, info_flag=False): + table_name = table_name.lower() + if not self.exists_table(table_name): + if info_flag: + self.logger.error(f'table: {table_name} not exists in database.') + return + self.cursor.execute(f'drop table {table_name}') + self.connection.commit() + if info_flag: + self.logger.info(f'table: {table_name} dropped') + return + + def insert_into_table(self, table_name, data_dict_list, batch=1000): + """ + + :param batch: + :param table_name: + :param data_dict_list: + :return: + """ + table_name = table_name.lower() + if not self.exists_table(table_name): + self.logger.error(f'table: {table_name} not exists in database.') + return + + table_structure = data_dict_list[0] + content_list = [] + for key, values in table_structure.items(): + content_list.append(f'{key.replace(" ", "_").replace("-", "_")}'.lower()) + + sql = f"INSERT INTO {table_name}({','.join(content_list)}) VALUES({','.join(['%s'] * len(content_list))})" + target = [list(item.values()) for item in data_dict_list] + + for i in range(0, len(target), batch): + self.cursor.executemany(sql, target[i: i+batch]) + self.connection.commit() + self.logger.info(f'Inserted into table {table_name}, rows: {min(i+batch, len(target))}/{len(target)}') + + def close(self): + self.connection.close()