You can not select more than 25 topics Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.
 
 
 
 
 
 
alpha_tools/import_alpha_txt.py

148 lines
4.6 KiB

# -*- coding: utf-8 -*-
import os
import psycopg2
class AlphaTxtImporter:
def __init__(self, txt_file='alpha.txt'):
self.database_name = "alpha"
self.txt_file = txt_file
self.connection = None
def get_connection(self):
"""获取数据库连接"""
if self.connection is None or self.connection.closed:
self.connection = psycopg2.connect(
host="localhost",
port="5432",
database=self.database_name,
user="jack",
password="aaaAAA111"
)
return self.connection
def close_connection(self):
"""关闭数据库连接"""
if self.connection and not self.connection.closed:
self.connection.close()
def load_alpha_from_txt(self):
"""从txt文件加载alpha表达式"""
if not os.path.exists(self.txt_file):
print(f"未找到 {self.txt_file} 文件")
return []
alpha_list = []
with open(self.txt_file, 'r', encoding='utf-8') as file:
for line_num, line in enumerate(file, 1):
line = line.strip()
if line and not line.startswith('#'): # 跳过空行和注释行
alpha_list.append(line)
print(f"{line_num} 行: {line}")
return alpha_list
def import_to_database(self, alpha_list):
"""导入alpha表达式到数据库"""
if not alpha_list:
print("没有要导入的alpha表达式")
return 0
conn = self.get_connection()
cursor = conn.cursor()
imported_count = 0
duplicate_count = 0
for alpha in alpha_list:
try:
# 使用 ON CONFLICT 避免重复插入
cursor.execute('''
INSERT INTO alpha_prepare (alpha, unused)
VALUES (%s, %s) ON CONFLICT (alpha) DO NOTHING
''', (alpha, True))
if cursor.rowcount > 0:
imported_count += 1
print(f"✅ 导入: {alpha}")
else:
duplicate_count += 1
print(f" 跳过(重复): {alpha}")
except Exception as e:
print(f"❌ 导入失败: {alpha} - 错误: {e}")
continue
conn.commit()
cursor.close()
return imported_count, duplicate_count
def check_existing_data(self):
"""检查数据库中已存在的数据"""
conn = self.get_connection()
cursor = conn.cursor()
cursor.execute("SELECT COUNT(*) FROM alpha_prepare")
total_count = cursor.fetchone()[0]
cursor.execute("SELECT COUNT(*) FROM alpha_prepare WHERE unused = TRUE")
unused_count = cursor.fetchone()[0]
cursor.close()
print(f"数据库统计:")
print(f" - 总alpha数量: {total_count}")
print(f" - 未使用的alpha数量: {unused_count}")
return total_count, unused_count
def run_import(self):
"""执行导入过程"""
print(f"开始从 {self.txt_file} 导入alpha表达式到 PostgreSQL...")
# 检查现有数据
print("\n=== 数据库当前状态 ===")
total_count, unused_count = self.check_existing_data()
# 加载txt文件
print(f"\n=== 加载 {self.txt_file} ===")
alpha_list = self.load_alpha_from_txt()
print(f"从文件加载了 {len(alpha_list)} 个alpha表达式")
if not alpha_list:
print("没有找到有效的alpha表达式,导入终止")
return
# 导入到数据库
print(f"\n=== 导入到数据库 ===")
imported_count, duplicate_count = self.import_to_database(alpha_list)
# 显示结果
print(f"\n=== 导入结果 ===")
print(f"文件中的表达式数量: {len(alpha_list)}")
print(f"成功导入: {imported_count}")
print(f"跳过重复: {duplicate_count}")
# 再次检查数据库状态
print(f"\n=== 导入后数据库状态 ===")
new_total_count, new_unused_count = self.check_existing_data()
print(f"新增记录: {new_total_count - total_count}")
print(f"新增未使用记录: {new_unused_count - unused_count}")
print("\n导入完成!")
def main():
# 可以指定不同的txt文件
txt_file = 'alpha.txt' # 可以改为其他文件名
importer = AlphaTxtImporter(txt_file)
try:
importer.run_import()
finally:
importer.close_connection()
if __name__ == "__main__":
main()