提高 Django 批量更新数据的性能

基于Django ORM的数据批量更新实现，目的：解决Django批量数据更新时的性能问题

问题

Django提供的原生批量更新API，示例代码如下：

Test.objects.filter(conditions='name').update(value='new name')

这个API在批量更新数据时有两个问题：

1. 频繁的获取数据库连接
2. 不支持数据更新条件的差异化

实现

在Java技术体系里JDBC提供的预编译语句和batch操作可以大大的提升数据操作性能，这里借鉴了JDBC思想的同时还参考了django-bulk-update，下面是一个简单的实现：

# coding: utf8

import itertools
import logging
from django.db import connections, models


def grouper(iterable, size):
    it = iter(iterable)
    while True:
        chunk = tuple(itertools.islice(it, size))
        if not chunk:
            return
        yield chunk


def gen_update_sql(obj):
    # 目前不支持组合主键更新
    meta = obj._meta
    pk_field = meta.get_field(meta.pk.name)
    db_table = meta.db_table
    conditions = '{pk} = %s'.format(pk=pk_field.column)
    fields = [(f.column, f.attname) for f in meta.concrete_fields if not f.primary_key]
    values = ', '.join(['%s = %s' % (f[0], '%s') for f in fields])
    fields.append((pk_field.column, pk_field.attname))
    update_sql = 'update {db_table} set {values} where {conditions}'.format(db_table=db_table, values=values, conditions=conditions)
    return update_sql, fields


def bulk_update(objs, batch_size=10, using='db_tag'):
    bool = False;
    if not objs:
        return bool
    connection = connections[using]
    try:
        with connection.cursor() as cursor:
            update_sql, fields = gen_update_sql(objs[0])
            params = []
            for objs_batch in grouper(objs, batch_size):
                params[:] = []
                for obj in objs_batch:
                    params.append([getattr(obj, f[1]) for f in fields])
                cursor.executemany(update_sql, params)
                connection.commit()
        bool = True
    except Exception as e:
        logging.error('batch update error, msg = %s', e.message)
    finally:
        connection.close()
    return bool

扩展

上述代码还有很大的扩展空间，如：支持更新字段的黑白名单；数据更新条件的丰富；数据分片（分库分表）等；

代码用到了Python几个比较犀利的特性：
1. 迭代器
2. 生成器
3. 推导式