Update airbyte protocol migration (#20745)

* Extract MigrationContainer from AirbyteMessageMigrator

* Add ConfiguredAirbyteCatalogMigrations

* Add ConfiguredAirbyteCatalog to AirbyteMessageMigrations

* Enable ConfiguredAirbyteCatalog migration

* Fix tests

* Remove extra this.

* Add missing docs

* Typo

Co-authored-by: Edward Gao <[email protected]>

Author: gosusnp

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

@@ -6,15 +6,14 @@
 
 import com.google.common.annotations.VisibleForTesting;
 import io.airbyte.commons.protocol.migrations.AirbyteMessageMigration;
+import io.airbyte.commons.protocol.migrations.MigrationContainer;
 import io.airbyte.commons.version.Version;
+import io.airbyte.protocol.models.ConfiguredAirbyteCatalog;
 import jakarta.annotation.PostConstruct;
 import jakarta.inject.Singleton;
-import java.util.Collection;
-import java.util.Collections;
 import java.util.List;
+import java.util.Optional;
 import java.util.Set;
-import java.util.SortedMap;
-import java.util.TreeMap;
 
 /**
  * AirbyteProtocol Message Migrator
@@ -25,104 +24,59 @@
 @Singleton
 public class AirbyteMessageMigrator {
 
-  private final List<AirbyteMessageMigration<?, ?>> migrationsToRegister;
-  private final SortedMap<String, AirbyteMessageMigration<?, ?>> migrations = new TreeMap<>();
-  private String mostRecentMajorVersion = "";
+  private final MigrationContainer<AirbyteMessageMigration<?, ?>> migrationContainer;
 
-  public AirbyteMessageMigrator(List<AirbyteMessageMigration<?, ?>> migrations) {
-    migrationsToRegister = migrations;
-  }
-
-  public AirbyteMessageMigrator() {
-    this(Collections.emptyList());
+  public AirbyteMessageMigrator(final List<AirbyteMessageMigration<?, ?>> migrations) {
+    migrationContainer = new MigrationContainer<>(migrations);
   }
 
   @PostConstruct
   public void initialize() {
-    migrationsToRegister.forEach(this::registerMigration);
+    migrationContainer.initialize();
   }
 
   /**
    * Downgrade a message from the most recent version to the target version by chaining all the
    * required migrations
    */
-  public <PreviousVersion, CurrentVersion> PreviousVersion downgrade(final CurrentVersion message, final Version target) {
-    if (target.getMajorVersion().equals(mostRecentMajorVersion)) {
-      return (PreviousVersion) message;
-    }
-
-    Object result = message;
-    Object[] selectedMigrations = selectMigrations(target).toArray();
-    for (int i = selectedMigrations.length; i > 0; --i) {
-      result = applyDowngrade((AirbyteMessageMigration<?, ?>) selectedMigrations[i - 1], result);
-    }
-    return (PreviousVersion) result;
+  public <PreviousVersion, CurrentVersion> PreviousVersion downgrade(final CurrentVersion message,
+                                                                     final Version target,
+                                                                     final Optional<ConfiguredAirbyteCatalog> configuredAirbyteCatalog) {
+    return migrationContainer.downgrade(message, target, (migration, msg) -> applyDowngrade(migration, msg, configuredAirbyteCatalog));
   }
 
   /**
    * Upgrade a message from the source version to the most recent version by chaining all the required
    * migrations
    */
-  public <PreviousVersion, CurrentVersion> CurrentVersion upgrade(final PreviousVersion message, final Version source) {
-    if (source.getMajorVersion().equals(mostRecentMajorVersion)) {
-      return (CurrentVersion) message;
-    }
-
-    Object result = message;
-    for (var migration : selectMigrations(source)) {
-      result = applyUpgrade(migration, result);
-    }
-    return (CurrentVersion) result;
+  public <PreviousVersion, CurrentVersion> CurrentVersion upgrade(final PreviousVersion message,
+                                                                  final Version source,
+                                                                  final Optional<ConfiguredAirbyteCatalog> configuredAirbyteCatalog) {
+    return migrationContainer.upgrade(message, source, (migration, msg) -> applyUpgrade(migration, msg, configuredAirbyteCatalog));
   }
 
   public Version getMostRecentVersion() {
-    return new Version(mostRecentMajorVersion, "0", "0");
-  }
-
-  private Collection<AirbyteMessageMigration<?, ?>> selectMigrations(final Version version) {
-    final Collection<AirbyteMessageMigration<?, ?>> results = migrations.tailMap(version.getMajorVersion()).values();
-    if (results.isEmpty()) {
-      throw new RuntimeException("Unsupported migration version " + version.serialize());
-    }
-    return results;
+    return migrationContainer.getMostRecentVersion();
   }
 
   // Helper function to work around type casting
-  private <PreviousVersion, CurrentVersion> PreviousVersion applyDowngrade(final AirbyteMessageMigration<PreviousVersion, CurrentVersion> migration,
-                                                                           final Object message) {
-    return migration.downgrade((CurrentVersion) message);
+  private static <PreviousVersion, CurrentVersion> PreviousVersion applyDowngrade(final AirbyteMessageMigration<PreviousVersion, CurrentVersion> migration,
+                                                                                  final Object message,
+                                                                                  final Optional<ConfiguredAirbyteCatalog> configuredAirbyteCatalog) {
+    return migration.downgrade((CurrentVersion) message, configuredAirbyteCatalog);
   }
 
   // Helper function to work around type casting
-  private <PreviousVersion, CurrentVersion> CurrentVersion applyUpgrade(final AirbyteMessageMigration<PreviousVersion, CurrentVersion> migration,
-                                                                        final Object message) {
-    return migration.upgrade((PreviousVersion) message);
-  }
-
-  /**
-   * Store migration in a sorted map key by the major of the lower version of the migration.
-   *
-   * The goal is to be able to retrieve the list of migrations to apply to get to/from a given
-   * version. We are only keying on the lower version because the right side (most recent version of
-   * the migration range) is always current version.
-   */
-  @VisibleForTesting
-  void registerMigration(final AirbyteMessageMigration<?, ?> migration) {
-    final String key = migration.getPreviousVersion().getMajorVersion();
-    if (!migrations.containsKey(key)) {
-      migrations.put(key, migration);
-      if (migration.getCurrentVersion().getMajorVersion().compareTo(mostRecentMajorVersion) > 0) {
-        mostRecentMajorVersion = migration.getCurrentVersion().getMajorVersion();
-      }
-    } else {
-      throw new RuntimeException("Trying to register a duplicated migration " + migration.getClass().getName());
-    }
+  private static <PreviousVersion, CurrentVersion> CurrentVersion applyUpgrade(final AirbyteMessageMigration<PreviousVersion, CurrentVersion> migration,
+                                                                               final Object message,
+                                                                               final Optional<ConfiguredAirbyteCatalog> configuredAirbyteCatalog) {
+    return migration.upgrade((PreviousVersion) message, configuredAirbyteCatalog);
   }
 
   // Used for inspection of the injection
   @VisibleForTesting
   Set<String> getMigrationKeys() {
-    return migrations.keySet();
+    return migrationContainer.getMigrationKeys();
   }
 
 }

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

@@ -6,6 +6,8 @@
 
 import io.airbyte.commons.version.Version;
 import io.airbyte.protocol.models.AirbyteMessage;
+import io.airbyte.protocol.models.ConfiguredAirbyteCatalog;
+import java.util.Optional;
 
 /**
  * Wraps message migration from a fixed version to the most recent version
@@ -20,12 +22,12 @@ public AirbyteMessageVersionedMigrator(final AirbyteMessageMigrator migrator, fi
     this.version = version;
   }
 
-  public OriginalMessageType downgrade(final AirbyteMessage message) {
-    return migrator.downgrade(message, version);
+  public OriginalMessageType downgrade(final AirbyteMessage message, final Optional<ConfiguredAirbyteCatalog> configuredAirbyteCatalog) {
+    return migrator.downgrade(message, version, configuredAirbyteCatalog);
   }
 
-  public AirbyteMessage upgrade(final OriginalMessageType message) {
-    return migrator.upgrade(message, version);
+  public AirbyteMessage upgrade(final OriginalMessageType message, final Optional<ConfiguredAirbyteCatalog> configuredAirbyteCatalog) {
+    return migrator.upgrade(message, version, configuredAirbyteCatalog);
   }
 
   public Version getVersion() {

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

@@ -0,0 +1,37 @@
+/*
+ * 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 AirbyteProtocolVersionedMigratorFactory {
+
+  private final AirbyteMessageMigrator airbyteMessageMigrator;
+  private final ConfiguredAirbyteCatalogMigrator configuredAirbyteCatalogMigrator;
+
+  public AirbyteProtocolVersionedMigratorFactory(final AirbyteMessageMigrator airbyteMessageMigrator,
+                                                 final ConfiguredAirbyteCatalogMigrator configuredAirbyteCatalogMigrator) {
+    this.airbyteMessageMigrator = airbyteMessageMigrator;
+    this.configuredAirbyteCatalogMigrator = configuredAirbyteCatalogMigrator;
+  }
+
+  public <T> AirbyteMessageVersionedMigrator<T> getAirbyteMessageMigrator(final Version version) {
+    return new AirbyteMessageVersionedMigrator<>(airbyteMessageMigrator, version);
+  }
+
+  public final VersionedProtocolSerializer getProtocolSerializer(final Version version) {
+    return new VersionedProtocolSerializer(configuredAirbyteCatalogMigrator, version);
+  }
+
+  public Version getMostRecentVersion() {
+    return airbyteMessageMigrator.getMostRecentVersion();
+  }
+
+}

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

@@ -0,0 +1,68 @@
+/*
+ * Copyright (c) 2022 Airbyte, Inc., all rights reserved.
+ */
+
+package io.airbyte.commons.protocol;
+
+import com.google.common.annotations.VisibleForTesting;
+import io.airbyte.commons.protocol.migrations.ConfiguredAirbyteCatalogMigration;
+import io.airbyte.commons.protocol.migrations.MigrationContainer;
+import io.airbyte.commons.version.Version;
+import jakarta.annotation.PostConstruct;
+import jakarta.inject.Singleton;
+import java.util.List;
+import java.util.Set;
+
+@Singleton
+public class ConfiguredAirbyteCatalogMigrator {
+
+  private final MigrationContainer<ConfiguredAirbyteCatalogMigration<?, ?>> migrationContainer;
+
+  public ConfiguredAirbyteCatalogMigrator(final List<ConfiguredAirbyteCatalogMigration<?, ?>> migrations) {
+    migrationContainer = new MigrationContainer<>(migrations);
+  }
+
+  @PostConstruct
+  public void initialize() {
+    migrationContainer.initialize();
+  }
+
+  /**
+   * Downgrade a message from the most recent version to the target version by chaining all the
+   * required migrations
+   */
+  public <PreviousVersion, CurrentVersion> PreviousVersion downgrade(final CurrentVersion message, final Version target) {
+    return migrationContainer.downgrade(message, target, ConfiguredAirbyteCatalogMigrator::applyDowngrade);
+  }
+
+  /**
+   * Upgrade a message from the source version to the most recent version by chaining all the required
+   * migrations
+   */
+  public <PreviousVersion, CurrentVersion> CurrentVersion upgrade(final PreviousVersion message, final Version source) {
+    return migrationContainer.upgrade(message, source, ConfiguredAirbyteCatalogMigrator::applyUpgrade);
+  }
+
+  public Version getMostRecentVersion() {
+    return migrationContainer.getMostRecentVersion();
+  }
+
+  // Helper function to work around type casting
+  private static <PreviousVersion, CurrentVersion> PreviousVersion applyDowngrade(final ConfiguredAirbyteCatalogMigration<PreviousVersion, CurrentVersion> migration,
+                                                                                  final Object message) {
+    return migration.downgrade((CurrentVersion) message);
+  }
+
+  // Helper function to work around type casting
+  private static <PreviousVersion, CurrentVersion> CurrentVersion applyUpgrade(final ConfiguredAirbyteCatalogMigration<PreviousVersion, CurrentVersion> migration,
+                                                                               final Object message) {
+    return migration.upgrade((PreviousVersion) message);
+  }
+
+  // Used for inspection of the injection
+  @VisibleForTesting
+  Set<String> getMigrationKeys() {
+    return migrationContainer.getMigrationKeys();
+  }
+
+}

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

@@ -0,0 +1,17 @@
+/*
+ * Copyright (c) 2022 Airbyte, Inc., all rights reserved.
+ */
+
+package io.airbyte.commons.protocol;
+
+import io.airbyte.commons.json.Jsons;
+import io.airbyte.protocol.models.ConfiguredAirbyteCatalog;
+
+public class DefaultProtocolSerializer implements ProtocolSerializer {
+
+  @Override
+  public String serialize(ConfiguredAirbyteCatalog configuredAirbyteCatalog) {
+    return Jsons.serialize(configuredAirbyteCatalog);
+  }
+
+}

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

@@ -0,0 +1,13 @@
+/*
+ * Copyright (c) 2022 Airbyte, Inc., all rights reserved.
+ */
+
+package io.airbyte.commons.protocol;
+
+import io.airbyte.protocol.models.ConfiguredAirbyteCatalog;
+
+public interface ProtocolSerializer {
+
+  String serialize(final ConfiguredAirbyteCatalog configuredAirbyteCatalog);
+
+}

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

@@ -0,0 +1,32 @@
+/*
+ * Copyright (c) 2022 Airbyte, Inc., all rights reserved.
+ */
+
+package io.airbyte.commons.protocol;
+
+import io.airbyte.commons.json.Jsons;
+import io.airbyte.commons.version.Version;
+import io.airbyte.protocol.models.ConfiguredAirbyteCatalog;
+
+/**
+ * Serialize a ConfiguredAirbyteCatalog to the specified version
+ * <p>
+ * This Serializer expects a ConfiguredAirbyteCatalog from the Current version of the platform,
+ * converts it to the target protocol version before serializing it.
+ */
+public class VersionedProtocolSerializer implements ProtocolSerializer {
+
+  private final ConfiguredAirbyteCatalogMigrator configuredAirbyteCatalogMigrator;
+  private final Version protocolVersion;
+
+  public VersionedProtocolSerializer(final ConfiguredAirbyteCatalogMigrator configuredAirbyteCatalogMigrator, final Version protocolVersion) {
+    this.configuredAirbyteCatalogMigrator = configuredAirbyteCatalogMigrator;
+    this.protocolVersion = protocolVersion;
+  }
+
+  @Override
+  public String serialize(final ConfiguredAirbyteCatalog configuredAirbyteCatalog) {
+    return Jsons.serialize(configuredAirbyteCatalogMigrator.downgrade(configuredAirbyteCatalog, protocolVersion));
+  }
+
+}