sbt-idp/cope2n-api/fwd_api/management/commands/migrate-database-010224.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

RETAILER_LIST = [
'',
'Amazon.sg',
'Gain City (TV/AV)',
'Harvey Norman (TV/AV)',
'KRIS SHOP',
'Lazada (Samsung Brand Store)',
'M1 Shop',
'Mohamed Mustafa & Samsuddin CO (TV/AV)',
'Parisilk (TV/AV)',
'Shopee (Samsung Brand Store)',
'StarHub Shop',
'M1 Shop',
'Ectri',
'Whandising',
'3 Mobile',
'Mister Mobile',
'Lazada',
'Mister Mobile',
'Samsung Experience Store',
'A & Samsuddin Co.',
'Parisilk',
'Samsung Brand Store',
'Shopee',
'M1 Shop',
'Onephone Online',
'3 Mobile',
'Samsung Experience Store',
'Challenger',
'Eas Marketing',
'Ing Mobile',
'Starhub Shop',
'Mister Mobile',
'Onephone Online',
'Starho'
]
RETAILER_LIST = list(set(RETAILER_LIST))

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)
        time_cost = {"imei": [], "invoice": [], "all": []}
        if request.ai_inference_profile is None:
            time_cost["imei"] = [-1 for _ in range(len(images))]
            time_cost["invoice"] = [-1]
            time_cost["all"] = [-1]
        else:
            for k, v in request.ai_inference_profile.items():
                time_cost[k.split("_")[0]].append(v["inference"][1][0] - v["inference"][0] + (v["postprocess"][1]-v["postprocess"][0]))
        for i, image in enumerate(images):
            # temp_imei_SAP_20240127223644_a493434edbf84fc08aeb87ef6cdde102_0.jpg
            try:
                # image.index_in_request = int(image.file_name.split(".")[0].split("_")[-1]) if len(image.file_name.split(".")[0].split("_")) > 4 else 0
                image.doc_type = image.file_name.split(".")[0].split("_")[1] if len(image.file_name.split(".")[0].split("_")) > 4 else "all"
                image.processing_time = time_cost[image.doc_type][image.index_in_request]
                if not request.predict_result:
                    raise KeyError(f"Key predict_result not found in {request.request_id}")
                if request.predict_result.get("status", 200) != 200:
                    raise AttributeError(f"Failed request: {request.request_id}")
                _predict_result = copy.deepcopy(predict_result_to_ready(request.predict_result))
                _feedback_result = copy.deepcopy(request.feedback_result)
                _reviewed_result = copy.deepcopy(request.reviewed_result)

                if image.doc_type == "invoice":
                    _predict_result["imei_number"] = []
                    if _feedback_result:
                        _feedback_result["imei_number"] = []
                    else:
                        None
                    if _reviewed_result:
                        _reviewed_result["imei_number"] = []
                    else:
                        None
                else:
                    _predict_result = {"retailername": None, "sold_to_party": None, "purchase_date": [], "imei_number": [_predict_result["imei_number"][image.index_in_request]]}
                    _feedback_result = {"retailername": None, "sold_to_party": None, "purchase_date": None, "imei_number": [_feedback_result["imei_number"][image.index_in_request]]} if _feedback_result else None
                    if isinstance(_reviewed_result, dict) and (len(_reviewed_result.get("imei_number", [])) or any(element in RETAILER_LIST for element in _reviewed_result.get("imei_number", []))) :
                        _reviewed_result = {"retailername": None, "sold_to_party": None, "purchase_date": None, "imei_number": [_reviewed_result["imei_number"][image.index_in_request]]} if _reviewed_result else None
                    else:
                        _reviewed_result = None
                        request.reviewed_result = None
                        request.is_reviewed = False
                        request.save()
                image.predict_result = _predict_result
                image.feedback_result = _feedback_result
                image.reviewed_result = _reviewed_result
                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']

        white_list = [
            'SAP_20240202130136_592fc76aca3248d596e2bedcc4340fee',
            'SAP_20240202164711_b07f7acb723946d595f6184b70a45387',
            'SAP_20240203210818_8c6f6dd2ae42434ab40c690022a1ff6e',
            'SAP_20240203141331_c64f71dc006d437dae558c6523de7e7c',
            'SAP_20240201212124_55d135704e354cf5bb029a6ee1501f9f',
            'SAP_20240202214418_c06c1669ce7743bb85d0f999bf022b8d',
            'SAP_20240203033349_0e3db9ca289f4c9db96a6a07feba0597',
            'SAP_20240203195254_d899295781474d24a36fee4d5cad0fd1',
            'SAP_20240203192612_d470b4300528481a886a4629d5e538ab',
            'SAP_20240203202654_a2399fb471b942dcb9f27a1bb09f026f',
            'SAP_20240131093605_3aec91f880234096932059b515e31bc2',
            'SAP_20240202164256_9b1d8c5531514b6dbe9857b33ef1dd93',
            'SAP_20240202194336_78d55dfff1684368b32405eb3bb16800',
            'SAP_20240202194901_8b86c75484f14411be7ddd63b0f71abe',
            'SAP_20240204140236_91b02047d6124224b98c1bca7ef5808a',
            'SAP_20240204140236_91b02047d6124224b98c1bca7ef5808a',
            'SAP_20240205085813_05939ccf8d6a44068a7394d40360c957',
            'SAP_20240205085813_05939ccf8d6a44068a7394d40360c957',
            'SAP_20240205121225_27d61c27da0445a5963f6dcf8bebc169',
            'SAP_20240205133855_3c355550a1d7449aab04176c70b0338d',
            'SAP_20240205133855_3c355550a1d7449aab04176c70b0338d',
            'SAP_20240205234354_08e4fbba404f440ba87ee08394d11a93',
            'SAP_20240206001800_e2e326de5adc4a12b9c7f713b96c70b9',
            'SAP_20240202193951_17a709bb3f4c4c3a9ef34b777fcd8b0d',
            'SAP_20240204161121_6e80153cf05f45b5aa46d08fd52ed454',
            'SAP_20240131181202_94dc906cd86e4a89b08386c5a737f843',
            'SAP_20240203175643_6a19fbd9839643219fd21af2d46ff67b',
            'SAP_20240131211332_7f937cfc26614230b43e7e658a7e9002',
            'SAP_20240131215607_6c02781a83804fce8ee4c44e8aade1e6',
            'SAP_20240131232749_5b8cfe62bb5847baa5c2e8e006d8fd4c',
            'SAP_20240201143938_a4da2d697e6f4eab86a97344ab6ef612',
            'SAP_20240201222839_9711372c80a84df6aefd2cf494f881a6',
            'SAP_20240201223812_e55a3ffd4e0940b9a3d1ab4058b88006',
            'SAP_20240202052454_35dedef1afc44cf3a2ee5e8f03a2e871',
            'SAP_20240202123705_10ed0532243d4fad971ab4ea9982848a',
            'SAP_20240202160851_e62bf43cecc24ff6988e41b59727b60a',
            'SAP_20240202173309_2bcaa76aadfd4f8c87f72a64344fc818',
            'SAP_20240202174053_2e5b97362b14492b91235e9a259bfa42',
            'SAP_20240202222341_f132f1fde9974bdd87a734c04c5b91a8',
            'SAP_20240203105436_38192f4a3c2745828b7bab702f73b5a8',
            'SAP_20240203112502_ca87267f28be4538b5cdc8c49cd951fe',
            'SAP_20240203141104_a5c879a0f7d94e36822d0f10f5f446be',
            'SAP_20240203143937_3730bb0fa62f4ea792dbaf6b75c190fa',
            'SAP_20240203144238_d5ef7897e23b4bbc954854746c315f09',
            'SAP_20240203151137_4de4ed9e1cd94c60af3e6588ab418cab',
            'SAP_20240203153540_9214a78c5a2142ccbf18dea79a7f7d32',
            'SAP_20240203174245_d463b37dce2a48b48cff9b10333da146',
            'SAP_20240203185611_4892b09d22f84418b4472bbbcab554f4',
            'SAP_20240203192059_82bcb05bc9cd48ab94c724783a78995a',
            'SAP_20240203200323_7e1944a9e6f3477793a7d8682efe12a8',
            'SAP_20240203202429_780d51302e2a402683d0e8b9f7a5206c',
            'SAP_20240203213530_e9e38035733348a1a3050b269a988ed4',
            'SAP_20240203220513_6939d446931044feaeb0ea39a6409024',
            'SAP_20240203224040_476f975502084050b829397c0d543626',
            'SAP_20240203230738_f4343b5dbc8d4fac98ee08c8b631e738',
            'SAP_20240204013646_d7d89fe2d22a4e5f8575ba3af0785e31',
            'SAP_20240204081015_4c7f0d11a24a4bcebfb1adc50ef4f34a',
            'SAP_20240204094056_66c2606c43a44dde9879467556aed70a',
            'SAP_20240204110605_aa45d3f72c72465190b083a8a2ded18a',
            'SAP_20240204123216_658c952721be43f9b4d107a38aa396e1',
            'SAP_20240204123216_658c952721be43f9b4d107a38aa396e1',
            'SAP_20240204140721_30e6267e4aed4597b43e10320fa5f75c',
            'SAP_20240204163159_b34afc9534cd41ea93aebf000bfb0fb4',
            'SAP_20240204163832_c7051196f83e4f68bc09668ea273370b',
            'SAP_20240204200851_4597724b85f74907b396658642c950e1',
            'SAP_20240204201303_156a0bf95ac2434083f9ff519d215b8a',
            'SAP_20240205122953_b2f78a10898048a09d21779372d910be',
            'SAP_20240205185633_14afd3692c5c47f4a228689c89268595',
            'SAP_20240205200646_f9b7bf31903649d1be41996174047b57',
            'SAP_20240205204209_9b8ee17983404658a0121717945c7ea5',
            'SAP_20240205235953_b92aaccd88214282880ec70034fbe0fc',
            'SAP_20240203180224_7bbf6e08a37f4a968a27852b0f9348a4',
            'SAP_20240204095520_f306f5b596db4c788f49d24da909ba3b',
            'SAP_20240204131511_ba42a26ce05e4d279f3c4d956723ce0c',
            'SAP_20240204153410_469368284b884c29b1213ea976956891',
            'SAP_20240129173923_7809ec94fa014fd88406f954a5b85c65',
            'SAP_20240131173059_3afe5da631fa44ec969baf8974092641',
            'SAP_20240202213429_737af53cb1194c31a50094361e3090f2',
            'SAP_20240203001712_5568f19671b941aab2fea2d93f6b688f',
            'SAP_20240203155754_882632f2d2b640bba9a6ea246cbace0e',
            'SAP_20240204090000_7272e17a7c4d468c8bb074c69a67edc6',
            'SAP_20240205140231_5a5be97738a648dda0520b24a1957336',
            'SAP_20240205172847_c36b8542b66148f6b473259438ca5280',
            'SAP_20240205173225_e5fe5e9db09b4eee92a0b044eecea233',
            'SAP_20240205183822_95dd07be90874db0ac1827dfeb6317e2',
        ]

        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')
                # start_date = timezone.make_aware(start_date)
                # end_date = timezone.make_aware(end_date)
            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()):
            if request.request_id not in white_list:
                continue
            self.process_request(request)
        self.stdout.write(self.style.SUCCESS('Sample Django management command executed successfully!'))