Browse Source

优化UID导入

zhangdongming 2 weeks ago
parent
commit
e6e2ac6178
1 changed files with 58 additions and 46 deletions
  1. 58 46
      AdminController/UIDBurnManageController.py

+ 58 - 46
AdminController/UIDBurnManageController.py

@@ -16,6 +16,7 @@ from datetime import datetime
 from typing import Dict, Any
 from uuid import uuid4
 
+import xlrd
 from django.core.paginator import Paginator, EmptyPage, PageNotAnInteger
 from django.db import transaction
 from django.db.models import Q
@@ -269,7 +270,7 @@ class UIDBurnManageView(View):
             base_dir = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
             upload_dir = os.path.join(base_dir, 'static', 'uploaded_files')
             os.makedirs(upload_dir, exist_ok=True)
-            file_path = os.path.join(upload_dir, f"{task_id}.xlsx")
+            file_path = os.path.join(upload_dir, f"{task_id}.xls")
 
             with open(file_path, 'wb+') as destination:
                 for chunk in excel_file.chunks():
@@ -296,72 +297,85 @@ class UIDBurnManageView(View):
 
     @classmethod
     def _process_import_batch_async(cls, task_id, file_path, redis_key, batch_number):
-        """后台线程处理批量导入任务(适配新表结构)"""
+        """后台线程处理批量导入任务(兼容 xls 和 xlsx,读取第3列UID)"""
 
         redis_obj = RedisObject()
-
         try:
-            # 获取并更新任务状态为处理中
             task_data = json.loads(redis_obj.get_data(redis_key))
             task_data['status'] = 'processing'
             redis_obj.set_data(redis_key, json.dumps(task_data))
 
-            # 1. 读取Excel文件获取总行数
-            wb = load_workbook(file_path)
-            ws = wb.active
-            total_rows = ws.max_row
-
-            # 更新总行数和开始时间
-            task_data['total'] = total_rows
+            # 判断文件类型并读取数据
+            file_ext = os.path.splitext(file_path)[-1].lower()
+            uid_rows = []
+
+            if file_ext == '.xls':
+                # 使用 xlrd 读取 .xls 文件
+                workbook = xlrd.open_workbook(file_path)
+                sheet = workbook.sheet_by_index(0)
+                total_rows = sheet.nrows
+                for row_idx in range(0, total_rows):  # 改为从第0行开始
+                    row = sheet.row_values(row_idx)
+                    if len(row) > 2 and row[2]:
+                        uid = str(row[2]).strip()
+                        if uid:
+                            uid_rows.append(uid)
+
+            elif file_ext == '.xlsx':
+                # 使用 openpyxl 读取 .xlsx 文件
+                workbook = load_workbook(file_path, read_only=True)
+                sheet = workbook.active
+                for row in sheet.iter_rows(min_row=1, values_only=True):  # 改为从第1行开始
+                    if len(row) > 2 and row[2]:
+                        uid = str(row[2]).strip()
+                        if uid:
+                            uid_rows.append(uid)
+            task_data['total'] = len(uid_rows)
             task_data['start_time'] = int(time.time())
             redis_obj.set_data(redis_key, json.dumps(task_data))
 
-            # 2. 创建批次记录
+            # 创建批次记录
             current_time = int(time.time())
             with transaction.atomic():
                 batch = BurnBatch(
                     batch_number=batch_number,
                     purpose='批次导入',
                     created_time=current_time,
-                    manager='system',  # 默认系统导入
-                    total_uid=0  # 初始为0,处理完成后更新
+                    manager='system',
+                    total_uid=0
                 )
                 batch.save()
                 batch_id = batch.id
 
-            # 3. 分批处理UID数据(每批500条)
+            # 分批处理UID
             batch_size = 500
             processed = 0
             success_count = 0
             uids_batch = []
 
-            for row in ws.iter_rows(min_row=1, values_only=True):
-                if row[0]:
-                    uid = str(row[0]).strip()
-                    if uid:
-                        uids_batch.append(uid)
-                        processed += 1
-
-                        # 每处理1000条更新一次进度
-                        if processed % 1000 == 0:
-                            progress = min(99, int((processed / total_rows) * 100))
-                            task_data['progress'] = progress
-                            task_data['processed'] = processed
-                            task_data['last_update'] = int(time.time())
-                            redis_obj.set_data(redis_key, json.dumps(task_data))
-
-                        # 处理批次
-                        if len(uids_batch) >= batch_size:
-                            success = cls._import_uids_batch(
-                                uids_batch,
-                                batch_id,
-                                current_time,
-                                redis_key
-                            )
-                            success_count += success
-                            uids_batch = []
+            for uid in uid_rows:
+                uids_batch.append(uid)
+                processed += 1
+
+                # 每处理1000条更新进度
+                if processed % 1000 == 0:
+                    progress = min(99, int((processed / len(uid_rows)) * 100))
+                    task_data['progress'] = progress
+                    task_data['processed'] = processed
+                    task_data['last_update'] = int(time.time())
+                    redis_obj.set_data(redis_key, json.dumps(task_data))
+
+                if len(uids_batch) >= batch_size:
+                    success = cls._import_uids_batch(
+                        uids_batch,
+                        batch_id,
+                        current_time,
+                        redis_key
+                    )
+                    success_count += success
+                    uids_batch = []
 
-            # 处理最后一批
+            # 最后一批处理
             if uids_batch:
                 success = cls._import_uids_batch(
                     uids_batch,
@@ -371,11 +385,10 @@ class UIDBurnManageView(View):
                 )
                 success_count += success
 
-            # 更新批次总UID
+            # 更新数据库记录
             with transaction.atomic():
                 BurnBatch.objects.filter(id=batch_id).update(total_uid=success_count)
 
-            # 更新最终状态
             task_data['status'] = 'completed'
             task_data['progress'] = 100
             task_data['processed'] = processed
@@ -383,16 +396,15 @@ class UIDBurnManageView(View):
             task_data['end_time'] = int(time.time())
             redis_obj.set_data(redis_key, json.dumps(task_data))
 
-            LOGGER.info(f"处理批量导入完成,任务ID: {task_id}, 处理UID数量: {processed}, 成功导入: {success_count}")
+            LOGGER.info(f"批次导入完成,任务ID: {task_id},处理UID: {processed},成功导入: {success_count}")
 
-            # 清理临时文件
             try:
                 os.remove(file_path)
             except Exception as e:
                 LOGGER.warning(f"删除临时文件失败: {str(e)}")
 
         except Exception as e:
-            LOGGER.error(f"处理批量导入失败: {str(e)}")
+            LOGGER.error(f"处理导入任务失败: {str(e)}")
             task_data = {
                 'status': 'failed',
                 'error': str(e),