【Django】QuerySetをcsvに出力するコマンド

大量のレコードが格納されているテーブルにselect分を発行して、結果をcsvに出力し、gzipする、というコードを書いたのでメモ。

models.py

from logging import getLogger
from django.db import models
import uuid

logger = getLogger(__name__)


class Analytics(models.Model):

    hit_type = models.CharField(max_length=20)

    category = models.CharField(max_length=255, null=True, blank=True)
    action = models.CharField(max_length=255, null=True)
    label = models.CharField(max_length=255, null=True, blank=True)
    value = models.IntegerField(null=True)
    url = models.CharField(max_length=255, null=True)
    created = models.DateTimeField(auto_now_add=True, db_index=True)
    tracking_user = models.ForeignKey('app.TrackingUser', related_name='analytics')

    def __str__(self):
        return '%s - %s' % (self.category, self.tracking_user.uuid)


class TrackingUser(models.Model):

    uuid = models.UUIDField(primary_key=True, default=uuid.uuid4, editable=False)
    created = models.DateTimeField(auto_now_add=True)

    def __str__(self):
        return str(self.uuid)

上記のようなウェブサイト上のユーザの行動をトラッキングしたデータを格納するModelがあったんだけれど、今回はAnalyticsのレコードを1ヵ月分取得し、csvに格納したあと、deleteできるコマンドを作りたかった。
ちなみに、Analyticsのデータは1日で10万件ほど増えていくので、1ヵ月分だと300万件くらい。
んで、色々やって見た結果、以下のコードが良い感じだったッス。

delete_analytics.py

from datetime import datetime, timedelta, date
from api.models import Analytics
from pytz import timezone, utc
from django.core.management.base import BaseCommand
from django.db import transaction
from dateutil.relativedelta import relativedelta
import csv
import gzip
import shutil
import os

class Command(BaseCommand):
    def add_arguments(self, parser):
        # Named (optional) arguments
        parser.add_argument('--target_date', metavar='target_date', type=str, nargs=None,
                            help='target date')

    @transaction.atomic
    def handle(self, *args, **options):
        if options['target_date']:
            delete_start_date = datetime.strptime(options['target_date'], '%Y-%m-%d').replace(day=1)
        else:
            this_month = date.today().replace(day=1)
            delete_start_date = this_month - relativedelta(months=3)
        delete_end_date = delete_start_date + relativedelta(months=1)

        queryset = Analytics.objects.select_related('user', 'tracking_user').filter(created__range=(delete_start_date, delete_end_date))

        # csvへの出力
        filename_date = delete_start_date.strftime('%Y%m')
        filename = filename_date + "_event.csv"

        model = queryset.model
        writer = csv.writer(open(filename, 'w'))

        ## csvのヘッダー部の書き込み
        headers = []
        for field in model._meta.fields:
            headers.append(field.name)
        writer.writerow(headers)

        ## 取得したレコードを書き込む
        for obj in queryset:
            row = []
            for field in headers:
                val = getattr(obj, field)
                if callable(val):
                    val = val()
                row.append(val)
            writer.writerow(row)

        # csvをgzipで圧縮
        with open(filename, 'rb') as gzip_in:
            with gzip.open(filename + ".gz", 'wb') as gzip_out:
                shutil.copyfileobj(gzip_in, gzip_out)

        os.remove(filename)

        queryset.delete()

manage.pyがあるディレクトリで./manage.py delete_analytics.py -target_date=2017-06-01とすると、2017年6月分のデータが201706_event.csvに出力されて、さらにgzipで圧縮、その後、deleteされる。


今回キモだったのは以下の部分。

queryset = Analytics.objects.select_related( 'tracking_user').filter(created__range=(delete_start_date, delete_end_date))

select_related('tracking_user').を書かないとリレーション先もselectしにいってしまうので時間がかかってしまう。
書いた場合は5分でコマンド終了したけれど、書かなかったら終了までに50分かかったのdeath。

追記

サーバ上で動かしたらメモリ不足で落ちてしまった...
何か対処を考えて、対応できたらまた書きます

参考リンク

django recipe: dump your queryset out as a csv file . palewire