edr-platform/apps/edr-freight-api/src/migrations/1791000000001-AddCompanyProfileIdToBookings.ts

import { MigrationInterface, QueryRunner } from 'typeorm';

export class AddCompanyProfileIdToBookings1791000000001
  implements MigrationInterface
{
  name = 'AddCompanyProfileIdToBookings1791000000001';

  public async up(queryRunner: QueryRunner): Promise<void> {
    await queryRunner.query(`
      ALTER TABLE freight.bookings
        ADD COLUMN IF NOT EXISTS company_profile_id UUID;
    `);

    await queryRunner.query(`
      CREATE INDEX IF NOT EXISTS idx_bookings_company_profile_id
        ON freight.bookings(company_profile_id);
    `);

    await queryRunner.query(`
      DO $$
      BEGIN
        IF NOT EXISTS (
          SELECT 1 FROM pg_constraint WHERE conname = 'FK_bookings_company_profile_id'
        ) THEN
          ALTER TABLE freight.bookings
            ADD CONSTRAINT "FK_bookings_company_profile_id"
            FOREIGN KEY (company_profile_id)
            REFERENCES freight.company_profiles(id);
        END IF;
      END $$;
    `);

    // Backfill by natural mapping: IMPORT → importer profile, EXPORT → exporter
    // profile, for each booking's own company.
    await queryRunner.query(`
      UPDATE freight.bookings b
        SET company_profile_id = cp.id
        FROM freight.company_profiles cp
        WHERE cp.company_id = b.company_id
          AND b.company_profile_id IS NULL
          AND (
            (b.trade_direction = 'IMPORT' AND cp.type = 'importer') OR
            (b.trade_direction = 'EXPORT' AND cp.type = 'exporter')
          );
    `);

    // Forwarder / single-profile companies: one profile per company, so the
    // mapping is unambiguous regardless of trade direction.
    await queryRunner.query(`
      UPDATE freight.bookings b
        SET company_profile_id = cp.id
        FROM freight.company_profiles cp
        JOIN freight.companies c ON c.id = cp.company_id
        WHERE cp.company_id = b.company_id
          AND c.type <> 'customer'
          AND b.company_profile_id IS NULL;
    `);

    // Remaining customer-owned rows (e.g. DOMESTIC, or a direction with no
    // matching profile): attribute to the company's importer profile, else its
    // exporter profile, so nothing disappears from the customer's list.
    await queryRunner.query(`
      UPDATE freight.bookings b
        SET company_profile_id = cp.id
        FROM (
          SELECT DISTINCT ON (company_id) company_id, id
            FROM freight.company_profiles
            ORDER BY company_id,
              CASE type
                WHEN 'importer' THEN 0
                WHEN 'exporter' THEN 1
                ELSE 2
              END
        ) cp
        WHERE cp.company_id = b.company_id
          AND b.company_id IS NOT NULL
          AND b.company_profile_id IS NULL;
    `);
  }

  public async down(queryRunner: QueryRunner): Promise<void> {
    await queryRunner.query(`
      ALTER TABLE freight.bookings
        DROP CONSTRAINT IF EXISTS "FK_bookings_company_profile_id";
    `);
    await queryRunner.query(`
      DROP INDEX IF EXISTS freight.idx_bookings_company_profile_id;
    `);
    await queryRunner.query(`
      ALTER TABLE freight.bookings
        DROP COLUMN IF EXISTS company_profile_id;
    `);
  }
}