From bd4fc25e6d7b74624f3acc20af4ab6773bd7714b Mon Sep 17 00:00:00 2001
From: Christoph Strobl <christoph.strobl@broadcom.com>
Date: Fri, 21 Mar 2025 14:45:43 +0100
Subject: [PATCH] allow keyId in raw format

---
 .../schema/IdentifiableJsonSchemaProperty.java | 18 +++++++++---------
 ...yableEncryptionCollectionCreationTests.java |  4 ++--
 2 files changed, 11 insertions(+), 11 deletions(-)

diff --git a/spring-data-mongodb/src/main/java/org/springframework/data/mongodb/core/schema/IdentifiableJsonSchemaProperty.java b/spring-data-mongodb/src/main/java/org/springframework/data/mongodb/core/schema/IdentifiableJsonSchemaProperty.java
index b77e97d48..d6d0dd494 100644
--- a/spring-data-mongodb/src/main/java/org/springframework/data/mongodb/core/schema/IdentifiableJsonSchemaProperty.java
+++ b/spring-data-mongodb/src/main/java/org/springframework/data/mongodb/core/schema/IdentifiableJsonSchemaProperty.java
@@ -1085,7 +1085,7 @@ public class IdentifiableJsonSchemaProperty<T extends JsonSchemaObject> implemen
 
 		private final JsonSchemaProperty targetProperty;
 		private final @Nullable String algorithm;
-		private final @Nullable String keyId;
+		private final @Nullable Object keyId;
 		private final @Nullable List<?> keyIds;
 
 		/**
@@ -1097,7 +1097,7 @@ public class IdentifiableJsonSchemaProperty<T extends JsonSchemaObject> implemen
 			this(target, null, null, null);
 		}
 
-		private EncryptedJsonSchemaProperty(JsonSchemaProperty target, @Nullable String algorithm, @Nullable String keyId,
+		private EncryptedJsonSchemaProperty(JsonSchemaProperty target, @Nullable String algorithm, @Nullable Object keyId,
 				@Nullable List<?> keyIds) {
 
 			Assert.notNull(target, "Target must not be null");
@@ -1152,6 +1152,10 @@ public class IdentifiableJsonSchemaProperty<T extends JsonSchemaObject> implemen
 			return new EncryptedJsonSchemaProperty(targetProperty, algorithm, keyId, null);
 		}
 
+		public EncryptedJsonSchemaProperty keyId(Object keyId) {
+			return new EncryptedJsonSchemaProperty(targetProperty, algorithm, keyId, null);
+		}
+
 		/**
 		 * @param keyId must not be {@literal null}.
 		 * @return new instance of {@link EncryptedJsonSchemaProperty}.
@@ -1179,11 +1183,7 @@ public class IdentifiableJsonSchemaProperty<T extends JsonSchemaObject> implemen
 			if (!ObjectUtils.isEmpty(keyId)) {
 				enc.append("keyId", keyId);
 			} else if (!ObjectUtils.isEmpty(keyIds)) {
-				if(keyIds.size() == 1) {
-					enc.append("keyId", keyIds.iterator().next());
-				} else {
-					enc.append("keyId", keyIds);
-				}
+				enc.append("keyId", keyIds);
 			}
 
 			Type type = extractPropertyType(propertySpecification);
@@ -1226,10 +1226,10 @@ public class IdentifiableJsonSchemaProperty<T extends JsonSchemaObject> implemen
 		}
 
 		public Object getKeyId() {
-			if(keyId != null) {
+			if (keyId != null) {
 				return keyId;
 			}
-			if(keyIds != null && keyIds.size() == 1) {
+			if (keyIds != null && keyIds.size() == 1) {
 				return keyIds.iterator().next();
 			}
 			return null;
diff --git a/spring-data-mongodb/src/test/java/org/springframework/data/mongodb/core/encryption/MongoQueryableEncryptionCollectionCreationTests.java b/spring-data-mongodb/src/test/java/org/springframework/data/mongodb/core/encryption/MongoQueryableEncryptionCollectionCreationTests.java
index ff5b1233b..da8ea9ccc 100644
--- a/spring-data-mongodb/src/test/java/org/springframework/data/mongodb/core/encryption/MongoQueryableEncryptionCollectionCreationTests.java
+++ b/spring-data-mongodb/src/test/java/org/springframework/data/mongodb/core/encryption/MongoQueryableEncryptionCollectionCreationTests.java
@@ -113,9 +113,9 @@ public class MongoQueryableEncryptionCollectionCreationTests {
 						range().min(-1L).max(1L).contention(0)));
 
 		CollectionOptions schemaOptions = CollectionOptions.encrypted(MongoJsonSchema.builder()
-				.property(queryable(encrypted(int32("encryptedInt")).keys(key1),
+				.property(queryable(encrypted(int32("encryptedInt")).keyId(key1),
 						new QueryCharacteristics(List.of(range().min(5).max(100).contention(1)))))
-				.property(queryable(encrypted(int64("nested.encryptedLong")).keys(key2),
+				.property(queryable(encrypted(int64("nested.encryptedLong")).keyId(key2),
 						new QueryCharacteristics(List.of(range().min(-1L).max(1L).contention(0)))))
 				.build());