lhyt_shop/utils/custom_modelviewset.py

# coding=utf-8

from rest_framework.viewsets import ModelViewSet

from utils import response_error, response_ok
from utils.exceptions import CustomError
from rest_framework.permissions import IsAuthenticated

from django.db import transaction

class CustomModelViewSet(ModelViewSet):
    permission_classes = [IsAuthenticated, ]
    def create(self, request, *args, **kwargs):
        try:
            with transaction.atomic():
                super(CustomModelViewSet, self).create(request, *args, **kwargs)
        except CustomError as e:
            return response_error(e.get_error_msg())
        except Exception as e:
            return response_error(str(e))
        return response_ok()      

    def update(self, request, *args, **kwargs):
        try:
            with transaction.atomic():
                super(CustomModelViewSet, self).update(request, *args, **kwargs)
        except CustomError as e:
            return response_error(e.get_error_msg())
        except Exception as e:
            return response_error(str(e))
        return response_ok()        

    def destroy(self, request, *args, **kwargs):
        try:
            with transaction.atomic():
                super(CustomModelViewSet, self).destroy(request, *args, **kwargs)
        except CustomError as e:
            return response_error(e.get_error_msg())
        except Exception as e:
            return response_error(str(e))
        return response_ok()