Allow for deleting connection in an unexpected state

airbyte-workers/src/main/java/io/airbyte/workers/temporal/TemporalClient.java

@@ -28,6 +28,7 @@
 import io.airbyte.workers.temporal.scheduling.state.WorkflowState;
 import io.airbyte.workers.temporal.spec.SpecWorkflow;
 import io.airbyte.workers.temporal.sync.SyncWorkflow;
+import io.temporal.api.common.v1.WorkflowExecution;
 import io.temporal.api.workflowservice.v1.ListOpenWorkflowExecutionsRequest;
 import io.temporal.api.workflowservice.v1.ListOpenWorkflowExecutionsResponse;
 import io.temporal.client.BatchRequest;
@@ -252,9 +253,32 @@ public void submitConnectionUpdaterAsync(final UUID connectionId) {
   }
 
   public void deleteConnection(final UUID connectionId) {
-    final ConnectionManagerWorkflow connectionManagerWorkflow = getConnectionUpdateWorkflow(connectionId);
-
-    connectionManagerWorkflow.deleteConnection();
+    try {
+      final ConnectionManagerWorkflow connectionManagerWorkflow = getConnectionUpdateWorkflow(connectionId);
+      connectionManagerWorkflow.deleteConnection();
+    } catch (final IllegalStateException e) {
+      log.info("Connection in an illegal state; Creating new workflow and sending delete signal");
+
+      final ConnectionManagerWorkflow connectionManagerWorkflow = getWorkflowOptionsWithWorkflowId(ConnectionManagerWorkflow.class,
+          TemporalJobType.CONNECTION_UPDATER, getConnectionManagerName(connectionId));
+
+      final ConnectionUpdaterInput input = ConnectionUpdaterInput.builder()
+          .connectionId(connectionId)
+          .jobId(null)
+          .attemptId(null)
+          .fromFailure(false)
+          .attemptNumber(1)
+          .workflowState(null)
+          .resetConnection(false)
+          .fromJobResetFailure(false)
+          .build();
+
+      final BatchRequest signalRequest = client.newSignalWithStartRequest();
+      signalRequest.add(connectionManagerWorkflow::run, input);
+      signalRequest.add(connectionManagerWorkflow::deleteConnection);
+      client.signalWithStart(signalRequest);
+      log.info("New start request and delete signal submitted");
+    }
   }
 
   public void update(final UUID connectionId) {
@@ -440,11 +464,13 @@ private <T> T getWorkflowOptionsWithWorkflowId(final Class<T> workflowClass, fin
     return client.newWorkflowStub(workflowClass, TemporalUtils.getWorkflowOptionsWithWorkflowId(jobType, name));
   }
 
-  private <T> T getExistingWorkflow(final Class<T> workflowClass, final String name) {
+  @VisibleForTesting
+  public <T> T getExistingWorkflow(final Class<T> workflowClass, final String name) {
     return client.newWorkflowStub(workflowClass, name);
   }
 
-  private ConnectionManagerWorkflow getConnectionUpdateWorkflow(final UUID connectionId) {
+  @VisibleForTesting
+  ConnectionManagerWorkflow getConnectionUpdateWorkflow(final UUID connectionId) {
     final boolean workflowReachable = isWorkflowReachable(getConnectionManagerName(connectionId));
 
     if (!workflowReachable) {
@@ -506,7 +532,7 @@ boolean isWorkflowStateRunning(final String workflowName) {
   }
 
   @VisibleForTesting
-  static String getConnectionManagerName(final UUID connectionId) {
+  public static String getConnectionManagerName(final UUID connectionId) {
     return "connection_manager_" + connectionId;
   }
 

airbyte-workers/src/test/java/io/airbyte/workers/temporal/TemporalClientTest.java

@@ -10,8 +10,10 @@
 import static org.junit.jupiter.api.Assertions.assertTrue;
 import static org.mockito.ArgumentMatchers.any;
 import static org.mockito.ArgumentMatchers.anyString;
+import static org.mockito.ArgumentMatchers.eq;
 import static org.mockito.Mockito.doNothing;
 import static org.mockito.Mockito.doReturn;
+import static org.mockito.Mockito.doThrow;
 import static org.mockito.Mockito.mock;
 import static org.mockito.Mockito.spy;
 import static org.mockito.Mockito.times;
@@ -39,8 +41,11 @@
 import io.airbyte.workers.temporal.scheduling.ConnectionManagerWorkflow.JobInformation;
 import io.airbyte.workers.temporal.spec.SpecWorkflow;
 import io.airbyte.workers.temporal.sync.SyncWorkflow;
+import io.temporal.client.BatchRequest;
 import io.temporal.client.WorkflowClient;
 import io.temporal.serviceclient.WorkflowServiceStubs;
+import io.temporal.workflow.Functions.Func;
+import io.temporal.workflow.Functions.Proc1;
 import java.io.IOException;
 import java.nio.file.Files;
 import java.nio.file.Path;
@@ -51,6 +56,7 @@
 import org.junit.jupiter.api.DisplayName;
 import org.junit.jupiter.api.Nested;
 import org.junit.jupiter.api.Test;
+import org.mockito.Mockito;
 
 class TemporalClientTest {
 
@@ -257,4 +263,28 @@ public void migrateCalled() {
 
   }
 
+  @SuppressWarnings("unchecked")
+  @Test
+  @DisplayName("Test delete connection method.")
+  void testDeleteConnection() {
+    final ConnectionManagerWorkflow mConnectionManagerWorkflow = mock(ConnectionManagerWorkflow.class);
+
+    doReturn(true).when(temporalClient).isWorkflowReachable(anyString());
+    when(workflowClient.newWorkflowStub(any(Class.class), anyString())).thenReturn(mConnectionManagerWorkflow);
+
+    final JobSyncConfig syncConfig = new JobSyncConfig()
+        .withSourceDockerImage(IMAGE_NAME1)
+        .withSourceDockerImage(IMAGE_NAME2)
+        .withSourceConfiguration(Jsons.emptyObject())
+        .withDestinationConfiguration(Jsons.emptyObject())
+        .withOperationSequence(List.of())
+        .withConfiguredAirbyteCatalog(new ConfiguredAirbyteCatalog());
+
+    temporalClient.submitSync(JOB_ID, ATTEMPT_ID, syncConfig, CONNECTION_ID);
+    temporalClient.deleteConnection(CONNECTION_ID);
+
+    verify(workflowClient, Mockito.never()).newSignalWithStartRequest();
+    verify(mConnectionManagerWorkflow).deleteConnection();
+  }
+
 }

airbyte-workers/src/test/java/io/airbyte/workers/temporal/scheduling/ConnectionManagerWorkflowTest.java

@@ -4,13 +4,18 @@
 
 package io.airbyte.workers.temporal.scheduling;
 
+import static io.airbyte.workers.temporal.TemporalClient.getConnectionManagerName;
+import static org.junit.jupiter.api.Assertions.assertFalse;
+import static org.junit.jupiter.api.Assertions.assertThrows;
+import static org.junit.jupiter.api.Assertions.assertTrue;
 import static org.mockito.Mockito.atLeastOnce;
 
 import io.airbyte.config.FailureReason.FailureOrigin;
 import io.airbyte.config.FailureReason.FailureType;
 import io.airbyte.config.StandardSyncInput;
 import io.airbyte.scheduler.models.IntegrationLauncherConfig;
 import io.airbyte.scheduler.models.JobRunConfig;
+import io.airbyte.workers.temporal.TemporalClient;
 import io.airbyte.workers.temporal.TemporalJobType;
 import io.airbyte.workers.temporal.scheduling.activities.AutoDisableConnectionActivity;
 import io.airbyte.workers.temporal.scheduling.activities.AutoDisableConnectionActivity.AutoDisableConnectionActivityInput;
@@ -42,10 +47,14 @@
 import io.airbyte.workers.temporal.scheduling.testsyncworkflow.SyncWorkflowWithActivityFailureException;
 import io.airbyte.workers.temporal.sync.SyncWorkflow;
 import io.temporal.client.WorkflowClient;
+import io.temporal.client.WorkflowNotFoundException;
 import io.temporal.client.WorkflowOptions;
 import io.temporal.failure.ApplicationFailure;
 import io.temporal.testing.TestWorkflowEnvironment;
 import io.temporal.worker.Worker;
+import java.io.IOException;
+import java.nio.file.Files;
+import java.nio.file.Path;
 import java.time.Duration;
 import java.util.ArrayList;
 import java.util.List;
@@ -409,6 +418,55 @@ public void deleteSync() throws InterruptedException {
       Mockito.verify(mConnectionDeletionActivity, Mockito.times(1)).deleteConnection(Mockito.any());
     }
 
+    @Test
+    @Timeout(value = 2,
+             unit = TimeUnit.SECONDS)
+    @DisplayName("Test that the connection is properly deleted even if the temporal workflow is in a bad state")
+    public void deleteConnectionInUnexpectedState() throws IOException {
+
+      final UUID connectionId = UUID.randomUUID();
+      final Path workspaceRoot = Files.createTempDirectory(Path.of("/tmp"), "temporal_client_test");
+      final TemporalClient temporalClient = new TemporalClient(client, workspaceRoot, testEnv.getWorkflowService(), null);
+
+      // assert that we cannot reach the existing workflow
+      ConnectionManagerWorkflow connectionManagerWorkflow =
+          temporalClient.getExistingWorkflow(ConnectionManagerWorkflow.class, getConnectionManagerName(connectionId));
+      assertThrows(WorkflowNotFoundException.class, connectionManagerWorkflow::getState);
+
+      temporalClient.deleteConnection(connectionId);
+      testEnv.sleep(Duration.ofMinutes(1L));
+
+      // assert that we spin up a new workflow and delete the connection
+      connectionManagerWorkflow = temporalClient.getExistingWorkflow(ConnectionManagerWorkflow.class, getConnectionManagerName(connectionId));
+      assertTrue(connectionManagerWorkflow.getState().isDeleted());
+    }
+
+    @Test
+    @Timeout(value = 2,
+             unit = TimeUnit.SECONDS)
+    @DisplayName("Test that the connection is properly deleted")
+    public void deleteConnection() throws IOException {
+
+      final UUID connectionId = UUID.randomUUID();
+      final Path workspaceRoot = Files.createTempDirectory(Path.of("/tmp"), "temporal_client_test");
+      final TemporalClient temporalClient = new TemporalClient(client, workspaceRoot, testEnv.getWorkflowService(), null);
+
+      temporalClient.submitConnectionUpdaterAsync(connectionId);
+      testEnv.sleep(Duration.ofMinutes(1L));
+
+      // assert connection is not deleted
+      ConnectionManagerWorkflow connectionManagerWorkflow =
+          temporalClient.getExistingWorkflow(ConnectionManagerWorkflow.class, getConnectionManagerName(connectionId));
+      assertFalse(connectionManagerWorkflow.getState().isDeleted());
+
+      temporalClient.deleteConnection(connectionId);
+      testEnv.sleep(Duration.ofMinutes(1L));
+
+      // assert that it is deleted
+      connectionManagerWorkflow = temporalClient.getExistingWorkflow(ConnectionManagerWorkflow.class, getConnectionManagerName(connectionId));
+      assertTrue(connectionManagerWorkflow.getState().isDeleted());
+    }
+
   }
 
   @Nested