[controller][admin-tool][compat] Add pub-sub encryption key URN to stores (#2962)

Author: minhmo1620

build.gradle

@@ -328,10 +328,7 @@ subprojects {
       // Protocol changes should pin the current version via this override and remove the override in a follow-up PR
       // when actually using the new protocol. Example to pin KME to v12 when introducing v13:
       // project(':internal:venice-common').file('src/main/resources/avro/KafkaMessageEnvelope/v12', PathValidation.DIRECTORY)
-      def versionOverrides = [
-          project(':internal:venice-common').file('src/main/resources/avro/StoreMetaValue/v46', PathValidation.DIRECTORY),
-          project(':services:venice-controller').file('src/main/resources/avro/AdminOperation/v101', PathValidation.DIRECTORY),
-      ]
+      def versionOverrides = []
 
       def schemaDirs = [sourceDir]
       sourceDir.eachDir { typeDir ->

clients/venice-admin-tool/src/main/java/com/linkedin/venice/AdminTool.java

@@ -1431,6 +1431,7 @@ static UpdateStoreQueryParams getUpdateStoreQueryParams(CommandLine cmd) {
     integerParam(cmd, Arg.LATEST_SUPERSET_SCHEMA_ID, p -> params.setLatestSupersetSchemaId(p), argSet);
     booleanParam(cmd, Arg.ENABLE_COMPACTION, p -> params.setCompactionEnabled(p), argSet);
     longParam(cmd, Arg.COMPACTION_THRESHOLD_MILLISECONDS, p -> params.setCompactionThresholdMilliseconds(p), argSet);
+    genericParam(cmd, Arg.PUB_SUB_ENCRYPTION_KEY_URN, s -> s, params::setPubSubEncryptionKeyUrn, argSet);
     longParam(cmd, Arg.MIN_COMPACTION_LAG_SECONDS, p -> params.setMinCompactionLagSeconds(p), argSet);
     longParam(cmd, Arg.MAX_COMPACTION_LAG_SECONDS, p -> params.setMaxCompactionLagSeconds(p), argSet);
     integerParam(cmd, Arg.MAX_RECORD_SIZE_BYTES, params::setMaxRecordSizeBytes, argSet);

clients/venice-admin-tool/src/main/java/com/linkedin/venice/Arg.java

@@ -303,6 +303,9 @@ public enum Arg {
   COMPACTION_THRESHOLD_MILLISECONDS(
       "compaction-threshold-milliseconds", "ctms", true, "Set compaction threshold in milliseconds"
   ),
+  PUB_SUB_ENCRYPTION_KEY_URN(
+      "pub-sub-encryption-key-urn", "psekurn", true, "Set the PubSub encryption key URN for an encryption-enabled store"
+  ),
   MIN_COMPACTION_LAG_SECONDS(
       "min-compaction-lag-seconds", "mcls", true, "Min compaction lag seconds for version topic of hybrid stores"
   ),

clients/venice-admin-tool/src/main/java/com/linkedin/venice/Command.java

@@ -118,6 +118,7 @@
 import static com.linkedin.venice.Arg.POSITION;
 import static com.linkedin.venice.Arg.PRINCIPAL;
 import static com.linkedin.venice.Arg.PROGRESS_INTERVAL;
+import static com.linkedin.venice.Arg.PUB_SUB_ENCRYPTION_KEY_URN;
 import static com.linkedin.venice.Arg.PUSH_ID;
 import static com.linkedin.venice.Arg.PUSH_STREAM_SOURCE_ADDRESS;
 import static com.linkedin.venice.Arg.READABILITY;
@@ -334,9 +335,9 @@ public enum Command {
           PUSH_STREAM_SOURCE_ADDRESS, BACKUP_VERSION_RETENTION_DAY, REPLICATION_FACTOR,
           NATIVE_REPLICATION_SOURCE_FABRIC, REPLICATE_ALL_CONFIGS, ACTIVE_ACTIVE_REPLICATION_ENABLED, REGIONS_FILTER,
           DISABLE_META_STORE, DISABLE_DAVINCI_PUSH_STATUS_STORE, STORAGE_PERSONA, STORE_VIEW_CONFIGS,
-          LATEST_SUPERSET_SCHEMA_ID, ENABLE_COMPACTION, COMPACTION_THRESHOLD_MILLISECONDS, MIN_COMPACTION_LAG_SECONDS,
-          MAX_COMPACTION_LAG_SECONDS, MAX_RECORD_SIZE_BYTES, MAX_NEARLINE_RECORD_SIZE_BYTES, THROUGHPUT_QUOTA_IN_BYTES,
-          THROUGHPUT_QUOTA_IN_RECORDS, UNUSED_SCHEMA_DELETION_ENABLED, BLOB_TRANSFER_ENABLED,
+          LATEST_SUPERSET_SCHEMA_ID, ENABLE_COMPACTION, COMPACTION_THRESHOLD_MILLISECONDS, PUB_SUB_ENCRYPTION_KEY_URN,
+          MIN_COMPACTION_LAG_SECONDS, MAX_COMPACTION_LAG_SECONDS, MAX_RECORD_SIZE_BYTES, MAX_NEARLINE_RECORD_SIZE_BYTES,
+          THROUGHPUT_QUOTA_IN_BYTES, THROUGHPUT_QUOTA_IN_RECORDS, UNUSED_SCHEMA_DELETION_ENABLED, BLOB_TRANSFER_ENABLED,
           SEPARATE_REALTIME_TOPIC_ENABLED, NEARLINE_PRODUCER_COMPRESSION_ENABLED, NEARLINE_PRODUCER_COUNT_PER_WRITER,
           TARGET_SWAP_REGION, TARGET_SWAP_REGION_WAIT_TIME, DAVINCI_HEARTBEAT_REPORTED, ENABLE_STORE_MIGRATION,
           GLOBAL_RT_DIV_ENABLED, ENUM_SCHEMA_EVOLUTION_ALLOWED, STORE_LIFECYCLE_HOOKS_LIST,

clients/venice-admin-tool/src/test/java/com/linkedin/venice/TestAdminTool.java

@@ -157,6 +157,17 @@ public void testAdminUpdateStoreRejectsEncryptionArgument() {
     expectThrows(ParseException.class, () -> AdminTool.getCommandLine(args));
   }
 
+  @Test
+  public void testAdminUpdateStoreAcceptsPubSubEncryptionKeyUrn() throws ParseException, IOException {
+    String pubSubEncryptionKeyUrn = "keyUrn:abc";
+    String[] args = { "--update-store", "--url", "http://localhost:7036", "--cluster", "test-cluster", "--store",
+        "testStore", "--pub-sub-encryption-key-urn", pubSubEncryptionKeyUrn };
+
+    UpdateStoreQueryParams params = AdminTool.getUpdateStoreQueryParams(AdminTool.getCommandLine(args));
+
+    assertEquals(params.getPubSubEncryptionKeyUrn(), Optional.of(pubSubEncryptionKeyUrn));
+  }
+
   @Test
   public void testAdminUpdateStoreArgThroughputQuota() throws ParseException, IOException {
     String[] args = { "--update-store", "--url", "http://localhost:7036", "--cluster", "test-cluster", "--store",

internal/venice-common/src/main/java/com/linkedin/venice/controllerapi/ControllerApiConstants.java

@@ -269,6 +269,7 @@ public class ControllerApiConstants {
 
   public static final String COMPACTION_ENABLED = "compaction_enabled";
   public static final String COMPACTION_THRESHOLD_MILLISECONDS = "compaction_threshold_milliseconds";
+  public static final String PUB_SUB_ENCRYPTION_KEY_URN = "pub_sub_encryption_key_urn";
   public static final String MIN_COMPACTION_LAG_SECONDS = "min_compaction_lag_seconds";
   public static final String MAX_COMPACTION_LAG_SECONDS = "max_compaction_lag_seconds";
 

internal/venice-common/src/main/java/com/linkedin/venice/controllerapi/UpdateStoreQueryParams.java

@@ -58,6 +58,7 @@
 import static com.linkedin.venice.controllerapi.ControllerApiConstants.PARTITION_COUNT;
 import static com.linkedin.venice.controllerapi.ControllerApiConstants.PERSONA_NAME;
 import static com.linkedin.venice.controllerapi.ControllerApiConstants.PREVIOUS_CURRENT_VERSION;
+import static com.linkedin.venice.controllerapi.ControllerApiConstants.PUB_SUB_ENCRYPTION_KEY_URN;
 import static com.linkedin.venice.controllerapi.ControllerApiConstants.PUSH_STREAM_SOURCE_ADDRESS;
 import static com.linkedin.venice.controllerapi.ControllerApiConstants.READ_COMPUTATION_ENABLED;
 import static com.linkedin.venice.controllerapi.ControllerApiConstants.READ_QUOTA_IN_CU;
@@ -140,6 +141,7 @@ public UpdateStoreQueryParams() {
    * @param srcStore The original store
    */
   public UpdateStoreQueryParams(StoreInfo srcStore, boolean storeMigrating) {
+    String pubSubEncryptionKeyUrn = srcStore.getPubSubEncryptionKeyUrn();
     // Copy everything except for currentVersion, daVinciPushStatusStoreEnabled, latestSuperSetValueSchemaId,
     // storeMetaSystemStoreEnabled, storeMetadataSystemStoreEnabled
     UpdateStoreQueryParams updateStoreQueryParams =
@@ -201,6 +203,10 @@ public UpdateStoreQueryParams(StoreInfo srcStore, boolean storeMigrating) {
                     .collect(Collectors.toMap(Map.Entry::getKey, e -> e.getValue().toString())))
             .setFlinkVeniceViewsEnabled(srcStore.isFlinkVeniceViewsEnabled());
 
+    if (pubSubEncryptionKeyUrn != null && !pubSubEncryptionKeyUrn.trim().isEmpty()) {
+      updateStoreQueryParams.setPubSubEncryptionKeyUrn(pubSubEncryptionKeyUrn);
+    }
+
     if (srcStore.getReplicationMetadataVersionId() != -1) {
       updateStoreQueryParams.setReplicationMetadataVersionID(srcStore.getReplicationMetadataVersionId());
     }
@@ -813,6 +819,14 @@ public Optional<Long> getMinCompactionLagSeconds() {
     return getLong(MIN_COMPACTION_LAG_SECONDS);
   }
 
+  public UpdateStoreQueryParams setPubSubEncryptionKeyUrn(String pubSubEncryptionKeyUrn) {
+    return putString(PUB_SUB_ENCRYPTION_KEY_URN, pubSubEncryptionKeyUrn);
+  }
+
+  public Optional<String> getPubSubEncryptionKeyUrn() {
+    return getString(PUB_SUB_ENCRYPTION_KEY_URN);
+  }
+
   public Optional<String> getViewName() {
     return getString(STORE_VIEW_NAME);
   }

internal/venice-common/src/main/java/com/linkedin/venice/meta/ReadOnlyStore.java

@@ -1090,6 +1090,7 @@ public StoreProperties cloneStoreProperties() {
     storeProperties.setCompactionEnabled(isCompactionEnabled());
     storeProperties.setCompactionThresholdMilliseconds(getCompactionThresholdMilliseconds());
     storeProperties.setEncryptionEnabled(isEncryptionEnabled());
+    storeProperties.setPubSubEncryptionKeyUrn(getPubSubEncryptionKeyUrn());
     storeProperties.setMinCompactionLagSeconds(getMinCompactionLagSeconds());
     storeProperties.setMaxCompactionLagSeconds(getMaxCompactionLagSeconds());
     storeProperties.setMaxRecordSizeBytes(getMaxRecordSizeBytes());
@@ -1707,6 +1708,16 @@ public void setEncryptionEnabled(boolean encryptionEnabled) {
     throw new UnsupportedOperationException();
   }
 
+  @Override
+  public String getPubSubEncryptionKeyUrn() {
+    return this.delegate.getPubSubEncryptionKeyUrn();
+  }
+
+  @Override
+  public void setPubSubEncryptionKeyUrn(String pubSubEncryptionKeyUrn) {
+    throw new UnsupportedOperationException();
+  }
+
   @Override
   public long getMinCompactionLagSeconds() {
     return this.delegate.getMinCompactionLagSeconds();

internal/venice-common/src/main/java/com/linkedin/venice/meta/Store.java

@@ -351,6 +351,10 @@ default void setVersionTargetRegionPromoted(int versionNumber, boolean targetReg
 
   void setEncryptionEnabled(boolean encryptionEnabled);
 
+  String getPubSubEncryptionKeyUrn();
+
+  void setPubSubEncryptionKeyUrn(String pubSubEncryptionKeyUrn);
+
   long getMinCompactionLagSeconds();
 
   void setMinCompactionLagSeconds(long minCompactionLagSeconds);

internal/venice-common/src/main/java/com/linkedin/venice/meta/StoreInfo.java

@@ -79,6 +79,7 @@ public static StoreInfo fromStore(Store store) {
     storeInfo.setCompactionEnabled(store.isCompactionEnabled());
     storeInfo.setCompactionThreshold(store.getCompactionThresholdMilliseconds());
     storeInfo.setEncryptionEnabled(store.isEncryptionEnabled());
+    storeInfo.setPubSubEncryptionKeyUrn(store.getPubSubEncryptionKeyUrn());
     storeInfo.setMinCompactionLagSeconds(store.getMinCompactionLagSeconds());
     storeInfo.setMaxCompactionLagSeconds(store.getMaxCompactionLagSeconds());
     storeInfo.setMaxRecordSizeBytes(store.getMaxRecordSizeBytes());
@@ -371,6 +372,8 @@ public static StoreInfo fromStore(Store store) {
 
   private boolean encryptionEnabled;
 
+  private String pubSubEncryptionKeyUrn = "";
+
   private long minCompactionLagSeconds;
 
   private long maxCompactionLagSeconds;
@@ -916,6 +919,14 @@ public void setEncryptionEnabled(boolean encryptionEnabled) {
     this.encryptionEnabled = encryptionEnabled;
   }
 
+  public String getPubSubEncryptionKeyUrn() {
+    return pubSubEncryptionKeyUrn;
+  }
+
+  public void setPubSubEncryptionKeyUrn(String pubSubEncryptionKeyUrn) {
+    this.pubSubEncryptionKeyUrn = pubSubEncryptionKeyUrn;
+  }
+
   public long getMinCompactionLagSeconds() {
     return minCompactionLagSeconds;
   }