[Python] DRF 대용량 CSV 데이터 안전하게 만들기

Django에서 만들어져있는 CSV 파일을 스트리밍으로 내려주는 방법은 FileResponse 또는 StreamingHttpResponse 클래스를 쓰면 간단하게 구현이 가능하다. 하지만 현시점의 데이터를 여러 조건으로 필터링하여 CSV 파일을 만들어서 내려주는 경우 CPU 100% + OOM이 발생하여 서버가 아무 일도 할 수 없게되는 현상이 발생할 수 있다. 이런 구현은 동시에 1개의 요청도 처리하기 어려운데 여러 다운로드 요청이 오게 되면 더 빠른 속도로 서버를 조질수 있다.

요청 하나로 서버 터지는 구현

스트리밍으로 주는데 왜 터지지? 라는 생각이 들 수 있는데, REST framework의 serializer.data에 접근하면 queryset을 수행하고 모든 결과를 수집한 이후에야 반복문이 돌기 시작하기 때문이다.

class CSVDownloadViewSet(viewsets.GenericViewSet):
    def create_stream(self, queryset):
        for data in self.get_serializer(querset, many=True).data:  # 여기서 메모리가 가득 차게된다.
            ...
            yield data

    @action(...)
    def download(request, *args, **kwargs):
        return StreamingHttpResponse(
            self.create_stream(self.filter_queryset(self.get_queryset())),
            ...  # Header 추가
        )

Offset - Limit

단점: 데이터 양이 많아지면 많아질수록 다운로드 속도가 매우매우 느려진다.

참고 - https://jojoldu.tistory.com/528

class CSVDownloadViewSet(viewsets.GenericViewSet):
    CHUNK_SIZE = 1000  # 동적으로 변경 가능하게 하는 구현도 좋다.  EX) django-constance

    def create_stream(self, queryset):
        curr, end = queryset.aggregate(Min("id"), Max("id")).values()
        while curr < end:
            yield from _create_stream(queryset[curr:self.CHUNK_SIZE])
            curr += self.CHUNK_SIZE

    def _create_stream(self, queryset):
        for data in self.get_serializer(queryset, many=True).data:
            ...
            yield data

    @action(...)
    def download(request, *args, **kwargs):
        return StreamingHttpResponse(
            self.create_stream(self.filter_queryset(self.get_queryset())),
            ...  # Header 추가
        )

No Offset

Offset - Limit의 단점을 개선한 구현

class CSVDownloadViewSet(viewsets.GenericViewSet):
    CHUNK_SIZE = 1000  # 동적으로 변경 가능하게 하는 구현도 좋다.  EX) django-constance

    def create_stream(self, queryset):
        queryset = queryset.order_by("id")
        curr, end = queryset.aggregate(Min("id"), Max("id")).values()
        while curr < end:
            yield from self._create_stream(queryset.filter(id__gte=curr)[: self.CHUNK_SIZE])
            curr += self.CHUNK_SIZE

    def _create_stream(self, queryset):
        for data in self.get_serializer(queryset, many=True).data:
            ...
            yield data

    @action(...)
    def download(request, *args, **kwargs):
        return StreamingHttpResponse(
            self.create_stream(self.filter_queryset(self.get_queryset())),
            ...  # Header 추가
        )

추가

  • CSV를 만들 때 데이터 안에 ,, "가 들어가 있는 경우 "'로 변경하고 데이터를 "로 감싸주면 데이터가 옆으로 밀리는 현상을 개선할 수 있다.
  • 중간에 ID가 비어있는 경우가 많은 경우 chunk queryset의 가장 큰 ID 값을 curr에 넣어주면 된다.

Written by@EHX
Software Developer, Back-End Engineer

GitHubFacebook