migration fix for code review

bcgov · Mar 15, 2024 · 6b5cce7 · 6b5cce7
1 parent 4212d1c
commit 6b5cce7
Show file tree

Hide file tree

Showing 3 changed files with 18 additions and 5 deletions.
diff --git a/jobs/payment-jobs/tasks/ejv_partner_distribution_task.py b/jobs/payment-jobs/tasks/ejv_partner_distribution_task.py
@@ -30,7 +30,7 @@
 from pay_api.models import RefundsPartial as RefundsPartialModel
 from pay_api.models import Receipt as ReceiptModel
 from pay_api.models import db
-from pay_api.utils.enums import DisbursementStatus, EjvFileType, InvoiceStatus, PaymentMethod
+from pay_api.utils.enums import DisbursementStatus, EjvFileType, EJVLinkType, InvoiceStatus, PaymentMethod
 from sqlalchemy import Date, cast
 
 from tasks.common.cgi_ejv import CgiEjv
@@ -258,8 +258,8 @@ def _create_ejv_file_for_partner(cls, batch_type: str):  # pylint:disable=too-ma
 
             # Create ejv invoice/partial_refund link records and set invoice status
             sequence = 1
-            sequence = cls._process_items(invoices, ejv_header_model, sequence, 'invoice')
-            cls._process_items(refund_partial_items, ejv_header_model, sequence, 'refund')
+            sequence = cls._process_items(invoices, ejv_header_model, sequence, EJVLinkType.INVOICE.value)
+            cls._process_items(refund_partial_items, ejv_header_model, sequence, EJVLinkType.REFUND.value)
 
             db.session.flush()
 

diff --git a/pay-api/migrations/versions/2024_03_06_04b8a7bed74e_ejv_link_generic.py b/pay-api/migrations/versions/2024_03_06_04b8a7bed74e_ejv_link_generic.py
@@ -8,6 +8,7 @@
 from alembic import op
 import sqlalchemy as sa
 from sqlalchemy.dialects import postgresql
+from pay_api.utils.enums import EJVLinkType
 
 # revision identifiers, used by Alembic.
 revision = '04b8a7bed74e'
@@ -17,16 +18,21 @@
 
 
 def upgrade():
-    op.add_column('ejv_invoice_links', sa.Column('link_type', sa.String(length=20), nullable=False))
+    op.add_column('ejv_invoice_links', sa.Column('link_type', sa.String(length=20), nullable=True))
     op.drop_index(op.f('ix_ejv_invoice_links_invoice_id'), table_name='ejv_invoice_links')
     op.alter_column('ejv_invoice_links', 'invoice_id', new_column_name='link_id',
                     existing_type=sa.Integer(), nullable=True)
     op.rename_table('ejv_invoice_links', 'ejv_links')
+    # Update statement to set link_type to 'invoice' for all existing records as there is no data for partial refunds yet.
+    op.execute("UPDATE ejv_links SET link_type = :link_type", {'link_type': EJVLinkType.INVOICE.value})
+
+    # Index for link_type + link_id
+    op.create_index('ix_ejv_links_link_type_link_id', 'ejv_links', ['link_type', 'link_id'], unique=False)
 
 def downgrade():
+    op.drop_index('ix_ejv_links_link_type_link_id', table_name='ejv_links')
     op.rename_table('ejv_links', 'ejv_invoice_links')
     op.alter_column('ejv_invoice_links', 'link_id', new_column_name='invoice_id', 
                     existing_type=sa.Integer(), nullable=False)
     op.create_index(op.f('ix_ejv_invoice_links_invoice_id'), 'ejv_invoice_links', ['invoice_id'], unique=False)
     op.drop_column('ejv_invoice_links', 'link_type')
-
diff --git a/pay-api/src/pay_api/utils/enums.py b/pay-api/src/pay_api/utils/enums.py
@@ -375,3 +375,10 @@ class QueueSources(Enum):
     PAY_JOBS = 'pay-jobs'
     PAY_QUEUE = 'pay-queue'
     FTP_POLLER = 'ftp-poller'
+
+
+class EJVLinkType(Enum):
+    """EJV link types for ejv_link table."""
+
+    INVOICE = 'invoice'
+    REFUND = 'refund'