Use MessageMigration for Source Connection Check. (#17656)

* More AirbyteVersion references fix

* Propagate protocol version from sourceDef to SchedulerClient

* Propagate protocol version to LauncherConfig

* Add VersionedMigratorFactory

* Update VersionedAirbyteStreamFactory

* Fix Version Json serialization/deserialization

* Plug message migration in CheckConnection for Sources

Author: gosusnp

airbyte-commons-protocol/src/main/java/io/airbyte/commons/protocol/AirbyteMessageSerDeProvider.java

@@ -7,7 +7,7 @@
 import com.google.common.annotations.VisibleForTesting;
 import io.airbyte.commons.protocol.serde.AirbyteMessageDeserializer;
 import io.airbyte.commons.protocol.serde.AirbyteMessageSerializer;
-import io.airbyte.commons.version.AirbyteVersion;
+import io.airbyte.commons.version.Version;
 import jakarta.annotation.PostConstruct;
 import jakarta.inject.Singleton;
 import java.util.Collections;
@@ -51,14 +51,14 @@ public void initialize() {
   /**
    * Returns the Deserializer for the version if known else empty
    */
-  public Optional<AirbyteMessageDeserializer<?>> getDeserializer(final AirbyteVersion version) {
+  public Optional<AirbyteMessageDeserializer<?>> getDeserializer(final Version version) {
     return Optional.ofNullable(deserializers.get(version.getMajorVersion()));
   }
 
   /**
    * Returns the Serializer for the version if known else empty
    */
-  public Optional<AirbyteMessageSerializer<?>> getSerializer(final AirbyteVersion version) {
+  public Optional<AirbyteMessageSerializer<?>> getSerializer(final Version version) {
     return Optional.ofNullable(serializers.get(version.getMajorVersion()));
   }
 

airbyte-commons-protocol/src/main/java/io/airbyte/commons/protocol/AirbyteMessageVersionedMigratorFactory.java

@@ -0,0 +1,26 @@
+/*
+ * Copyright (c) 2022 Airbyte, Inc., all rights reserved.
+ */
+
+package io.airbyte.commons.protocol;
+
+import io.airbyte.commons.version.Version;
+import jakarta.inject.Singleton;
+
+/**
+ * Factory to build AirbyteMessageVersionedMigrator
+ */
+@Singleton
+public class AirbyteMessageVersionedMigratorFactory {
+
+  private final AirbyteMessageMigrator migrator;
+
+  public AirbyteMessageVersionedMigratorFactory(final AirbyteMessageMigrator migrator) {
+    this.migrator = migrator;
+  }
+
+  public <T> AirbyteMessageVersionedMigrator<T> getVersionedMigrator(final Version version) {
+    return new AirbyteMessageVersionedMigrator<>(this.migrator, version);
+  }
+
+}

airbyte-commons-protocol/src/main/java/io/airbyte/commons/protocol/migrations/AirbyteMessageMigrationV0.java

@@ -7,13 +7,13 @@
 import io.airbyte.commons.json.Jsons;
 import io.airbyte.commons.version.Version;
 import io.airbyte.protocol.models.AirbyteMessage;
+import jakarta.inject.Singleton;
 
 /**
  * Demo migration to illustrate the template. This should be deleted once we added the v0 to v1
  * migration.
  */
-// NOTE, to actually wire this migration, uncomment the annotation
-// @Singleton
+@Singleton
 public class AirbyteMessageMigrationV0
     implements AirbyteMessageMigration<AirbyteMessage, io.airbyte.protocol.models.v0.AirbyteMessage> {
 

airbyte-commons-protocol/src/main/java/io/airbyte/commons/protocol/serde/AirbyteMessageDeserializer.java

@@ -5,12 +5,12 @@
 package io.airbyte.commons.protocol.serde;
 
 import com.fasterxml.jackson.databind.JsonNode;
-import io.airbyte.commons.version.AirbyteVersion;
+import io.airbyte.commons.version.Version;
 
 public interface AirbyteMessageDeserializer<T> {
 
   T deserialize(final JsonNode json);
 
-  AirbyteVersion getTargetVersion();
+  Version getTargetVersion();
 
 }

airbyte-commons-protocol/src/main/java/io/airbyte/commons/protocol/serde/AirbyteMessageSerializer.java

@@ -4,12 +4,12 @@
 
 package io.airbyte.commons.protocol.serde;
 
-import io.airbyte.commons.version.AirbyteVersion;
+import io.airbyte.commons.version.Version;
 
 public interface AirbyteMessageSerializer<T> {
 
   String serialize(final T message);
 
-  AirbyteVersion getTargetVersion();
+  Version getTargetVersion();
 
 }

airbyte-commons-protocol/src/test/java/io/airbyte/commons/protocol/AirbyteMessageMigratorMicronautTest.java

@@ -9,6 +9,7 @@
 import io.micronaut.test.extensions.junit5.annotation.MicronautTest;
 import jakarta.inject.Inject;
 import java.util.HashSet;
+import java.util.List;
 import org.junit.jupiter.api.Test;
 
 @MicronautTest
@@ -21,7 +22,7 @@ class AirbyteMessageMigratorMicronautTest {
   void testMigrationInjection() {
     // This should contain the list of all the supported majors of the airbyte protocol except the most
     // recent one since the migrations themselves are keyed on the lower version.
-    assertEquals(new HashSet<>(), messageMigrator.getMigrationKeys());
+    assertEquals(new HashSet<>(List.of("0")), messageMigrator.getMigrationKeys());
   }
 
 }

airbyte-commons-protocol/src/test/java/io/airbyte/commons/protocol/AirbyteMessageSerDeProviderTest.java

@@ -11,7 +11,7 @@
 
 import io.airbyte.commons.protocol.serde.AirbyteMessageDeserializer;
 import io.airbyte.commons.protocol.serde.AirbyteMessageSerializer;
-import io.airbyte.commons.version.AirbyteVersion;
+import io.airbyte.commons.version.Version;
 import java.util.Optional;
 import org.junit.jupiter.api.BeforeEach;
 import org.junit.jupiter.api.Test;
@@ -29,55 +29,55 @@ class AirbyteMessageSerDeProviderTest {
   void beforeEach() {
     serDeProvider = new AirbyteMessageSerDeProvider();
 
-    deserV0 = buildDeserializer(new AirbyteVersion("0.1.0"));
-    deserV1 = buildDeserializer(new AirbyteVersion("1.1.0"));
+    deserV0 = buildDeserializer(new Version("0.1.0"));
+    deserV1 = buildDeserializer(new Version("1.1.0"));
     serDeProvider.registerDeserializer(deserV0);
     serDeProvider.registerDeserializer(deserV1);
 
-    serV0 = buildSerializer(new AirbyteVersion("0.2.0"));
-    serV1 = buildSerializer(new AirbyteVersion("1.0.0"));
+    serV0 = buildSerializer(new Version("0.2.0"));
+    serV1 = buildSerializer(new Version("1.0.0"));
     serDeProvider.registerSerializer(serV0);
     serDeProvider.registerSerializer(serV1);
   }
 
   @Test
   void testGetDeserializer() {
-    assertEquals(Optional.of(deserV0), serDeProvider.getDeserializer(new AirbyteVersion("0.1.0")));
-    assertEquals(Optional.of(deserV0), serDeProvider.getDeserializer(new AirbyteVersion("0.2.0")));
-    assertEquals(Optional.of(deserV1), serDeProvider.getDeserializer(new AirbyteVersion("1.1.0")));
-    assertEquals(Optional.empty(), serDeProvider.getDeserializer(new AirbyteVersion("2.0.0")));
+    assertEquals(Optional.of(deserV0), serDeProvider.getDeserializer(new Version("0.1.0")));
+    assertEquals(Optional.of(deserV0), serDeProvider.getDeserializer(new Version("0.2.0")));
+    assertEquals(Optional.of(deserV1), serDeProvider.getDeserializer(new Version("1.1.0")));
+    assertEquals(Optional.empty(), serDeProvider.getDeserializer(new Version("2.0.0")));
   }
 
   @Test
   void testGetSerializer() {
-    assertEquals(Optional.of(serV0), serDeProvider.getSerializer(new AirbyteVersion("0.1.0")));
-    assertEquals(Optional.of(serV1), serDeProvider.getSerializer(new AirbyteVersion("1.0.0")));
-    assertEquals(Optional.empty(), serDeProvider.getSerializer(new AirbyteVersion("3.2.0")));
+    assertEquals(Optional.of(serV0), serDeProvider.getSerializer(new Version("0.1.0")));
+    assertEquals(Optional.of(serV1), serDeProvider.getSerializer(new Version("1.0.0")));
+    assertEquals(Optional.empty(), serDeProvider.getSerializer(new Version("3.2.0")));
   }
 
   @Test
   void testRegisterDeserializerShouldFailOnVersionCollision() {
-    AirbyteMessageDeserializer<?> deser = buildDeserializer(new AirbyteVersion("0.2.0"));
+    AirbyteMessageDeserializer<?> deser = buildDeserializer(new Version("0.2.0"));
     assertThrows(RuntimeException.class, () -> {
       serDeProvider.registerDeserializer(deser);
     });
   }
 
   @Test
   void testRegisterSerializerShouldFailOnVersionCollision() {
-    AirbyteMessageSerializer<?> ser = buildSerializer(new AirbyteVersion("0.5.0"));
+    AirbyteMessageSerializer<?> ser = buildSerializer(new Version("0.5.0"));
     assertThrows(RuntimeException.class, () -> {
       serDeProvider.registerSerializer(ser);
     });
   }
 
-  private <T> AirbyteMessageDeserializer<T> buildDeserializer(AirbyteVersion version) {
+  private <T> AirbyteMessageDeserializer<T> buildDeserializer(Version version) {
     final AirbyteMessageDeserializer<T> deser = mock(AirbyteMessageDeserializer.class);
     when(deser.getTargetVersion()).thenReturn(version);
     return deser;
   }
 
-  private <T> AirbyteMessageSerializer<T> buildSerializer(AirbyteVersion version) {
+  private <T> AirbyteMessageSerializer<T> buildSerializer(Version version) {
     final AirbyteMessageSerializer<T> ser = mock(AirbyteMessageSerializer.class);
     when(ser.getTargetVersion()).thenReturn(version);
     return ser;

airbyte-commons-worker/src/main/java/io/airbyte/workers/internal/DefaultAirbyteStreamFactory.java

@@ -32,7 +32,7 @@ public class DefaultAirbyteStreamFactory implements AirbyteStreamFactory {
 
   private final MdcScope.Builder containerLogMdcBuilder;
   private final AirbyteProtocolPredicate protocolValidator;
-  private final Logger logger;
+  protected final Logger logger;
 
   public DefaultAirbyteStreamFactory() {
     this(MdcScope.DEFAULT_BUILDER);

airbyte-commons-worker/src/main/java/io/airbyte/workers/internal/VersionedAirbyteStreamFactory.java

@@ -7,8 +7,11 @@
 import com.fasterxml.jackson.databind.JsonNode;
 import io.airbyte.commons.json.Jsons;
 import io.airbyte.commons.logging.MdcScope;
+import io.airbyte.commons.protocol.AirbyteMessageSerDeProvider;
 import io.airbyte.commons.protocol.AirbyteMessageVersionedMigrator;
+import io.airbyte.commons.protocol.AirbyteMessageVersionedMigratorFactory;
 import io.airbyte.commons.protocol.serde.AirbyteMessageDeserializer;
+import io.airbyte.commons.version.Version;
 import io.airbyte.protocol.models.AirbyteMessage;
 import java.util.stream.Stream;
 import org.slf4j.Logger;
@@ -26,19 +29,23 @@ public class VersionedAirbyteStreamFactory<T> extends DefaultAirbyteStreamFactor
 
   private final AirbyteMessageDeserializer<T> deserializer;
   private final AirbyteMessageVersionedMigrator<T> migrator;
+  private final Version protocolVersion;
 
-  public VersionedAirbyteStreamFactory(final AirbyteMessageDeserializer<T> deserializer,
-                                       final AirbyteMessageVersionedMigrator<T> migrator) {
-    this(deserializer, migrator, MdcScope.DEFAULT_BUILDER);
+  public VersionedAirbyteStreamFactory(final AirbyteMessageSerDeProvider serDeProvider,
+                                       final AirbyteMessageVersionedMigratorFactory migratorFactory,
+                                       final Version protocolVersion) {
+    this(serDeProvider, migratorFactory, protocolVersion, MdcScope.DEFAULT_BUILDER);
   }
 
-  public VersionedAirbyteStreamFactory(final AirbyteMessageDeserializer<T> deserializer,
-                                       final AirbyteMessageVersionedMigrator<T> migrator,
+  public VersionedAirbyteStreamFactory(final AirbyteMessageSerDeProvider serDeProvider,
+                                       final AirbyteMessageVersionedMigratorFactory migratorFactory,
+                                       final Version protocolVersion,
                                        final MdcScope.Builder containerLogMdcBuilder) {
     // TODO AirbyteProtocolPredicate needs to be updated to be protocol version aware
     super(new AirbyteProtocolPredicate(), LOGGER, containerLogMdcBuilder);
-    this.deserializer = deserializer;
-    this.migrator = migrator;
+    this.deserializer = (AirbyteMessageDeserializer<T>) serDeProvider.getDeserializer(protocolVersion).orElseThrow();
+    this.migrator = migratorFactory.getVersionedMigrator(protocolVersion);
+    this.protocolVersion = protocolVersion;
   }
 
   @Override
@@ -47,7 +54,7 @@ protected Stream<AirbyteMessage> toAirbyteMessage(final JsonNode json) {
       final io.airbyte.protocol.models.v0.AirbyteMessage message = migrator.upgrade(deserializer.deserialize(json));
       return Stream.of(convert(message));
     } catch (RuntimeException e) {
-      LOGGER.warn("Failed to upgrade a message from version {}: {}", migrator.getVersion(), Jsons.serialize(json));
+      logger.warn("Failed to upgrade a message from version {}: {}", protocolVersion, Jsons.serialize(json), e);
       return Stream.empty();
     }
   }

airbyte-commons-worker/src/main/resources/workers_models/IntegrationLauncherConfig.yaml

@@ -16,3 +16,6 @@ properties:
     type: integer
   dockerImage:
     type: string
+  protocolVersion:
+    type: object
+    existingJavaType: io.airbyte.commons.version.Version

airbyte-commons/src/main/java/io/airbyte/commons/version/Version.java

@@ -4,13 +4,17 @@
 
 package io.airbyte.commons.version;
 
+import com.fasterxml.jackson.databind.annotation.JsonDeserialize;
+import com.fasterxml.jackson.databind.annotation.JsonSerialize;
 import com.google.common.base.Preconditions;
 import java.util.Objects;
 
 /**
  * A semVer Version class that allows "dev" as a version.
  */
 @SuppressWarnings({"PMD.AvoidFieldNameMatchingTypeName", "PMD.ConstructorCallsOverridableMethod"})
+@JsonDeserialize(using = VersionDeserializer.class)
+@JsonSerialize(using = VersionSerializer.class)
 public class Version {
 
   public static final String DEV_VERSION_PREFIX = "dev";

airbyte-commons/src/main/java/io/airbyte/commons/version/VersionDeserializer.java

@@ -0,0 +1,31 @@
+/*
+ * Copyright (c) 2022 Airbyte, Inc., all rights reserved.
+ */
+
+package io.airbyte.commons.version;
+
+import com.fasterxml.jackson.core.JacksonException;
+import com.fasterxml.jackson.core.JsonParser;
+import com.fasterxml.jackson.databind.DeserializationContext;
+import com.fasterxml.jackson.databind.JsonNode;
+import com.fasterxml.jackson.databind.deser.std.StdDeserializer;
+import java.io.IOException;
+
+public class VersionDeserializer extends StdDeserializer<Version> {
+
+  public VersionDeserializer() {
+    this(null);
+  }
+
+  public VersionDeserializer(Class<?> vc) {
+    super(vc);
+  }
+
+  @Override
+  public Version deserialize(JsonParser p, DeserializationContext ctxt) throws IOException, JacksonException {
+    final JsonNode node = p.getCodec().readTree(p);
+    final String v = node.get("version").asText();
+    return new Version(v);
+  }
+
+}

airbyte-commons/src/main/java/io/airbyte/commons/version/VersionSerializer.java

@@ -0,0 +1,29 @@
+/*
+ * Copyright (c) 2022 Airbyte, Inc., all rights reserved.
+ */
+
+package io.airbyte.commons.version;
+
+import com.fasterxml.jackson.core.JsonGenerator;
+import com.fasterxml.jackson.databind.SerializerProvider;
+import com.fasterxml.jackson.databind.ser.std.StdSerializer;
+import java.io.IOException;
+
+public class VersionSerializer extends StdSerializer<Version> {
+
+  public VersionSerializer() {
+    this(null);
+  }
+
+  public VersionSerializer(Class<Version> t) {
+    super(t);
+  }
+
+  @Override
+  public void serialize(Version value, JsonGenerator gen, SerializerProvider provider) throws IOException {
+    gen.writeStartObject();
+    gen.writeStringField("version", value.version);
+    gen.writeEndObject();
+  }
+
+}

airbyte-commons/src/test/java/io/airbyte/commons/version/VersionTest.java

@@ -0,0 +1,28 @@
+/*
+ * Copyright (c) 2022 Airbyte, Inc., all rights reserved.
+ */
+
+package io.airbyte.commons.version;
+
+import static org.junit.jupiter.api.Assertions.assertEquals;
+
+import io.airbyte.commons.json.Jsons;
+import org.junit.jupiter.api.Test;
+
+class VersionTest {
+
+  @Test
+  void testJsonSerializationDeserialization() {
+    final String jsonString = """
+                              {"version": "1.2.3"}
+                              """;
+    final Version expectedVersion = new Version("1.2.3");
+
+    final Version deserializedVersion = Jsons.deserialize(jsonString, Version.class);
+    assertEquals(expectedVersion, deserializedVersion);
+
+    final Version deserializedVersionLoop = Jsons.deserialize(Jsons.serialize(deserializedVersion), Version.class);
+    assertEquals(expectedVersion, deserializedVersionLoop);
+  }
+
+}

airbyte-config/config-models/src/main/resources/types/JobCheckConnectionConfig.yaml

@@ -15,3 +15,6 @@ properties:
     existingJavaType: com.fasterxml.jackson.databind.JsonNode
   dockerImage:
     type: string
+  protocolVersion:
+    type: object
+    existingJavaType: io.airbyte.commons.version.Version