sbt-idp/cope2n-api/fwd_api/management/commands/migrate-database-fill-invoice-no.py

# myapp/management/commands/mycustomcommand.py
from django.core.management.base import BaseCommand
from tqdm import tqdm
from fwd_api.models import SubscriptionRequestFile, SubscriptionRequest
from fwd_api.exception.exceptions import InvalidException
from fwd_api.utils.accuracy import predict_result_to_ready
import traceback
import copy
from django.utils import timezone

KEY = "invoice_no"

class Command(BaseCommand):
    help = 'Refactor database for image level'

    def add_arguments(self, parser):
        # Add your command-line arguments here
        parser.add_argument('start', type=str, help='start date, sample: 2023-01-02T00:00:00+0700')
        parser.add_argument('end', type=str, help='end date, sample: 2023-01-03T00:00:00+0700')

    def process_request(self, request):
        if len(request.request_id.split(".")[0].split("_")) < 2:
            return
        images = SubscriptionRequestFile.objects.filter(request=request)
        
        if not request.predict_result:
            self.stdout.write(self.style.WARNING(f"Key predict_result not found in {request.request_id}"))
            return
        if request.predict_result.get("status", 200) != 200:
            self.stdout.write(self.style.WARNING(f"Key predict_result not found in {request.request_id}"))
            return
        
        if isinstance(request.predict_result, dict):
            invoice_no_exited = False
            for value in request.predict_result["content"]["document"][0]["content"]:
                if value["label"] == "invoice_no":
                    invoice_no_exited = True
            if not invoice_no_exited:
                request.predict_result["content"]["document"][0]["content"].insert(4, {"label": "invoice_no","value": None})
        if isinstance(request.feedback_result, dict):
            request.feedback_result[KEY] = None
        if isinstance(request.reviewed_result, dict):
            _reviewed_result = copy.deepcopy(request.reviewed_result)
            _reviewed_result[KEY] = None
            request.reviewed_result = _reviewed_result
        request.save()
        
        for i, image in enumerate(images):
            try:
                if isinstance(image.predict_result, dict):
                    # print(f"[INFO]: dict ... {request.predict_result}")
                    image.predict_result[KEY] = None
                if isinstance(image.feedback_result, dict):
                    image.feedback_result[KEY] = None
                if isinstance(image.reviewed_result, dict):
                    image.reviewed_result[KEY] = None
                image.save()
            except Exception as e:
                self.stdout.write(self.style.ERROR(f"Request: {request.request_id} failed with {e}"))
                print(traceback.format_exc())
                continue

    def handle(self, *args, **options):
        start = options['start']
        end = options['end']

        if start or end:
            try:
                start_date = timezone.datetime.strptime(start, '%Y-%m-%dT%H:%M:%S%z') # We care only about day precision only
                end_date = timezone.datetime.strptime(end, '%Y-%m-%dT%H:%M:%S%z')
            except Exception as e:
                print(f"[INFO]: start: {start}")
                print(f"[INFO]: end: {end}")
                raise InvalidException(excArgs="Date format")
            subcription_iter = SubscriptionRequest.objects.filter(created_at__range=(start_date, end_date))
        else:
            subcription_iter = SubscriptionRequest.objects.all()

        for request in tqdm(subcription_iter.iterator()):
            self.process_request(request)
        self.stdout.write(self.style.SUCCESS('Sample Django management command executed successfully!'))