refactor: share tbank-sync claim and generation lock (WealthLedger/backend#166) (#188)

src/modules/investments/instruments/infrastructure/candle.repository.ts

@@ -2,8 +2,15 @@ import { Injectable } from '@nestjs/common';
 import { createHash, randomUUID } from 'node:crypto';
 
 import type { Prisma } from '../../../../../generated/prisma/client';
-import type { TbankSyncAttempt } from '../../../../integrations/tbank/tbank.api';
 
+import {
+  claimSyncWork,
+  lockConnectionGeneration,
+} from '../../../../integrations/tbank/tbank-sync-claim';
+import {
+  isJoinableRunningAttempt,
+  type TbankSyncAttempt,
+} from '../../../../integrations/tbank/tbank.api';
 import { PrismaService } from '../../../../prisma/prisma.service';
 import {
   type CandleRange,
@@ -62,34 +69,30 @@ export class CandleRepository {
     priority = 1,
   ) {
     return this.prisma.$transaction(async (tx) => {
-      const current = await tx.tbankConnection.updateMany({
-        where: { ...owner, invalidToken: false },
-        data: { membershipRevision: { increment: 0 } },
-      });
-      if (!current.count) return null;
+      const current = await lockConnectionGeneration(tx, owner);
+      if (!current) return null;
       const scope = { ...owner, scopeKey };
       const previous = await tx.tbankSyncWork.findUnique({
         where: { userId_revision_scopeKey: scope },
       });
-      if (
-        previous?.state === 'running' &&
-        previous.processRunId === runId &&
-        active.has(previous.attemptId)
-      ) {
-        const previousIntent = previous.intent as Prisma.JsonObject;
+      const joinable = isJoinableRunningAttempt(previous, runId, active)
+        ? previous
+        : null;
+      if (joinable) {
+        const previousIntent = joinable.intent as Prisma.JsonObject;
         const previousNow =
           typeof previousIntent.now === 'string' ? previousIntent.now : '';
-        const nextPriority = Math.min(previous.priority, priority);
+        const nextPriority = Math.min(joinable.priority, priority);
         if (
-          nextPriority === previous.priority &&
+          nextPriority === joinable.priority &&
           (!intent.force ||
             (previousIntent.force === true &&
               (!intent.now || intent.now <= previousNow)))
         ) {
-          return { work: previous, launch: false };
+          return { work: joinable, launch: false };
         }
         const work = await tx.tbankSyncWork.update({
-          where: { id: previous.id },
+          where: { id: joinable.id },
           data: {
             priority: nextPriority,
             intent: {
@@ -163,15 +166,10 @@ export class CandleRepository {
       ).values(),
     ];
     return this.prisma.$transaction(async (tx) => {
-      const valid = await tx.tbankSyncWork.updateMany({
-        where: {
-          ...attempt,
-          state: 'running',
-          connection: { invalidToken: false },
-        },
-        data: { progress: { lastWindow: key } },
+      const valid = await claimSyncWork(tx, attempt, {
+        progress: { lastWindow: key },
       });
-      if (!valid.count) return false;
+      if (!valid) return false;
       const identity = {
         workId: attempt.id,
         attemptId: attempt.attemptId,
@@ -196,18 +194,11 @@ export class CandleRepository {
     completeWork = true,
   ) {
     return this.prisma.$transaction(async (tx) => {
-      const claimed = await tx.tbankSyncWork.updateMany({
-        where: {
-          ...attempt,
-          state: 'running',
-          connection: { invalidToken: false },
-        },
-        data: {
-          ...(completeWork ? { state: 'complete' } : {}),
-          errorCode: null,
-        },
+      const claimed = await claimSyncWork(tx, attempt, {
+        ...(completeWork ? { state: 'complete' } : {}),
+        errorCode: null,
       });
-      if (!claimed.count) return false;
+      if (!claimed) return false;
       const staged = await tx.tbankSyncStage.findMany({
         where: {
           workId: attempt.id,
@@ -311,28 +302,17 @@ export class CandleRepository {
         },
       });
       if (!origin) return false;
-      const result = await tx.tbankSyncWork.updateMany({
-        where: {
-          ...attempt,
-          state: 'running',
-          connection: { invalidToken: false },
-        },
-        data: { state: 'complete', errorCode: null },
+      const result = await claimSyncWork(tx, attempt, {
+        state: 'complete',
+        errorCode: null,
       });
-      return result.count > NO_UPDATED_ROWS;
+      return result > NO_UPDATED_ROWS;
     });
   }
 
   fail(attempt: TbankSyncAttempt, errorCode: string) {
     return this.prisma.$transaction(async (tx) => {
-      await tx.tbankSyncWork.updateMany({
-        where: {
-          ...attempt,
-          state: 'running',
-          connection: { invalidToken: false },
-        },
-        data: { state: 'failed', errorCode },
-      });
+      await claimSyncWork(tx, attempt, { state: 'failed', errorCode });
       await tx.tbankSyncStage.deleteMany({
         where: { workId: attempt.id, attemptId: attempt.attemptId },
       });

src/modules/investments/instruments/infrastructure/catalog-discovery.repository.ts

@@ -5,6 +5,7 @@ import { isDeepStrictEqual } from 'node:util';
 import type { Prisma } from '../../../../../generated/prisma/client';
 
 import { TbankErrorCode } from '../../../../integrations/tbank/tbank-error-code';
+import { claimSyncWork, lockConnectionGeneration } from '../../../../integrations/tbank/tbank-sync-claim';
 import { isJoinableRunningAttempt, type TbankBond, type TbankShare, type TbankSyncAttempt } from '../../../../integrations/tbank/tbank.api';
 import { PrismaService } from '../../../../prisma/prisma.service';
 
@@ -159,24 +160,18 @@ export class CatalogDiscoveryRepository {
     initiator: TbankSyncAttempt = attempt,
   ) {
     return this.prisma.$transaction(async (tx) => {
-      const origin = await tx.tbankSyncWork.updateMany({
-        where: {
-          ...initiator,
-          state: { in: ['running', 'complete'] },
-          connection: { invalidToken: false },
-        },
-        data: { progress: {} },
+      const origin = await claimSyncWork(
+        tx,
+        initiator,
+        { progress: {} },
+        { state: { in: ['running', 'complete'] } },
+      );
+      if (!origin) return false;
+      const claimed = await claimSyncWork(tx, attempt, {
+        state: 'complete',
+        progress: {},
       });
-      if (!origin.count) return false;
-      const claimed = await tx.tbankSyncWork.updateMany({
-        where: {
-          ...attempt,
-          state: 'running',
-          connection: { invalidToken: false },
-        },
-        data: { state: 'complete', progress: {} },
-      });
-      if (!claimed.count) return false;
+      if (!claimed) return false;
       const payload = JSON.parse(
         JSON.stringify(value),
       ) as Prisma.InputJsonObject;
@@ -284,11 +279,8 @@ export class CatalogDiscoveryRepository {
     active: ReadonlySet<string>,
   ) {
     return this.prisma.$transaction(async (tx) => {
-      const current = await tx.tbankConnection.updateMany({
-        where: { ...owner, invalidToken: false },
-        data: { membershipRevision: { increment: 0 } },
-      });
-      if (!current.count) return null;
+      const current = await lockConnectionGeneration(tx, owner);
+      if (!current) return null;
       const scope = { ...owner, scopeKey: catalogScope(input) };
       const previous = await tx.tbankSyncWork.findUnique({
         where: { userId_revision_scopeKey: scope },
@@ -341,15 +333,11 @@ export class CatalogDiscoveryRepository {
     now: Date,
   ) {
     return this.prisma.$transaction(async (tx) => {
-      const claimed = await tx.tbankSyncWork.updateMany({
-        where: {
-          ...attempt,
-          state: 'running',
-          connection: { invalidToken: false },
-        },
-        data: { state: 'complete', progress: {} },
+      const claimed = await claimSyncWork(tx, attempt, {
+        state: 'complete',
+        progress: {},
       });
-      if (!claimed.count) return false;
+      if (!claimed) return false;
       for (const instrument of instruments) {
         // Discovery can create a minimal identity, but cannot overwrite verified metadata.
         const identity = { provider, uid: instrument.uid };
@@ -461,6 +449,7 @@ export class CatalogDiscoveryRepository {
 
   finish(attempt: TbankSyncAttempt, success: boolean) {
     return this.prisma.$transaction(async (tx) => {
+      // Fails sibling attempts whose connection was invalidated mid-flight.
       await tx.tbankSyncWork.updateMany({
         where: {
           ...attempt,
@@ -469,16 +458,9 @@ export class CatalogDiscoveryRepository {
         },
         data: { state: 'failed', errorCode: TbankErrorCode.InvalidToken },
       });
-      await tx.tbankSyncWork.updateMany({
-        where: {
-          ...attempt,
-          state: 'running',
-          connection: { invalidToken: false },
-        },
-        data: {
-          state: success ? 'complete' : 'failed',
-          errorCode: success ? null : TbankErrorCode.Unavailable,
-        },
+      await claimSyncWork(tx, attempt, {
+        state: success ? 'complete' : 'failed',
+        errorCode: success ? null : TbankErrorCode.Unavailable,
       });
     });
   }