优化Django自定义管理命令批量填充多对多关联表的方法

问题背景 #

我正在学习Python和Django，已完成从CSV文件导入大量数据至数据库的工作，但其中一份描述多对多关联关系的CSV导入效率极低——6000行数据耗时约6分钟。

模型定义如下：

# models.py
class Provider(models.Model):
    name = models.CharField(max_length=25)
    product = models.ManyToManyField("Product", related_name="vendor_products")

class Product(models.Model):
    name = models.CharField(max_length=25)

Django自动生成的中间表为core_provider_product，CSV数据格式如下：

provider_id,product_id
3,1001
1,1003
3,1007
2,1001
1,1002
1,1001
# ... 更多数据

当前使用的自定义管理命令虽能运行但效率极低：

class Command(BaseCommand):
    # add_arguments 代码省略

    def handle(self, *args, **options):
        for file in options["csvfile"]:
            with open(file) as f:
                reader = csv.reader(f)
                for row in reader:
                    myprovider = row[0]
                    myprovider_obj = Provider.objects.get(pk=myprovider)
                    myproduct = row[1]
                    myproduct_obj = Product.objects.get(pk=myproduct)
                    myprovider_obj.product.add(myproduct_obj)

优化思路与实现 #

核心问题分析 #

当前代码每循环一行就执行2次数据库get查询，6000行数据会产生12000次独立查询，这是性能瓶颈的核心；同时逐行调用add方法会频繁触发数据库事务提交，进一步拖慢速度。

优化方案（直接操作中间表+批量处理） #

from django.db import transaction
import csv

class Command(BaseCommand):
    # add_arguments 代码省略

    def handle(self, *args, **options):
        for file in options["csvfile"]:
            with open(file) as f:
                reader = csv.DictReader(f)
                # 收集所有需要的ID与原始数据
                provider_ids = set()
                product_ids = set()
                valid_associations = []
                
                for row in reader:
                    try:
                        provider_id = int(row['provider_id'])
                        product_id = int(row['product_id'])
                    except ValueError:
                        self.stdout.write(self.style.WARNING(f"无效行数据：{row}"))
                        continue
                    
                    provider_ids.add(provider_id)
                    product_ids.add(product_id)
                    valid_associations.append((provider_id, product_id))

                # 批量查询所有存在的Provider和Product，避免重复查询
                existing_providers = Provider.objects.filter(id__in=provider_ids).values_list('id', flat=True)
                existing_products = Product.objects.filter(id__in=product_ids).values_list('id', flat=True)
                
                # 过滤掉关联对象不存在的无效数据
                filtered_associations = [
                    (p_id, pr_id) for p_id, pr_id in valid_associations
                    if p_id in existing_providers and pr_id in existing_products
                ]

                # 直接操作多对多中间表，批量插入数据
                through_model = Provider.product.through
                with transaction.atomic():
                    through_model.objects.bulk_create(
                        [
                            through_model(provider_id=p_id, product_id=pr_id)
                            for p_id, pr_id in filtered_associations
                        ],
                        batch_size=1000  # 根据数据库性能调整，MySQL建议1000-2000
                    )
                
                self.stdout.write(self.style.SUCCESS(f"成功导入{len(filtered_associations)}条有效关联数据"))

额外优化细节 #

• 使用transaction.atomic()包裹批量操作，减少数据库提交次数，同时保证数据原子性（要么全成功，要么全回滚）。
• 用csv.DictReader替代csv.reader，通过字段名访问数据，避免索引错误，提升代码可读性。
• 提前过滤无效数据（比如非整数ID、关联对象不存在的行），避免无效数据库操作。
• 若CSV文件极大，可分块读取处理，避免一次性加载所有数据占用过多内存。

内容的提问来源于stack exchange，提问作者givi