fix(migration): introduce ServiceLoader SPI for extension migration providers (#27760)

Author: manerow

openmetadata-service/src/main/java/org/openmetadata/service/migration/api/MigrationProcess.java

@@ -40,12 +40,14 @@
  * Even for pure Java migrations, the directory structure MUST exist - SQL files can be empty but the
  * directory structure is mandatory.
  *
- * <p>Java migrations must follow this naming convention:
+ * <p>Native OpenMetadata Java migrations must follow this naming convention:
  * {@code org.openmetadata.service.migration.[dbPackageName].[versionPackage].Migration}
  * Example: {@code org.openmetadata.service.migration.postgres.v120.Migration}
  *
- * In collate:
- * {@code io.collate.service.migration.[dbPackageName].[versionPackage].Migration}
+ * <p>Migrations that ship outside of OpenMetadata (extension migration directories) are resolved
+ * by implementations of {@link MigrationProcessExtensionProvider} registered via
+ * {@code java.util.ServiceLoader}. When no provider handles a given extension version, the
+ * workflow falls back to {@link MigrationProcessImpl} (SQL changes only, no Java data migration).
  *
  * <p>Java migrations should extend {@code MigrationProcessImpl} and override required methods,
  * particularly {@code runDataMigration()} and {@code getMigrationOps()}.

openmetadata-service/src/main/java/org/openmetadata/service/migration/api/MigrationProcessExtensionProvider.java

@@ -0,0 +1,41 @@
+/*
+ *  Copyright 2021 Collate
+ *  Licensed under the Apache License, Version 2.0 (the "License");
+ *  you may not use this file except in compliance with the License.
+ *  You may obtain a copy of the License at
+ *  http://www.apache.org/licenses/LICENSE-2.0
+ *  Unless required by applicable law or agreed to in writing, software
+ *  distributed under the License is distributed on an "AS IS" BASIS,
+ *  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ *  See the License for the specific language governing permissions and
+ *  limitations under the License.
+ */
+
+package org.openmetadata.service.migration.api;
+
+import java.util.Optional;
+import org.openmetadata.service.migration.utils.MigrationFile;
+
+/**
+ * SPI for resolving a {@link MigrationProcess} for migration directories that ship outside of
+ * OpenMetadata. Commercial distributions or downstream forks register implementations via
+ * {@code META-INF/services/org.openmetadata.service.migration.api.MigrationProcessExtensionProvider}
+ * to plug in their own Java migration classes without OpenMetadata having to know about them.
+ *
+ * <p>The migration workflow only consults providers for files where {@code MigrationFile.isExtension}
+ * is true. If no provider returns a present value, the workflow falls back to {@link
+ * MigrationProcessImpl}, which runs the SQL changes from the version's directory and performs no
+ * Java-level data migration.
+ */
+public interface MigrationProcessExtensionProvider {
+
+  /**
+   * Resolve a migration process for the given extension migration file.
+   *
+   * @param file the extension migration file (guaranteed {@code file.isExtension == true})
+   * @return the {@link MigrationProcess} to run for this version, or {@link Optional#empty()} if
+   *     this provider does not handle the version (the workflow will try the next provider, or
+   *     fall back to {@link MigrationProcessImpl}).
+   */
+  Optional<MigrationProcess> provide(MigrationFile file);
+}

openmetadata-service/src/main/java/org/openmetadata/service/migration/api/MigrationWorkflow.java

@@ -14,11 +14,13 @@
 import java.util.Map;
 import java.util.Objects;
 import java.util.Optional;
+import java.util.ServiceLoader;
 import java.util.Set;
 import java.util.UUID;
 import java.util.regex.Matcher;
 import java.util.regex.Pattern;
 import java.util.stream.Stream;
+import java.util.stream.StreamSupport;
 import lombok.Getter;
 import lombok.extern.slf4j.Slf4j;
 import org.flywaydb.core.api.configuration.ClassicConfiguration;
@@ -162,6 +164,7 @@ public List<MigrationFile> getMigrationFilesFromPath(
   private List<MigrationProcess> filterAndGetMigrationsToRun(
       List<MigrationFile> availableMigrations) {
     List<MigrationFile> applyMigrations = resolveApplyMigrations(availableMigrations);
+    List<MigrationProcessExtensionProvider> extensionProviders = loadExtensionProviders();
     List<MigrationProcess> processes = new ArrayList<>();
     try {
       for (MigrationFile file : applyMigrations) {
@@ -172,29 +175,37 @@ private List<MigrationProcess> filterAndGetMigrationsToRun(
               file.version);
           continue;
         }
-        String extClazzName = null;
-        if (file.version.contains("collate")) {
-          extClazzName = file.getMigrationProcessExtClassName();
-        }
-        if (extClazzName != null) {
-          MigrationProcess collateProcess =
-              (MigrationProcess)
-                  Class.forName(extClazzName).getConstructor(MigrationFile.class).newInstance(file);
-          processes.add(collateProcess);
-        } else {
-          String clazzName = file.getMigrationProcessClassName();
-          MigrationProcess openMetadataProcess =
-              (MigrationProcess)
-                  Class.forName(clazzName).getConstructor(MigrationFile.class).newInstance(file);
-          processes.add(openMetadataProcess);
-        }
+        processes.add(resolveMigrationProcess(file, extensionProviders));
       }
     } catch (Exception e) {
       LOG.error("Failed to list and add migrations to run due to ", e);
     }
     return processes;
   }
 
+  private MigrationProcess resolveMigrationProcess(
+      MigrationFile file, List<MigrationProcessExtensionProvider> extensionProviders)
+      throws ReflectiveOperationException {
+    if (file.isExtension) {
+      // No provider handled this extension version: run SQL only, skip Java data migration.
+      // Critical: do not fall through to OM's same-version native migration class.
+      return extensionProviders.stream()
+          .map(provider -> provider.provide(file))
+          .flatMap(Optional::stream)
+          .findFirst()
+          .orElseGet(() -> new MigrationProcessImpl(file));
+    }
+    String clazzName = file.getMigrationProcessClassName();
+    return (MigrationProcess)
+        Class.forName(clazzName).getConstructor(MigrationFile.class).newInstance(file);
+  }
+
+  private List<MigrationProcessExtensionProvider> loadExtensionProviders() {
+    return StreamSupport.stream(
+            ServiceLoader.load(MigrationProcessExtensionProvider.class).spliterator(), false)
+        .toList();
+  }
+
   private static int compareVersions(String version1, String version2) {
     int[] v1Parts = parseVersion(version1);
     int[] v2Parts = parseVersion(version2);

openmetadata-service/src/main/java/org/openmetadata/service/migration/utils/MigrationFile.java

@@ -119,16 +119,12 @@ public String getMigrationProcessClassName() {
     return clazzName;
   }
 
-  public String getMigrationProcessExtClassName() {
-    String clazzName =
-        String.format(
-            "io.collate.service.migration.%s.%s.Migration", dbPackageName, getVersionPackageName());
-    try {
-      Class.forName(clazzName);
-    } catch (ClassNotFoundException e) {
-      return null;
+  public String getVersionPackageName() {
+    StringBuilder arrayAsString = new StringBuilder();
+    for (int versionNumber : versionNumbers) {
+      arrayAsString.append(versionNumber);
     }
-    return clazzName;
+    return "v" + arrayAsString;
   }
 
   public String getMigrationsFilePath() {
@@ -189,14 +185,6 @@ private int compareVersionNumbers(int[] another) {
     return 0;
   }
 
-  private String getVersionPackageName() {
-    StringBuilder arrayAsString = new StringBuilder();
-    for (int versionNumber : versionNumbers) {
-      arrayAsString.append(versionNumber);
-    }
-    return "v" + arrayAsString;
-  }
-
   public boolean isReprocessing() {
     return reprocessing;
   }

openmetadata-service/src/test/java/org/openmetadata/service/migration/api/MigrationWorkflowTest.java

@@ -218,6 +218,35 @@ void loadMigrationsReprocessesCurrentExtensionDespiteHigherExecutedCoreVersions(
     assertEquals(List.of("0.0.1", "0.0.1-collate"), getMigrationVersions(workflow));
   }
 
+  @Test
+  void loadMigrationsFallsBackToNoOpWhenNoExtensionProviderHandlesVersion() throws Exception {
+    // Regression: a Collate version directory without a registered Java provider must resolve to
+    // MigrationProcessImpl (no-op data migration), NOT to OM's same-numeric-version migration
+    // class. Otherwise OM data migrations like migrateThreadTasksToTaskEntity would be re-run for
+    // every Collate version that shares its major.minor.patch.
+    Path nativeRoot = Files.createDirectories(tempDir.resolve("native"));
+    Path extensionRoot = Files.createDirectories(tempDir.resolve("extension"));
+    createMigrationDir(extensionRoot, "1.12.1-collate", "SELECT 22;");
+    when(migrationDAO.getMigrationVersions()).thenReturn(List.of());
+
+    MigrationWorkflow workflow =
+        new MigrationWorkflow(
+            jdbi,
+            nativeRoot.toString(),
+            ConnectionType.POSTGRES,
+            extensionRoot.toString(),
+            null,
+            config,
+            false);
+
+    workflow.loadMigrations();
+
+    List<MigrationProcess> resolved = getMigrations(workflow);
+    assertEquals(1, resolved.size());
+    assertEquals("1.12.1-collate", resolved.get(0).getVersion());
+    assertEquals(MigrationProcessImpl.class, resolved.get(0).getClass());
+  }
+
   @Test
   void loadMigrationsFallsBackToRunningEverythingWhenMigrationLookupFails() throws Exception {
     Path nativeRoot = Files.createDirectories(tempDir.resolve("native"));
@@ -724,10 +753,14 @@ private void createMigrationDir(Path root, String version, String sql) throws Ex
 
   @SuppressWarnings("unchecked")
   private List<String> getMigrationVersions(MigrationWorkflow workflow) throws Exception {
+    return getMigrations(workflow).stream().map(MigrationProcess::getVersion).toList();
+  }
+
+  @SuppressWarnings("unchecked")
+  private List<MigrationProcess> getMigrations(MigrationWorkflow workflow) throws Exception {
     Field field = MigrationWorkflow.class.getDeclaredField("migrations");
     field.setAccessible(true);
-    List<MigrationProcess> migrations = (List<MigrationProcess>) field.get(workflow);
-    return migrations.stream().map(MigrationProcess::getVersion).toList();
+    return (List<MigrationProcess>) field.get(workflow);
   }
 
   @SuppressWarnings("unchecked")