close ssh in case of exception during check in Postgres connector

airbyte-integrations/bases/base-java/src/main/java/io/airbyte/integrations/base/IntegrationRunner.java

@@ -135,8 +135,7 @@ public void runInternal(final ITransaction transaction, final IntegrationConfig
         validateConfig(integration.spec().getConnectionSpecification(), config, "READ");
         final ConfiguredAirbyteCatalog catalog = parseConfig(parsed.getCatalogPath(), ConfiguredAirbyteCatalog.class);
         final Optional<JsonNode> stateOptional = parsed.getStatePath().map(IntegrationRunner::parseConfig);
-        final AutoCloseableIterator<AirbyteMessage> messageIterator = source.read(config, catalog, stateOptional.orElse(null));
-        try (messageIterator) {
+        try (final AutoCloseableIterator<AirbyteMessage> messageIterator = source.read(config, catalog, stateOptional.orElse(null))) {
           AirbyteSentry.executeWithTracing("ReadSource", () -> messageIterator.forEachRemaining(outputRecordCollector::accept));
         }
       }
@@ -145,8 +144,9 @@ public void runInternal(final ITransaction transaction, final IntegrationConfig
         final JsonNode config = parseConfig(parsed.getConfigPath());
         validateConfig(integration.spec().getConnectionSpecification(), config, "WRITE");
         final ConfiguredAirbyteCatalog catalog = parseConfig(parsed.getCatalogPath(), ConfiguredAirbyteCatalog.class);
-        final AirbyteMessageConsumer consumer = destination.getConsumer(config, catalog, outputRecordCollector);
-        AirbyteSentry.executeWithTracing("WriteDestination", () -> consumeWriteStream(consumer));
+        try (final AirbyteMessageConsumer consumer = destination.getConsumer(config, catalog, outputRecordCollector)) {
+          AirbyteSentry.executeWithTracing("WriteDestination", () -> consumeWriteStream(consumer));
+        }
       }
       default -> throw new IllegalStateException("Unexpected value: " + parsed.getCommand());
     }
@@ -158,8 +158,7 @@ public void runInternal(final ITransaction transaction, final IntegrationConfig
   static void consumeWriteStream(final AirbyteMessageConsumer consumer) throws Exception {
     // use a Scanner that only processes new line characters to strictly abide with the
     // https://jsonlines.org/ standard
-    final Scanner input = new Scanner(System.in).useDelimiter("[\r\n]+");
-    try (consumer) {
+    try(final Scanner input = new Scanner(System.in).useDelimiter("[\r\n]+")) {
       consumer.start();
       while (input.hasNext()) {
         final String inputString = input.next();

airbyte-integrations/bases/base-java/src/main/java/io/airbyte/integrations/base/ssh/SshWrappedSource.java

@@ -14,9 +14,11 @@
 import io.airbyte.protocol.models.ConfiguredAirbyteCatalog;
 import io.airbyte.protocol.models.ConnectorSpecification;
 import java.util.List;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
 
 public class SshWrappedSource implements Source {
-
+  private static final Logger LOGGER = LoggerFactory.getLogger(SshWrappedSource.class);
   private final Source delegate;
   private final List<String> hostKey;
   private final List<String> portKey;
@@ -46,7 +48,15 @@ public AirbyteCatalog discover(final JsonNode config) throws Exception {
   public AutoCloseableIterator<AirbyteMessage> read(final JsonNode config, final ConfiguredAirbyteCatalog catalog, final JsonNode state)
       throws Exception {
     final SshTunnel tunnel = SshTunnel.getInstance(config, hostKey, portKey);
-    return AutoCloseableIterators.appendOnClose(delegate.read(tunnel.getConfigInTunnel(), catalog, state), tunnel::close);
+    final AutoCloseableIterator<AirbyteMessage> delegateRead;
+    try {
+      delegateRead = delegate.read(tunnel.getConfigInTunnel(), catalog, state);
+    } catch (final Exception e) {
+      LOGGER.error("Exception occurred while getting the delegate read iterator, closing SSH tunnel", e);
+      tunnel.close();
+      throw e;
+    }
+    return AutoCloseableIterators.appendOnClose(delegateRead, tunnel::close);
   }
 
 }

airbyte-integrations/bases/base-java/src/test/java/io/airbyte/integrations/base/IntegrationRunnerTest.java

@@ -241,14 +241,13 @@ void testDestinationConsumerLifecycleSuccess() throws Exception {
         + Jsons.serialize(message2) + "\n"
         + Jsons.serialize(stateMessage)).getBytes()));
 
-    final AirbyteMessageConsumer airbyteMessageConsumerMock = mock(AirbyteMessageConsumer.class);
-    IntegrationRunner.consumeWriteStream(airbyteMessageConsumerMock);
-
-    final InOrder inOrder = inOrder(airbyteMessageConsumerMock);
-    inOrder.verify(airbyteMessageConsumerMock).accept(message1);
-    inOrder.verify(airbyteMessageConsumerMock).accept(message2);
-    inOrder.verify(airbyteMessageConsumerMock).accept(stateMessage);
-    inOrder.verify(airbyteMessageConsumerMock).close();
+    try (final AirbyteMessageConsumer airbyteMessageConsumerMock = mock(AirbyteMessageConsumer.class)) {
+      IntegrationRunner.consumeWriteStream(airbyteMessageConsumerMock);
+      final InOrder inOrder = inOrder(airbyteMessageConsumerMock);
+      inOrder.verify(airbyteMessageConsumerMock).accept(message1);
+      inOrder.verify(airbyteMessageConsumerMock).accept(message2);
+      inOrder.verify(airbyteMessageConsumerMock).accept(stateMessage);
+    }
   }
 
   @Test
@@ -267,15 +266,13 @@ void testDestinationConsumerLifecycleFailure() throws Exception {
             .withEmittedAt(EMITTED_AT));
     System.setIn(new ByteArrayInputStream((Jsons.serialize(message1) + "\n" + Jsons.serialize(message2)).getBytes()));
 
-    final AirbyteMessageConsumer airbyteMessageConsumerMock = mock(AirbyteMessageConsumer.class);
-    doThrow(new IOException("error")).when(airbyteMessageConsumerMock).accept(message1);
-
-    assertThrows(IOException.class, () -> IntegrationRunner.consumeWriteStream(airbyteMessageConsumerMock));
-
-    final InOrder inOrder = inOrder(airbyteMessageConsumerMock);
-    inOrder.verify(airbyteMessageConsumerMock).accept(message1);
-    inOrder.verify(airbyteMessageConsumerMock).close();
-    inOrder.verifyNoMoreInteractions();
+    try (final AirbyteMessageConsumer airbyteMessageConsumerMock = mock(AirbyteMessageConsumer.class)) {
+      doThrow(new IOException("error")).when(airbyteMessageConsumerMock).accept(message1);
+      assertThrows(IOException.class, () -> IntegrationRunner.consumeWriteStream(airbyteMessageConsumerMock));
+      final InOrder inOrder = inOrder(airbyteMessageConsumerMock);
+      inOrder.verify(airbyteMessageConsumerMock).accept(message1);
+      inOrder.verifyNoMoreInteractions();
+    }
   }
 
 }