feat(train-sets): implement multi-locomotive support for train sets

- Added TrainSetLocomotive entity to link multiple locomotives to a train set. - Updated TrainSet entity to include a OneToMany relationship with TrainSetLocomotive. - Modified the train scheduling logic to require at least two locomotives for a train set. - Enhanced the UI components to support selecting multiple locomotives. - Introduced new permissions for viewing customs clearance. - Updated migrations to create the train_set_locomotives table and backfill existing data. - Implemented utility functions for managing train numbers based on cargo type and direction. - Added tests for train number utilities to ensure correct functionality.
2026-09-07 22:25:42 +00:00 · 2026-06-24 23:54:45 +00:00
parent 15ab9f906e
commit e0c3044933
28 changed files with 817 additions and 118 deletions
--- a/apps/edr-freight-api/src/modules/train-scheduling/train-scheduling.service.ts
+++ b/apps/edr-freight-api/src/modules/train-scheduling/train-scheduling.service.ts
@@ -22,6 +22,7 @@ import { Container } from '../container-management/entities/container.entity';
 import { Locomotive } from '../locomotives/entities/locomotive.entity';
 import { LocomotivesRepository } from '../locomotives/locomotives.repository';
 import { Route } from '../routes/entities/route.entity';
+import { TrainSetLocomotive } from '../train-sets/entities/train-set-locomotive.entity';
 import { TrainSetWagon } from '../train-sets/entities/train-set-wagon.entity';
 import { TrainSet } from '../train-sets/entities/train-set.entity';
 import { TrainScheduleBooking } from '../train-schedules/entities/train-schedule-booking.entity';
@@ -79,8 +80,10 @@ import {
  pickBulkWagonType,
 } from './wagon-type-resolver.util';
 import { deriveScheduleDirection } from './derive-schedule-direction.util';
+import { pickLowestFreeNumber, pickTrainNumberPool } from './train-number.util';
 import {
  deriveTrainCapacityFromLocomotive,
+  minLocomotiveLimits,
  wagonTypeDimensionsFromEntity,
 } from './train-capacity.util';
 import {
@@ -288,31 +291,42 @@ export class TrainSchedulingService {

  async createContainerTrainSchedule(dto: CreateContainerTrainScheduleDto) {
    const route = await this.getActiveRoute(dto.routeId);
-    const locomotive = await this.selectOrValidateLocomotive(dto.locomotiveId, 0, 0);
+
+    const locomotiveIds = [...new Set(dto.locomotiveIds)];
+    if (locomotiveIds.length < 2) {
+      throw new BadRequestException('A train must be pulled by at least two locomotives');
+    }

    const createdScheduleId = await this.dataSource.transaction(async (manager) => {
-      const lockedLocomotive = await manager.getRepository(Locomotive).findOne({
-        where: { id: locomotive.id },
-        lock: { mode: 'pessimistic_write' },
-      });
-      if (!lockedLocomotive) {
-        throw new NotFoundException(`Locomotive ${locomotive.id} not found`);
-      }
-      if (lockedLocomotive.status !== 'AVAILABLE') {
-        throw new ConflictException(`Locomotive ${lockedLocomotive.code} is not available`);
+      // Lock and validate every locomotive: all must be AVAILABLE and at the origin yard.
+      const lockedLocomotives: Locomotive[] = [];
+      for (const locomotiveId of locomotiveIds) {
+        const locked = await manager.getRepository(Locomotive).findOne({
+          where: { id: locomotiveId },
+          lock: { mode: 'pessimistic_write' },
+        });
+        if (!locked) {
+          throw new NotFoundException(`Locomotive ${locomotiveId} not found`);
+        }
+        if (locked.status !== 'AVAILABLE') {
+          throw new ConflictException(`Locomotive ${locked.code} is not available`);
+        }
+        if (locked.currentYardId !== route.originYardId) {
+          throw new ConflictException(
+            `Locomotive ${locked.code} is at yard ${locked.currentYardId} but schedule originates from ${route.originYardId}`,
+          );
+        }
+        lockedLocomotives.push(locked);
      }

      const direction = deriveScheduleDirection(
        route.originYard ?? { country: null },
        route.destinationYard ?? { country: null },
      );
-      if (lockedLocomotive.currentYardId !== route.originYardId) {
-        throw new ConflictException(
-          `Locomotive ${lockedLocomotive.code} is at yard ${lockedLocomotive.currentYardId} but schedule originates from ${route.originYardId}`,
-        );
-      }

-      const trainSet = await this.buildEmptyTrainSet(manager, lockedLocomotive);
+      const trainSet = await this.buildEmptyTrainSet(manager, lockedLocomotives);
+      // Effective capacity is capped by the weakest locomotive in the set.
+      const limitLoco = minLocomotiveLimits(lockedLocomotives) ?? undefined;
      const schedule = manager.getRepository(TrainSchedule).create({
        trainSetId: trainSet.id,
        routeId: route.id,
@@ -322,11 +336,14 @@ export class TrainSchedulingService {
        status: TrainScheduleStatusEnum.Draft,
        direction,
        maxWagons: (
-          await this.resolveTrainLimitConfig(dto, lockedLocomotive)
+          await this.resolveTrainLimitConfig(dto, limitLoco)
        ).maxWagonsPerTrain,
      });
      const saved = await manager.getRepository(TrainSchedule).save(schedule);
-      await manager.getRepository(Locomotive).update(lockedLocomotive.id, { status: 'ASSIGNED' });
+      await manager.getRepository(Locomotive).update(
+        { id: In(lockedLocomotives.map((l) => l.id)) },
+        { status: 'ASSIGNED' },
+      );
      return saved.id;
    });

@@ -375,8 +392,9 @@ export class TrainSchedulingService {
      maxWagonsPerTrain: dto.maxWagonsPerTrain,
    };

-    const locomotive = schedule.trainSet.locomotive;
-    const limits = await this.resolveTrainLimitConfig(previewDto, locomotive ?? undefined);
+    const setLocomotives = this.locomotivesOfTrainSet(schedule.trainSet);
+    const limitLoco = minLocomotiveLimits(setLocomotives) ?? undefined;
+    const limits = await this.resolveTrainLimitConfig(previewDto, limitLoco);
    const validation = await this.validateBookingsForScheduling(
      previewDto,
      freightType ?? null,
@@ -408,17 +426,17 @@ export class TrainSchedulingService {
    const totalWeightTons = validation.summary.totalWeightTons;
    const totalLengthMeters = validation.summary.totalLengthMeters;

-    if (!locomotive) {
-      throw new BadRequestException('Schedule train set has no locomotive');
+    if (!limitLoco) {
+      throw new BadRequestException('Schedule train set has no locomotives');
    }
-    if (Number(locomotive.maxPullWeightTons) < totalWeightTons) {
+    if (limitLoco.maxPullWeightTons < totalWeightTons) {
      throw new BadRequestException(
-        `Locomotive ${locomotive.code} cannot pull ${totalWeightTons}T`,
+        `Train set locomotives cannot pull ${totalWeightTons}T`,
      );
    }
-    if (Number(locomotive.maxTrainLengthMeters) < totalLengthMeters) {
+    if (limitLoco.maxTrainLengthMeters < totalLengthMeters) {
      throw new BadRequestException(
-        `Locomotive ${locomotive.code} cannot support ${totalLengthMeters}m`,
+        `Train set locomotives cannot support ${totalLengthMeters}m`,
      );
    }

@@ -681,10 +699,12 @@ export class TrainSchedulingService {

    const now = new Date();
    await this.dataSource.transaction(async (manager) => {
+      const trainNumber = await this.assignTrainNumber(manager, schedule);
+
      await this.trainSchedulesRepository.updateStatus(
        scheduleId,
        TrainScheduleStatusEnum.Dispatched,
-        { actualDepartureAt: now },
+        { actualDepartureAt: now, trainNumber },
        manager,
      );
      if (schedule.trainSetId) {
@@ -718,6 +738,60 @@ export class TrainSchedulingService {
    return this.getTrainScheduleById(scheduleId);
  }

+  /**
+   * Assign a fixed train number on dispatch. The number is drawn from the pool
+   * for the train's dominant cargo type (container vs bulk) and trade direction
+   * (export = odd, import = even). Numbers recycle once a train ARRIVES, so the
+   * "used" set is every still-DISPATCHED schedule's number. Locked FOR UPDATE so
+   * concurrent dispatches can't grab the same number. Throws when the pool is
+   * exhausted. Idempotent: returns the existing number if already assigned.
+   */
+  private async assignTrainNumber(
+    manager: EntityManager,
+    schedule: TrainSchedule,
+  ): Promise<string> {
+    if (schedule.trainNumber) return schedule.trainNumber;
+
+    // Count container vs bulk wagons from the planned allocations.
+    let containerWagons = 0;
+    let bulkWagons = 0;
+    for (const wagon of schedule.trainSet?.wagons ?? []) {
+      const isBulk = (wagon.allocations ?? []).some((a) => a.loadType === 'BULK');
+      if (isBulk) bulkWagons += 1;
+      else containerWagons += 1;
+    }
+
+    const direction =
+      (schedule.direction as 'IMPORT' | 'EXPORT' | 'DOMESTIC' | null) ??
+      (schedule.originStation && schedule.destinationStation
+        ? deriveScheduleDirection(schedule.originStation, schedule.destinationStation)
+        : null);
+
+    const pool = pickTrainNumberPool(containerWagons, bulkWagons, direction);
+
+    // Lock the set of currently-active numbered schedules so two concurrent
+    // dispatches serialize and can't both claim the same lowest-free number.
+    const activeNumbered = await manager
+      .getRepository(TrainSchedule)
+      .createQueryBuilder('schedule')
+      .setLock('pessimistic_write')
+      .where('schedule.status = :status', { status: TrainScheduleStatusEnum.Dispatched })
+      .andWhere('schedule.train_number IS NOT NULL')
+      .getMany();
+
+    const usedNumbers = activeNumbered
+      .map((s) => s.trainNumber)
+      .filter((n): n is string => Boolean(n));
+
+    const number = pickLowestFreeNumber(pool.numbers, usedNumbers);
+    if (!number) {
+      throw new ConflictException(
+        `No free ${pool.cargo.toLowerCase()} ${pool.direction.toLowerCase()} train number available; a train must arrive to free one`,
+      );
+    }
+    return number;
+  }
+
  /** Open or close a schedule's booking window (staff override). */
  async setBookingWindow(scheduleId: string, status: 'OPEN' | 'CLOSED'): Promise<void> {
    await this.dataSource
@@ -931,16 +1005,12 @@ export class TrainSchedulingService {
        });
      }

-      if (schedule.trainSet?.locomotiveId) {
-        const loco = await manager
-          .getRepository(Locomotive)
-          .findOne({ where: { id: schedule.trainSet.locomotiveId } });
-        if (loco) {
-          await manager.getRepository(Locomotive).update(loco.id, {
-            status: 'AVAILABLE',
-            currentYardId: schedule.destinationStationId,
-          });
-        }
+      const arrivingLocoIds = this.locomotivesOfTrainSet(schedule.trainSet).map((l) => l.id);
+      if (arrivingLocoIds.length) {
+        await manager.getRepository(Locomotive).update(
+          { id: In(arrivingLocoIds) },
+          { status: 'AVAILABLE', currentYardId: schedule.destinationStationId },
+        );
      }

      for (const slot of schedule.trainSet?.wagons ?? []) {
@@ -982,7 +1052,7 @@ export class TrainSchedulingService {
  async getContainerTrainSchedules() {
    const schedules = await this.trainSchedulesRepository.findAll({
      relations: {
-        trainSet: { locomotive: true },
+        trainSet: { locomotive: true, locomotives: { locomotive: true } },
        route: true,
        originStation: true,
        destinationStation: true,
@@ -1013,10 +1083,11 @@ export class TrainSchedulingService {
      if (schedule.trainSetId) {
        await manager.getRepository(TrainSet).update(schedule.trainSetId, { status: 'CANCELLED' });
      }
-      if (schedule.trainSet?.locomotiveId) {
-        await manager.getRepository(Locomotive).update(schedule.trainSet.locomotiveId, {
-          status: 'AVAILABLE',
-        });
+      const cancelledLocoIds = this.locomotivesOfTrainSet(schedule.trainSet).map((l) => l.id);
+      if (cancelledLocoIds.length) {
+        await manager
+          .getRepository(Locomotive)
+          .update({ id: In(cancelledLocoIds) }, { status: 'AVAILABLE' });
      }
      for (const wagon of schedule.trainSet?.wagons ?? []) {
        if (wagon.physicalWagonId) {
@@ -1251,24 +1322,29 @@ export class TrainSchedulingService {
      }
    }

-    let assignedLocomotive: Locomotive | null = null;
+    let assignedLocomotives: Locomotive[] = [];
    if (targetScheduleId) {
      const targetSchedule =
        await this.trainSchedulesRepository.findByIdWithFullGraph(targetScheduleId);
-      assignedLocomotive = targetSchedule?.trainSet?.locomotive ?? null;
+      assignedLocomotives = this.locomotivesOfTrainSet(targetSchedule?.trainSet);
    }

-    if (assignedLocomotive) {
-      if (assignedLocomotive.currentYardId !== originYardId) {
+    if (assignedLocomotives.length) {
+      // Every locomotive of the set must sit at the origin yard, and the weakest
+      // one must still be able to pull the train (min limits across the set).
+      const offYard = assignedLocomotives.find((l) => l.currentYardId !== originYardId);
+      const setLimits = minLocomotiveLimits(assignedLocomotives);
+      if (offYard) {
        violations.push(
-          `Locomotive ${assignedLocomotive.code} is not at the schedule origin yard`,
+          `Locomotive ${offYard.code} is not at the schedule origin yard`,
        );
      } else if (
-        Number(assignedLocomotive.maxPullWeightTons) < totalWeightTons ||
-        Number(assignedLocomotive.maxTrainLengthMeters) < totalLengthMeters
+        setLimits &&
+        (setLimits.maxPullWeightTons < totalWeightTons ||
+          setLimits.maxTrainLengthMeters < totalLengthMeters)
      ) {
        violations.push(
-          'Assigned locomotive cannot support the total train weight and length',
+          'Assigned locomotives cannot support the total train weight and length',
        );
      }
    } else {
@@ -1818,6 +1894,22 @@ export class TrainSchedulingService {
    }
  }

+  /**
+   * All locomotives attached to a loaded train set. Prefers the `locomotives`
+   * link rows; falls back to the legacy single `locomotive` for train sets
+   * created before multi-loco support.
+   */
+  private locomotivesOfTrainSet(
+    trainSet: TrainSet | null | undefined,
+  ): Locomotive[] {
+    if (!trainSet) return [];
+    const linked = (trainSet.locomotives ?? [])
+      .map((link) => link.locomotive)
+      .filter((loco): loco is Locomotive => Boolean(loco));
+    if (linked.length) return linked;
+    return trainSet.locomotive ? [trainSet.locomotive] : [];
+  }
+
  async selectOrValidateLocomotive(
    locomotiveId: string,
    totalWeightTons: number,
@@ -1841,15 +1933,28 @@ export class TrainSchedulingService {
    return locomotive;
  }

-  private async buildEmptyTrainSet(manager: EntityManager, locomotive: Locomotive) {
+  private async buildEmptyTrainSet(manager: EntityManager, locomotives: Locomotive[]) {
+    const [primary] = locomotives;
    const trainSet = manager.getRepository(TrainSet).create({
-      locomotiveId: locomotive.id,
+      // `locomotiveId` retained as the primary locomotive for single-loco read paths.
+      locomotiveId: primary.id,
      totalWeightTons: 0,
      totalLengthMeters: 0,
      wagonCount: 0,
      status: 'DRAFT',
    });
-    return manager.getRepository(TrainSet).save(trainSet);
+    const saved = await manager.getRepository(TrainSet).save(trainSet);
+
+    const links = locomotives.map((loco, index) =>
+      manager.getRepository(TrainSetLocomotive).create({
+        trainSetId: saved.id,
+        locomotiveId: loco.id,
+        sequenceNo: index,
+      }),
+    );
+    await manager.getRepository(TrainSetLocomotive).save(links);
+
+    return saved;
  }

  private async getActiveRoute(routeId: string) {
@@ -1915,6 +2020,12 @@ export class TrainSchedulingService {
            currentYardId: schedule.trainSet.locomotive.currentYardId ?? null,
          }
        : null,
+      locomotives: this.locomotivesOfTrainSet(schedule.trainSet).map((loco) => ({
+        id: loco.id,
+        code: loco.code,
+        name: loco.name ?? null,
+        currentYardId: loco.currentYardId ?? null,
+      })),
      wagonCount: schedule.trainSet?.wagonCount ?? 0,
      totalWeightTons: roundTons(Number(schedule.trainSet?.totalWeightTons ?? 0)),
      totalLengthMeters: roundTons(Number(schedule.trainSet?.totalLengthMeters ?? 0)),
@@ -1952,7 +2063,7 @@ export class TrainSchedulingService {
        bookingWindowStatus: 'OPEN',
      },
      relations: {
-        trainSet: { locomotive: true },
+        trainSet: { locomotive: true, locomotives: { locomotive: true } },
        route: { milestones: true },
        originStation: true,
        destinationStation: true,
@@ -2099,6 +2210,15 @@ export class TrainSchedulingService {
                  ),
                }
              : null,
+            locomotives: this.locomotivesOfTrainSet(schedule.trainSet).map((loco) => ({
+              id: loco.id,
+              code: loco.code,
+              name: loco.name ?? null,
+              status: loco.status,
+              currentYardId: loco.currentYardId ?? null,
+              maxPullWeightTons: roundTons(Number(loco.maxPullWeightTons)),
+              maxTrainLengthMeters: roundTons(Number(loco.maxTrainLengthMeters)),
+            })),
            wagons: [...(schedule.trainSet.wagons ?? [])]
              .sort((a, b) => a.sequenceNo - b.sequenceNo)
              .map((wagon) => ({