airbyte-common-workers: Collect trace message on failed connection_status

airbyte-commons-worker/src/main/java/io/airbyte/workers/WorkerUtils.java

@@ -5,7 +5,7 @@
 package io.airbyte.workers;
 
 import com.fasterxml.jackson.databind.JsonNode;
-import io.airbyte.config.ConnectorJobOutput;
+import io.airbyte.commons.io.IOs;
 import io.airbyte.config.ConnectorJobOutput.OutputType;
 import io.airbyte.config.FailureReason;
 import io.airbyte.config.StandardSyncInput;
@@ -15,9 +15,11 @@
 import io.airbyte.protocol.models.AirbyteMessage.Type;
 import io.airbyte.protocol.models.AirbyteStreamNameNamespacePair;
 import io.airbyte.protocol.models.AirbyteTraceMessage;
-import io.airbyte.workers.exception.WorkerException;
 import io.airbyte.workers.helper.FailureHelper;
 import io.airbyte.workers.helper.FailureHelper.ConnectorCommand;
+import io.airbyte.workers.internal.AirbyteStreamFactory;
+import java.io.IOException;
+import java.io.InputStream;
 import java.time.Duration;
 import java.time.temporal.ChronoUnit;
 import java.util.ArrayList;
@@ -107,28 +109,43 @@ public static WorkerDestinationConfig syncToWorkerDestinationConfig(final Standa
         .withState(sync.getState());
   }
 
-  public static ConnectorJobOutput getJobFailureOutputOrThrow(final OutputType outputType,
-                                                              final Map<Type, List<AirbyteMessage>> messagesByType,
-                                                              final String defaultErrorMessage)
-      throws WorkerException {
-    final Optional<AirbyteTraceMessage> traceMessage =
-        messagesByType.getOrDefault(Type.TRACE, new ArrayList<>()).stream()
-            .map(AirbyteMessage::getTrace)
-            .filter(trace -> trace.getType() == AirbyteTraceMessage.Type.ERROR)
-            .findFirst();
+  private static ConnectorCommand getConnectorCommandFromOutputType(final OutputType outputType) {
+    return switch (outputType) {
+      case SPEC -> ConnectorCommand.SPEC;
+      case CHECK_CONNECTION -> ConnectorCommand.CHECK;
+      case DISCOVER_CATALOG_ID -> ConnectorCommand.DISCOVER;
+    };
+  }
+
+  private static Optional<AirbyteTraceMessage> getTraceMessageFromMessagesByType(final Map<Type, List<AirbyteMessage>> messagesByType) {
+    return messagesByType.getOrDefault(Type.TRACE, new ArrayList<>()).stream()
+        .map(AirbyteMessage::getTrace)
+        .filter(trace -> trace.getType() == AirbyteTraceMessage.Type.ERROR)
+        .findFirst();
+  }
+
+  public static Map<Type, List<AirbyteMessage>> getMessagesByType(Process process, AirbyteStreamFactory streamFactory, int timeOut)
+      throws IOException {
+    final Map<Type, List<AirbyteMessage>> messagesByType;
+    try (final InputStream stdout = process.getInputStream()) {
+      messagesByType = streamFactory.create(IOs.newBufferedReader(stdout))
+          .collect(Collectors.groupingBy(AirbyteMessage::getType));
+
+      WorkerUtils.gentleClose(process, timeOut, TimeUnit.MINUTES);
+      return messagesByType;
+    }
+  }
 
+  public static Optional<FailureReason> getJobFailureReasonFromMessages(final OutputType outputType,
+                                                                        final Map<Type, List<AirbyteMessage>> messagesByType) {
+    Optional<AirbyteTraceMessage> traceMessage = getTraceMessageFromMessagesByType(messagesByType);
     if (traceMessage.isPresent()) {
-      final ConnectorCommand connectorCommand = switch (outputType) {
-        case SPEC -> ConnectorCommand.SPEC;
-        case CHECK_CONNECTION -> ConnectorCommand.CHECK;
-        case DISCOVER_CATALOG_ID -> ConnectorCommand.DISCOVER;
-      };
-
-      final FailureReason failureReason = FailureHelper.connectorCommandFailure(traceMessage.get(), null, null, connectorCommand);
-      return new ConnectorJobOutput().withOutputType(outputType).withFailureReason(failureReason);
+      final ConnectorCommand connectorCommand = getConnectorCommandFromOutputType(outputType);
+      return Optional.of(FailureHelper.connectorCommandFailure(traceMessage.get(), null, null, connectorCommand));
+    } else {
+      return Optional.empty();
     }
 
-    throw new WorkerException(defaultErrorMessage);
   }
 
   public static Map<AirbyteStreamNameNamespacePair, JsonNode> mapStreamNamesToSchemas(final StandardSyncInput syncInput) {

airbyte-commons-worker/src/main/java/io/airbyte/workers/general/DefaultCheckConnectionWorker.java

@@ -9,11 +9,11 @@
 
 import datadog.trace.api.Trace;
 import io.airbyte.commons.enums.Enums;
-import io.airbyte.commons.io.IOs;
 import io.airbyte.commons.io.LineGobbler;
 import io.airbyte.commons.json.Jsons;
 import io.airbyte.config.ConnectorJobOutput;
 import io.airbyte.config.ConnectorJobOutput.OutputType;
+import io.airbyte.config.FailureReason;
 import io.airbyte.config.StandardCheckConnectionInput;
 import io.airbyte.config.StandardCheckConnectionOutput;
 import io.airbyte.config.StandardCheckConnectionOutput.Status;
@@ -27,14 +27,11 @@
 import io.airbyte.workers.internal.AirbyteStreamFactory;
 import io.airbyte.workers.internal.DefaultAirbyteStreamFactory;
 import io.airbyte.workers.process.IntegrationLauncher;
-import java.io.InputStream;
 import java.nio.file.Path;
 import java.util.ArrayList;
 import java.util.List;
 import java.util.Map;
 import java.util.Optional;
-import java.util.concurrent.TimeUnit;
-import java.util.stream.Collectors;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
 
@@ -67,38 +64,33 @@ public ConnectorJobOutput run(final StandardCheckConnectionInput input, final Pa
           jobRoot,
           WorkerConstants.SOURCE_CONFIG_JSON_FILENAME,
           Jsons.serialize(input.getConnectionConfiguration()));
+      final int exitCode = process.exitValue();
+      LOGGER.debug("Check connection job subprocess finished with exit code {}", exitCode);
+      final ConnectorJobOutput jobOutput = new ConnectorJobOutput().withOutputType(OutputType.CHECK_CONNECTION);
 
       LineGobbler.gobble(process.getErrorStream(), LOGGER::error);
 
-      final Map<Type, List<AirbyteMessage>> messagesByType;
-      try (final InputStream stdout = process.getInputStream()) {
-        messagesByType = streamFactory.create(IOs.newBufferedReader(stdout))
-            .collect(Collectors.groupingBy(AirbyteMessage::getType));
-
-        WorkerUtils.gentleClose(process, 1, TimeUnit.MINUTES);
-      }
-
-      final int exitCode = process.exitValue();
-      final Optional<AirbyteConnectionStatus> status = messagesByType
+      final Map<Type, List<AirbyteMessage>> messagesByType = WorkerUtils.getMessagesByType(process, streamFactory, 30);
+      final Optional<AirbyteConnectionStatus> connectionStatus = messagesByType
           .getOrDefault(Type.CONNECTION_STATUS, new ArrayList<>()).stream()
           .map(AirbyteMessage::getConnectionStatus)
           .findFirst();
 
-      if (status.isPresent() && exitCode == 0) {
-        final StandardCheckConnectionOutput output = new StandardCheckConnectionOutput()
-            .withStatus(Enums.convertTo(status.get().getStatus(), Status.class))
-            .withMessage(status.get().getMessage());
+      final Optional<FailureReason> failureReason = WorkerUtils.getJobFailureReasonFromMessages(OutputType.CHECK_CONNECTION, messagesByType);
+      failureReason.ifPresent(jobOutput::setFailureReason);
 
-        LOGGER.debug("Check connection job subprocess finished with exit code {}", exitCode);
+      if (connectionStatus.isPresent()) {
+        final StandardCheckConnectionOutput output = new StandardCheckConnectionOutput()
+            .withStatus(Enums.convertTo(connectionStatus.get().getStatus(), Status.class))
+            .withMessage(connectionStatus.get().getMessage());
         LOGGER.debug("Check connection job received output: {}", output);
-        LineGobbler.endSection("CHECK");
-        return new ConnectorJobOutput().withOutputType(OutputType.CHECK_CONNECTION).withCheckConnection(output);
+        jobOutput.setCheckConnection(output);
       } else {
-        final String message = String.format("Error checking connection, status: %s, exit code: %d", status, exitCode);
-        LOGGER.error(message);
-
-        return WorkerUtils.getJobFailureOutputOrThrow(OutputType.CHECK_CONNECTION, messagesByType, message);
+        if (failureReason.isEmpty()) {
+          throw new WorkerException("Error checking connection status: no status nor failure reason were outputted.");
+        }
       }
+      return jobOutput;
 
     } catch (final Exception e) {
       ApmTraceUtils.addExceptionToTrace(e);

airbyte-commons-worker/src/main/java/io/airbyte/workers/general/DefaultDiscoverCatalogWorker.java

@@ -10,11 +10,11 @@
 import static io.airbyte.metrics.lib.ApmTraceConstants.WORKER_OPERATION_NAME;
 
 import datadog.trace.api.Trace;
-import io.airbyte.commons.io.IOs;
 import io.airbyte.commons.io.LineGobbler;
 import io.airbyte.commons.json.Jsons;
 import io.airbyte.config.ConnectorJobOutput;
 import io.airbyte.config.ConnectorJobOutput.OutputType;
+import io.airbyte.config.FailureReason;
 import io.airbyte.config.StandardDiscoverCatalogInput;
 import io.airbyte.config.persistence.ConfigRepository;
 import io.airbyte.metrics.lib.ApmTraceUtils;
@@ -27,16 +27,13 @@
 import io.airbyte.workers.internal.AirbyteStreamFactory;
 import io.airbyte.workers.internal.DefaultAirbyteStreamFactory;
 import io.airbyte.workers.process.IntegrationLauncher;
-import java.io.InputStream;
 import java.nio.file.Path;
 import java.util.ArrayList;
 import java.util.HashMap;
 import java.util.List;
 import java.util.Map;
 import java.util.Optional;
 import java.util.UUID;
-import java.util.concurrent.TimeUnit;
-import java.util.stream.Collectors;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
 
@@ -73,43 +70,38 @@ public ConnectorJobOutput run(final StandardDiscoverCatalogInput discoverSchemaI
           jobRoot,
           WorkerConstants.SOURCE_CONFIG_JSON_FILENAME,
           Jsons.serialize(discoverSchemaInput.getConnectionConfiguration()));
+      final int exitCode = process.exitValue();
+      final String exitCodeMessage = String.format("Discover job subprocess finished with exit code %s", exitCode);
+      LOGGER.debug(exitCodeMessage);
+      final ConnectorJobOutput jobOutput = new ConnectorJobOutput().withOutputType(OutputType.DISCOVER_CATALOG_ID);
 
       LineGobbler.gobble(process.getErrorStream(), LOGGER::error);
 
-      final Map<Type, List<AirbyteMessage>> messagesByType;
-
-      try (final InputStream stdout = process.getInputStream()) {
-        messagesByType = streamFactory.create(IOs.newBufferedReader(stdout))
-            .collect(Collectors.groupingBy(AirbyteMessage::getType));
-
-        WorkerUtils.gentleClose(process, 30, TimeUnit.MINUTES);
-      }
+      final Map<Type, List<AirbyteMessage>> messagesByType = WorkerUtils.getMessagesByType(process, streamFactory, 30);
 
       final Optional<AirbyteCatalog> catalog = messagesByType
           .getOrDefault(Type.CATALOG, new ArrayList<>()).stream()
           .map(AirbyteMessage::getCatalog)
           .findFirst();
 
-      final int exitCode = process.exitValue();
-      if (exitCode == 0) {
-        if (catalog.isEmpty()) {
-          throw new WorkerException("Integration failed to output a catalog struct.");
-        }
+      final Optional<FailureReason> failureReason = WorkerUtils.getJobFailureReasonFromMessages(OutputType.DISCOVER_CATALOG_ID, messagesByType);
+      failureReason.ifPresent(jobOutput::setFailureReason);
 
+      if (catalog.isPresent()) {
         final UUID catalogId =
             configRepository.writeActorCatalogFetchEvent(catalog.get(),
                 // NOTE: sourceId is marked required in the OpenAPI config but the code generator doesn't enforce
                 // it, so we check again here.
                 discoverSchemaInput.getSourceId() == null ? null : UUID.fromString(discoverSchemaInput.getSourceId()),
                 discoverSchemaInput.getConnectorVersion(),
                 discoverSchemaInput.getConfigHash());
-        return new ConnectorJobOutput().withOutputType(OutputType.DISCOVER_CATALOG_ID).withDiscoverCatalogId(catalogId);
+        jobOutput.setDiscoverCatalogId(catalogId);
       } else {
-        return WorkerUtils.getJobFailureOutputOrThrow(
-            OutputType.DISCOVER_CATALOG_ID,
-            messagesByType,
-            String.format("Discover job subprocess finished with exit code %s", exitCode));
+        if (failureReason.isEmpty()) {
+          throw new WorkerException("Integration failed to output a catalog struct.");
+        }
       }
+      return jobOutput;
     } catch (final WorkerException e) {
       ApmTraceUtils.addExceptionToTrace(e);
       throw e;

airbyte-commons-worker/src/main/java/io/airbyte/workers/general/DefaultGetSpecWorker.java

@@ -9,10 +9,10 @@
 import static io.airbyte.metrics.lib.ApmTraceConstants.WORKER_OPERATION_NAME;
 
 import datadog.trace.api.Trace;
-import io.airbyte.commons.io.IOs;
 import io.airbyte.commons.io.LineGobbler;
 import io.airbyte.config.ConnectorJobOutput;
 import io.airbyte.config.ConnectorJobOutput.OutputType;
+import io.airbyte.config.FailureReason;
 import io.airbyte.config.JobGetSpecConfig;
 import io.airbyte.metrics.lib.ApmTraceUtils;
 import io.airbyte.protocol.models.AirbyteMessage;
@@ -23,14 +23,11 @@
 import io.airbyte.workers.internal.AirbyteStreamFactory;
 import io.airbyte.workers.internal.DefaultAirbyteStreamFactory;
 import io.airbyte.workers.process.IntegrationLauncher;
-import java.io.InputStream;
 import java.nio.file.Path;
 import java.util.ArrayList;
 import java.util.List;
 import java.util.Map;
 import java.util.Optional;
-import java.util.concurrent.TimeUnit;
-import java.util.stream.Collectors;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
 
@@ -40,7 +37,6 @@ public class DefaultGetSpecWorker implements GetSpecWorker {
 
   private final IntegrationLauncher integrationLauncher;
   private final AirbyteStreamFactory streamFactory;
-
   private Process process;
 
   public DefaultGetSpecWorker(final IntegrationLauncher integrationLauncher,
@@ -59,39 +55,32 @@ public ConnectorJobOutput run(final JobGetSpecConfig config, final Path jobRoot)
     ApmTraceUtils.addTagsToTrace(Map.of(JOB_ROOT_KEY, jobRoot, DOCKER_IMAGE_KEY, config.getDockerImage()));
     try {
       process = integrationLauncher.spec(jobRoot);
+      final int exitCode = process.exitValue();
+      final String exitCodeMessage = String.format("Spec job subprocess finished with exit code %s", exitCode);
+      LOGGER.debug(exitCodeMessage);
 
+      final ConnectorJobOutput jobOutput = new ConnectorJobOutput().withOutputType(OutputType.SPEC);
       LineGobbler.gobble(process.getErrorStream(), LOGGER::error);
 
-      final Map<Type, List<AirbyteMessage>> messagesByType;
-      try (final InputStream stdout = process.getInputStream()) {
-        messagesByType = streamFactory.create(IOs.newBufferedReader(stdout))
-            .collect(Collectors.groupingBy(AirbyteMessage::getType));
-
-        // todo (cgardens) - let's pre-fetch the images outside of the worker so we don't need account for
-        // this.
-        // retrieving spec should generally be instantaneous, but since docker images might not be pulled
-        // it could take a while longer depending on internet conditions as well.
-        WorkerUtils.gentleClose(process, 30, TimeUnit.MINUTES);
-      }
+      final Map<Type, List<AirbyteMessage>> messagesByType = WorkerUtils.getMessagesByType(process, streamFactory, 30);
 
       final Optional<ConnectorSpecification> spec = messagesByType
           .getOrDefault(Type.SPEC, new ArrayList<>()).stream()
           .map(AirbyteMessage::getSpec)
           .findFirst();
 
-      final int exitCode = process.exitValue();
-      if (exitCode == 0) {
-        if (spec.isEmpty()) {
-          throw new WorkerException("integration failed to output a spec struct.");
-        }
+      final Optional<FailureReason> failureReason = WorkerUtils.getJobFailureReasonFromMessages(OutputType.SPEC, messagesByType);
+      failureReason.ifPresent(jobOutput::setFailureReason);
 
-        return new ConnectorJobOutput().withOutputType(OutputType.SPEC).withSpec(spec.get());
+      if (spec.isPresent()) {
+        jobOutput.setSpec(spec.get());
       } else {
-        return WorkerUtils.getJobFailureOutputOrThrow(
-            OutputType.SPEC,
-            messagesByType,
-            String.format("Spec job subprocess finished with exit code %s", exitCode));
+        if (failureReason.isEmpty()) {
+          throw new WorkerException("Integration failed to output a spec struct and did not output a failure reason.");
+        }
       }
+
+      return jobOutput;
     } catch (final Exception e) {
       throw new WorkerException(String.format("Error while getting spec from image %s", config.getDockerImage()), e);
     }