From 63a56390a6ffee9f68478aa392ff03659e78e084 Mon Sep 17 00:00:00 2001
From: Diego Escobedo <diego@escobedo.com>
Date: Thu, 10 Nov 2022 15:14:41 -0800
Subject: [PATCH] migration fail retry with exception handling

---
 .../migrations/0069_transfer_invoice_to_fk.py | 86 +++++++++++--------
 ...historicalinvoice_old_customer_and_more.py | 37 --------
 2 files changed, 48 insertions(+), 75 deletions(-)
 delete mode 100644 backend/metering_billing/migrations/0070_remove_historicalinvoice_old_customer_and_more.py

diff --git a/backend/metering_billing/migrations/0069_transfer_invoice_to_fk.py b/backend/metering_billing/migrations/0069_transfer_invoice_to_fk.py
index 7becd7071..f518331aa 100644
--- a/backend/metering_billing/migrations/0069_transfer_invoice_to_fk.py
+++ b/backend/metering_billing/migrations/0069_transfer_invoice_to_fk.py
@@ -15,48 +15,58 @@ def migrate_jsonfields_to_fk(apps, schema_editor):
     for invoice in Invoice.objects.all():
         try:
             invoice_org = invoice.old_organization["company_name"]
+        except KeyError:
+            invoice_org = None
+        try:
+            invoice_customer = invoice.old_customer["customer_id"]
+        except:
+            invoice_customer = None
+        try:
+            invoice_sub = invoice.old_subscription["subscription_id"]
         except:
-            print(invoice.__dict__)
-            raise
-        invoice_customer = invoice.old_customer["customer_id"]
-        invoice_sub = invoice.old_subscription["subscription_id"]
+            invoice_sub = None
 
-        if invoice_org in invoice_org_dict:
-            invoice_org_object = invoice_org_dict[invoice_org]
-        else:
-            try:
-                invoice_org_object = Organization.objects.get(company_name=invoice_org)
-            except Organization.DoesNotExist:
-                invoice_org_object = None
-                print(f"Invoice {invoice.pk} has no organization")
-            invoice_org_dict[invoice_org] = invoice_org_object
-        invoice.organization = invoice_org_object
+        if invoice_org:
+            if invoice_org in invoice_org_dict:
+                invoice_org_object = invoice_org_dict[invoice_org]
+            else:
+                try:
+                    invoice_org_object = Organization.objects.get(
+                        company_name=invoice_org
+                    )
+                except Organization.DoesNotExist:
+                    invoice_org_object = None
+                    print(f"Invoice {invoice.pk} has no organization")
+                invoice_org_dict[invoice_org] = invoice_org_object
+            invoice.organization = invoice_org_object
 
-        if invoice_customer in invoice_customer_dict:
-            invoice_customer_object = invoice_customer_dict[invoice_customer]
-        else:
-            try:
-                invoice_customer_object = Customer.objects.get(
-                    organization=invoice_org_object, customer_id=invoice_customer
-                )
-            except:
-                invoice_customer_object = None
-                print(f"Invoice {invoice.pk} has no customer")
-            invoice_customer_dict[invoice_customer] = invoice_customer_object
-        invoice.customer = invoice_customer_object
+        if invoice_customer:
+            if invoice_customer in invoice_customer_dict:
+                invoice_customer_object = invoice_customer_dict[invoice_customer]
+            else:
+                try:
+                    invoice_customer_object = Customer.objects.get(
+                        organization=invoice_org_object, customer_id=invoice_customer
+                    )
+                except:
+                    invoice_customer_object = None
+                    print(f"Invoice {invoice.pk} has no customer")
+                invoice_customer_dict[invoice_customer] = invoice_customer_object
+            invoice.customer = invoice_customer_object
 
-        if invoice_sub in invoice_sub_dict:
-            invoice_sub_object = invoice_sub_dict[invoice_sub]
-        else:
-            try:
-                invoice_sub_object = Subscription.objects.get(
-                    organization=invoice_org_object, subscription_id=invoice_sub
-                )
-            except:
-                invoice_sub_object = None
-                print(f"Invoice {invoice.pk} has no subscription")
-            invoice_sub_dict[invoice_sub] = invoice_sub_object
-        invoice.subscription = invoice_sub_object
+        if invoice_sub:
+            if invoice_sub in invoice_sub_dict:
+                invoice_sub_object = invoice_sub_dict[invoice_sub]
+            else:
+                try:
+                    invoice_sub_object = Subscription.objects.get(
+                        organization=invoice_org_object, subscription_id=invoice_sub
+                    )
+                except:
+                    invoice_sub_object = None
+                    print(f"Invoice {invoice.pk} has no subscription")
+                invoice_sub_dict[invoice_sub] = invoice_sub_object
+            invoice.subscription = invoice_sub_object
 
         invoice.save()
 
diff --git a/backend/metering_billing/migrations/0070_remove_historicalinvoice_old_customer_and_more.py b/backend/metering_billing/migrations/0070_remove_historicalinvoice_old_customer_and_more.py
deleted file mode 100644
index f8469875b..000000000
--- a/backend/metering_billing/migrations/0070_remove_historicalinvoice_old_customer_and_more.py
+++ /dev/null
@@ -1,37 +0,0 @@
-# Generated by Django 4.0.5 on 2022-11-10 21:22
-
-from django.db import migrations
-
-
-class Migration(migrations.Migration):
-
-    dependencies = [
-        ("metering_billing", "0069_transfer_invoice_to_fk"),
-    ]
-
-    operations = [
-        migrations.RemoveField(
-            model_name="historicalinvoice",
-            name="old_customer",
-        ),
-        migrations.RemoveField(
-            model_name="historicalinvoice",
-            name="old_organization",
-        ),
-        migrations.RemoveField(
-            model_name="historicalinvoice",
-            name="old_subscription",
-        ),
-        migrations.RemoveField(
-            model_name="invoice",
-            name="old_customer",
-        ),
-        migrations.RemoveField(
-            model_name="invoice",
-            name="old_organization",
-        ),
-        migrations.RemoveField(
-            model_name="invoice",
-            name="old_subscription",
-        ),
-    ]