[controller][schema] Move migration schema normalization into

StoreSchemaService

Author: Ali Poursamadi

services/venice-controller/src/main/java/com/linkedin/venice/controller/ParentSchemaOrchestrator.java

@@ -90,8 +90,7 @@ SchemaEntry addValueSchema(
       DirectionalSchemaCompatibilityType expectedCompatibilityType) {
     parent.acquireAdminMessageLock(clusterName, storeName);
     try {
-      newValueSchemaStr = parent.getVeniceHelixAdmin()
-          .normalizeSchemaForMigration(clusterName, storeName, newValueSchemaStr);
+      newValueSchemaStr = storeSchemaService.normalizeSchemaForMigration(clusterName, storeName, newValueSchemaStr);
       final int newValueSchemaId = storeSchemaService.checkPreConditionForAddValueSchemaAndGetNewSchemaId(
           clusterName,
           storeName,
@@ -253,8 +252,7 @@ SchemaEntry addValueSchema(
       DirectionalSchemaCompatibilityType expectedCompatibilityType) {
     parent.acquireAdminMessageLock(clusterName, storeName);
     try {
-      newValueSchemaStr =
-          parent.getVeniceHelixAdmin().normalizeSchemaForMigration(clusterName, storeName, newValueSchemaStr);
+      newValueSchemaStr = storeSchemaService.normalizeSchemaForMigration(clusterName, storeName, newValueSchemaStr);
       Schema newValueSchema = AvroSchemaParseUtils.parseSchemaFromJSONStrictValidation(newValueSchemaStr);
 
       final Store store = parent.getVeniceHelixAdmin().getStore(clusterName, storeName);

services/venice-controller/src/main/java/com/linkedin/venice/controller/StoreSchemaService.java

@@ -8,8 +8,10 @@
 import com.linkedin.avroutil1.compatibility.RecordGenerationConfig;
 import com.linkedin.venice.exceptions.InvalidVeniceSchemaException;
 import com.linkedin.venice.exceptions.VeniceException;
+import com.linkedin.venice.helix.ZkStoreConfigAccessor;
 import com.linkedin.venice.meta.ReadWriteSchemaRepository;
 import com.linkedin.venice.meta.Store;
+import com.linkedin.venice.meta.StoreConfig;
 import com.linkedin.venice.meta.Version;
 import com.linkedin.venice.schema.AvroSchemaParseUtils;
 import com.linkedin.venice.schema.GeneratedSchemaID;
@@ -235,13 +237,67 @@ private void validateValueSchemaUsingRandomGenerator(String schemaStr, String cl
     }
   }
 
+  /**
+   * If {@code storeName} is migrating into {@code clusterName}, accept schemas that fail strict parse only because of
+   * {@code validateNumericDefaultValueTypes} (e.g. legacy {@code {"type":"float","default":0}}) by walking the JSON
+   * and coercing numeric defaults to the declared field type. The output is strict-parse-clean, which keeps downstream
+   * consumers that strict-parse (DaVinci's {@code SchemaUtils.annotateValueSchema}, VPJ, Samza producer) working.
+   *
+   * Re-strict-parses the coerced output as a defensive check so anything beyond the numeric-default tier (bad names,
+   * dangling content, union default not first branch) still fails loudly.
+   *
+   * For non-migration calls, and for migration calls whose input is already strict-clean, returns the input unchanged —
+   * so this can be wired into entry points idempotently.
+   *
+   * @return possibly-coerced schema string that is guaranteed to pass strict parsing.
+   */
+  String normalizeSchemaForMigration(String clusterName, String storeName, String schemaStr) {
+    ZkStoreConfigAccessor accessor = admin.getStoreConfigAccessor(clusterName);
+    if (!accessor.containsConfig(storeName)) {
+      return schemaStr;
+    }
+    StoreConfig cfg = accessor.getStoreConfig(storeName);
+    if (cfg == null || !clusterName.equals(cfg.getMigrationDestCluster())) {
+      return schemaStr;
+    }
+    // Migration context. If strict already passes, leave the string unchanged so we don't
+    // introduce gratuitous diffs against the source schema.
+    try {
+      AvroSchemaParseUtils.parseSchemaFromJSONStrictValidation(schemaStr);
+      return schemaStr;
+    } catch (Exception strictFailure) {
+      LOGGER.info(
+          "Strict parse failed for store {} migrating into cluster {}; attempting numeric-default coercion.",
+          storeName,
+          clusterName,
+          strictFailure);
+      String coerced = AvroSchemaParseUtils.coerceNumericDefaultsToFieldType(schemaStr);
+      // Defensive: anything LOOSE_NUMERICS would have been lenient about (union default not first
+      // branch, bad names, etc.) is outside the coercion scope and must still fail strict. When it
+      // does, surface the *original* strict failure too — it's the one the operator needs to see.
+      try {
+        AvroSchemaParseUtils.parseSchemaFromJSONStrictValidation(coerced);
+      } catch (Exception coercedFailure) {
+        coercedFailure.addSuppressed(strictFailure);
+        throw coercedFailure;
+      }
+      if (!coerced.equals(schemaStr)) {
+        LOGGER.info(
+            "Coerced numeric default(s) in value schema for store {} migrating into cluster {}.",
+            storeName,
+            clusterName);
+      }
+      return coerced;
+    }
+  }
+
   SchemaEntry addValueSchema(
       String clusterName,
       String storeName,
       String valueSchemaStr,
       DirectionalSchemaCompatibilityType expectedCompatibilityType) {
     admin.checkControllerLeadershipFor(clusterName);
-    valueSchemaStr = admin.normalizeSchemaForMigration(clusterName, storeName, valueSchemaStr);
+    valueSchemaStr = normalizeSchemaForMigration(clusterName, storeName, valueSchemaStr);
     ReadWriteSchemaRepository schemaRepository =
         admin.getHelixVeniceClusterResources(clusterName).getSchemaRepository();
     SchemaEntry schemaEntry = schemaRepository.addValueSchema(storeName, valueSchemaStr, expectedCompatibilityType);
@@ -260,7 +316,7 @@ SchemaEntry addValueSchema(
       int schemaId,
       DirectionalSchemaCompatibilityType compatibilityType) {
     admin.checkControllerLeadershipFor(clusterName);
-    valueSchemaStr = admin.normalizeSchemaForMigration(clusterName, storeName, valueSchemaStr);
+    valueSchemaStr = normalizeSchemaForMigration(clusterName, storeName, valueSchemaStr);
     ReadWriteSchemaRepository schemaRepository =
         admin.getHelixVeniceClusterResources(clusterName).getSchemaRepository();
     int newValueSchemaId =
@@ -318,8 +374,8 @@ SchemaEntry addSupersetSchema(
       String supersetSchemaStr,
       int supersetSchemaId) {
     admin.checkControllerLeadershipFor(clusterName);
-    valueSchema = admin.normalizeSchemaForMigration(clusterName, storeName, valueSchema);
-    supersetSchemaStr = admin.normalizeSchemaForMigration(clusterName, storeName, supersetSchemaStr);
+    valueSchema = normalizeSchemaForMigration(clusterName, storeName, valueSchema);
+    supersetSchemaStr = normalizeSchemaForMigration(clusterName, storeName, supersetSchemaStr);
     ReadWriteSchemaRepository schemaRepository =
         admin.getHelixVeniceClusterResources(clusterName).getSchemaRepository();
 
@@ -364,7 +420,7 @@ int checkPreConditionForAddValueSchemaAndGetNewSchemaId(
       String storeName,
       String valueSchemaStr,
       DirectionalSchemaCompatibilityType expectedCompatibilityType) {
-    valueSchemaStr = admin.normalizeSchemaForMigration(clusterName, storeName, valueSchemaStr);
+    valueSchemaStr = normalizeSchemaForMigration(clusterName, storeName, valueSchemaStr);
     AvroSchemaUtils.validateAvroSchemaStr(valueSchemaStr);
     AvroSchemaParseUtils.parseSchemaFromJSONStrictValidation(valueSchemaStr);
     validateValueSchemaUsingRandomGenerator(valueSchemaStr, clusterName, storeName);

services/venice-controller/src/main/java/com/linkedin/venice/controller/VeniceHelixAdmin.java

@@ -1176,7 +1176,7 @@ public void createStore(
     HelixVeniceClusterResources clusterResources = getHelixVeniceClusterResources(clusterName);
     LOGGER.info("Start creating store {} in cluster {} with owner {}", storeName, clusterName, owner);
     try (AutoCloseableLock ignore = clusterResources.getClusterLockManager().createStoreWriteLock(storeName)) {
-      valueSchema = normalizeSchemaForMigration(clusterName, storeName, valueSchema);
+      valueSchema = storeSchemaService.normalizeSchemaForMigration(clusterName, storeName, valueSchema);
       checkPreConditionForCreateStore(clusterName, storeName, keySchema, valueSchema, isSystemStore, true);
       VeniceControllerClusterConfig config = getHelixVeniceClusterResources(clusterName).getConfig();
       Store newStore = new ZKStore(
@@ -2257,61 +2257,6 @@ protected void checkPreConditionForCreateStore(
     new SchemaEntry(SchemaData.INVALID_VALUE_SCHEMA_ID, valueSchema);
   }
 
-  /**
-   * If {@code storeName} is migrating into {@code clusterName}, accept schemas that fail strict
-   * parse only because of {@code validateNumericDefaultValueTypes} (e.g. legacy
-   * {@code {"type":"float","default":0}}) by walking the JSON and coercing numeric defaults to the
-   * declared field type. The output is strict-parse-clean, which keeps downstream consumers that
-   * strict-parse (DaVinci's {@code SchemaUtils.annotateValueSchema}, VPJ, Samza producer) working.
-   *
-   * Re-strict-parses the coerced output as a defensive check so anything beyond the numeric-default
-   * tier (bad names, dangling content, union default not first branch) still fails loudly.
-   *
-   * For non-migration calls, and for migration calls whose input is already strict-clean, returns
-   * the input unchanged — so this can be wired into entry points idempotently.
-   *
-   * @return possibly-coerced schema string that is guaranteed to pass strict parsing.
-   */
-  String normalizeSchemaForMigration(String clusterName, String storeName, String schemaStr) {
-    ZkStoreConfigAccessor accessor = getStoreConfigAccessor(clusterName);
-    if (!accessor.containsConfig(storeName)) {
-      return schemaStr;
-    }
-    StoreConfig cfg = accessor.getStoreConfig(storeName);
-    if (cfg == null || !clusterName.equals(cfg.getMigrationDestCluster())) {
-      return schemaStr;
-    }
-    // Migration context. If strict already passes, leave the string unchanged so we don't
-    // introduce gratuitous diffs against the source schema.
-    try {
-      AvroSchemaParseUtils.parseSchemaFromJSONStrictValidation(schemaStr);
-      return schemaStr;
-    } catch (Exception strictFailure) {
-      LOGGER.info(
-          "Strict parse failed for store {} migrating into cluster {}; attempting numeric-default coercion.",
-          storeName,
-          clusterName,
-          strictFailure);
-      String coerced = AvroSchemaParseUtils.coerceNumericDefaultsToFieldType(schemaStr);
-      // Defensive: anything LOOSE_NUMERICS would have been lenient about (union default not first
-      // branch, bad names, etc.) is outside the coercion scope and must still fail strict. When it
-      // does, surface the *original* strict failure too — it's the one the operator needs to see.
-      try {
-        AvroSchemaParseUtils.parseSchemaFromJSONStrictValidation(coerced);
-      } catch (Exception coercedFailure) {
-        coercedFailure.addSuppressed(strictFailure);
-        throw coercedFailure;
-      }
-      if (!coerced.equals(schemaStr)) {
-        LOGGER.info(
-            "Coerced numeric default(s) in value schema for store {} migrating into cluster {}.",
-            storeName,
-            clusterName);
-      }
-      return coerced;
-    }
-  }
-
   void checkStoreGraveyardForRecreation(String clusterName, String storeName) {
     HelixVeniceClusterResources clusterResources = getHelixVeniceClusterResources(clusterName);
 

services/venice-controller/src/main/java/com/linkedin/venice/controller/VeniceParentHelixAdmin.java

@@ -916,7 +916,8 @@ public void createStore(
       Optional<String> accessPermissions) {
     acquireAdminMessageLock(clusterName, storeName);
     try {
-      valueSchema = getVeniceHelixAdmin().normalizeSchemaForMigration(clusterName, storeName, valueSchema);
+      valueSchema = getVeniceHelixAdmin().getStoreSchemaService()
+          .normalizeSchemaForMigration(clusterName, storeName, valueSchema);
       getVeniceHelixAdmin()
           .checkPreConditionForCreateStore(clusterName, storeName, keySchema, valueSchema, isSystemStore, false);
       LOGGER.info("Adding store: {} to cluster: {}", storeName, clusterName);

services/venice-controller/src/test/java/com/linkedin/venice/controller/AbstractTestVeniceParentHelixAdmin.java

@@ -102,9 +102,9 @@ public void setupInternalMocks() {
     SchemaEntry mockEntry = new SchemaEntry(0, TEST_SCHEMA);
     doReturn(mockEntry).when(internalAdmin).getKeySchema(anyString(), anyString());
     // Outside a store-migration context, schema normalization is a passthrough. Mirror that here so
-    // the mocked internal admin doesn't return null and blank out the value schema during createStore/
+    // the mocked schema service doesn't return null and blank out the value schema during createStore/
     // addValueSchema.
-    when(internalAdmin.normalizeSchemaForMigration(anyString(), anyString(), any()))
+    when(storeSchemaService.normalizeSchemaForMigration(anyString(), anyString(), any()))
         .thenAnswer(invocation -> invocation.getArgument(2));
 
     zkClient = mock(ZkClient.class);

services/venice-controller/src/test/java/com/linkedin/venice/controller/TestStoreSchemaService.java

@@ -4,7 +4,6 @@
 import static org.mockito.ArgumentMatchers.anyInt;
 import static org.mockito.ArgumentMatchers.anyString;
 import static org.mockito.ArgumentMatchers.eq;
-import static org.mockito.Mockito.doAnswer;
 import static org.mockito.Mockito.doReturn;
 import static org.mockito.Mockito.doThrow;
 import static org.mockito.Mockito.mock;
@@ -17,6 +16,7 @@
 import static org.testng.Assert.assertTrue;
 
 import com.linkedin.venice.exceptions.VeniceException;
+import com.linkedin.venice.helix.ZkStoreConfigAccessor;
 import com.linkedin.venice.meta.ReadWriteSchemaRepository;
 import com.linkedin.venice.meta.Store;
 import com.linkedin.venice.meta.Version;
@@ -55,6 +55,7 @@ public class TestStoreSchemaService {
   private VeniceHelixAdmin admin;
   private HelixVeniceClusterResources resources;
   private ReadWriteSchemaRepository schemaRepo;
+  private ZkStoreConfigAccessor storeConfigAccessor;
   private VeniceControllerClusterConfig config;
   private StoreSchemaService service;
 
@@ -63,10 +64,11 @@ public void setUp() {
     admin = mock(VeniceHelixAdmin.class);
     resources = mock(HelixVeniceClusterResources.class);
     schemaRepo = mock(ReadWriteSchemaRepository.class);
+    storeConfigAccessor = mock(ZkStoreConfigAccessor.class);
     config = mock(VeniceControllerClusterConfig.class);
     doReturn(resources).when(admin).getHelixVeniceClusterResources(CLUSTER);
-    doAnswer(invocation -> invocation.getArgument(2)).when(admin)
-        .normalizeSchemaForMigration(anyString(), anyString(), anyString());
+    doReturn(storeConfigAccessor).when(admin).getStoreConfigAccessor(anyString());
+    doReturn(false).when(storeConfigAccessor).containsConfig(anyString());
     doReturn(schemaRepo).when(resources).getSchemaRepository();
     doReturn(config).when(resources).getConfig();
     service = new StoreSchemaService(admin);

services/venice-controller/src/test/java/com/linkedin/venice/controller/TestVeniceHelixAdminWithoutCluster.java

@@ -477,7 +477,7 @@ public void testWriteEndOfPushThrowsForNonExistentStore() throws Exception {
     admin.writeEndOfPush("cluster", "missing_store", 1, false, null);
   }
 
-  // ---- normalizeSchemaForMigration ---------------------------------------------------------
+  // ---- StoreSchemaService.normalizeSchemaForMigration --------------------------------------
 
   private static final String CLEAN_VALUE_SCHEMA =
       "{\"type\":\"record\",\"name\":\"Clean\",\"fields\":[" + "{\"name\":\"v\",\"type\":\"string\"}]}";
@@ -491,23 +491,22 @@ public void testWriteEndOfPushThrowsForNonExistentStore() throws Exception {
   private static final String NON_NUMERIC_STRICT_VIOLATION_SCHEMA = "{\"type\":\"record\",\"name\":\"BadUnion\","
       + "\"fields\":[{\"name\":\"f\",\"type\":[\"int\",\"null\"],\"default\":null}]}";
 
-  private static VeniceHelixAdmin newNormalizeMock(String clusterName, String storeName, StoreConfig storeConfig) {
+  private static StoreSchemaService newNormalizeService(String clusterName, String storeName, StoreConfig storeConfig) {
     VeniceHelixAdmin admin = mock(VeniceHelixAdmin.class);
     ZkStoreConfigAccessor accessor = mock(ZkStoreConfigAccessor.class);
     doReturn(storeConfig != null).when(accessor).containsConfig(storeName);
     doReturn(storeConfig).when(accessor).getStoreConfig(storeName);
     doReturn(accessor).when(admin).getStoreConfigAccessor(clusterName);
-    doCallRealMethod().when(admin).normalizeSchemaForMigration(anyString(), anyString(), anyString());
-    return admin;
+    return new StoreSchemaService(admin);
   }
 
   @Test
   public void testNormalizeReturnsInputWhenStoreConfigAbsent() {
     String cluster = "venice-dest";
     String store = "legacy_store";
-    VeniceHelixAdmin admin = newNormalizeMock(cluster, store, null);
+    StoreSchemaService service = newNormalizeService(cluster, store, null);
     Assert.assertSame(
-        admin.normalizeSchemaForMigration(cluster, store, LEGACY_NUMERIC_DEFAULT_SCHEMA),
+        service.normalizeSchemaForMigration(cluster, store, LEGACY_NUMERIC_DEFAULT_SCHEMA),
         LEGACY_NUMERIC_DEFAULT_SCHEMA,
         "Non-migration context (no storeConfig) must return input by identity to avoid parse cost");
   }
@@ -519,9 +518,9 @@ public void testNormalizeReturnsInputWhenMigrationDestIsAnotherCluster() {
     StoreConfig cfg = mock(StoreConfig.class);
     doReturn("some_other_dest").when(cfg).getMigrationDestCluster();
 
-    VeniceHelixAdmin admin = newNormalizeMock(cluster, store, cfg);
+    StoreSchemaService service = newNormalizeService(cluster, store, cfg);
     Assert.assertSame(
-        admin.normalizeSchemaForMigration(cluster, store, LEGACY_NUMERIC_DEFAULT_SCHEMA),
+        service.normalizeSchemaForMigration(cluster, store, LEGACY_NUMERIC_DEFAULT_SCHEMA),
         LEGACY_NUMERIC_DEFAULT_SCHEMA,
         "When migrationDestCluster does not match this cluster, input must pass through unchanged");
   }
@@ -533,9 +532,9 @@ public void testNormalizeReturnsInputWhenStrictAlreadyPasses() {
     StoreConfig cfg = mock(StoreConfig.class);
     doReturn(cluster).when(cfg).getMigrationDestCluster();
 
-    VeniceHelixAdmin admin = newNormalizeMock(cluster, store, cfg);
+    StoreSchemaService service = newNormalizeService(cluster, store, cfg);
     Assert.assertSame(
-        admin.normalizeSchemaForMigration(cluster, store, CLEAN_VALUE_SCHEMA),
+        service.normalizeSchemaForMigration(cluster, store, CLEAN_VALUE_SCHEMA),
         CLEAN_VALUE_SCHEMA,
         "Strict-clean input under migration context must not be reserialized");
   }
@@ -547,8 +546,8 @@ public void testNormalizeReserializesLegacyNumericDefault() {
     StoreConfig cfg = mock(StoreConfig.class);
     doReturn(cluster).when(cfg).getMigrationDestCluster();
 
-    VeniceHelixAdmin admin = newNormalizeMock(cluster, store, cfg);
-    String normalized = admin.normalizeSchemaForMigration(cluster, store, LEGACY_NUMERIC_DEFAULT_SCHEMA);
+    StoreSchemaService service = newNormalizeService(cluster, store, cfg);
+    String normalized = service.normalizeSchemaForMigration(cluster, store, LEGACY_NUMERIC_DEFAULT_SCHEMA);
 
     Assert.assertNotEquals(normalized, LEGACY_NUMERIC_DEFAULT_SCHEMA, "Legacy schema must be reserialized");
     // The whole point: output must be strict-parse-clean so downstream consumers don't trip.
@@ -562,12 +561,12 @@ public void testNormalizeRejectsNonNumericStrictViolation() {
     StoreConfig cfg = mock(StoreConfig.class);
     doReturn(cluster).when(cfg).getMigrationDestCluster();
 
-    VeniceHelixAdmin admin = newNormalizeMock(cluster, store, cfg);
+    StoreSchemaService service = newNormalizeService(cluster, store, cfg);
     // Migration context, but the violation is outside the numeric-default tier — must propagate.
     // The original strict failure must ride along as a suppressed exception so the operator can
     // see what was actually wrong with the source schema, not just that post-coercion strict tripped.
     try {
-      admin.normalizeSchemaForMigration(cluster, store, NON_NUMERIC_STRICT_VIOLATION_SCHEMA);
+      service.normalizeSchemaForMigration(cluster, store, NON_NUMERIC_STRICT_VIOLATION_SCHEMA);
       Assert.fail("Expected normalize to throw on non-numeric strict violation");
     } catch (Exception coercedFailure) {
       Assert.assertTrue(