feat: add @RedisKey annotation to populate entity fields with Redis keys during search operations (#538)

Add new @RedisKey annotation that allows entities to capture their Redis key during retrieval operations. This feature enables entities to know their full Redis key, which is useful for operations that need the complete key including custom prefixes.

- Add @RedisKey field annotation for marking fields to receive Redis keys
- Implement populateRedisKey utility method with class-level caching for performance
- Integrate key population into all entity retrieval paths (adapters, repositories, queries)
- Add optimized metadata caching in RedisEnhancedPersistentEntity with lazy initialization
- Support both Document and Hash entities with default and custom key prefixes
- Add comprehensive test coverage for all entity types and key prefix scenarios

The implementation uses efficient caching strategies to avoid repeated field scanning for entities that don't use the @RedisKey feature, ensuring minimal performance impact.

Author: bsbodden

redis-om-spring/src/main/java/com/redis/om/spring/RedisEnhancedKeyValueAdapter.java

@@ -231,7 +231,19 @@ public <T> T get(Object id, String keyspace, Class<T> type) {
     data.setId(stringId);
     data.setKeyspace(stringKeyspace);
 
-    return readTimeToLiveIfSet(binId, converter.read(type, data));
+    T entity = readTimeToLiveIfSet(binId, converter.read(type, data));
+    if (entity != null) {
+      String redisKey = new String(binId);
+      // Use optimized method if we can get the persistent entity
+      RedisPersistentEntity<?> persistentEntity = converter.getMappingContext().getPersistentEntity(type);
+      if (persistentEntity instanceof RedisEnhancedPersistentEntity) {
+        ((RedisEnhancedPersistentEntity<?>) persistentEntity).populateRedisKey(entity, redisKey);
+      } else {
+        // Fallback to utility method
+        com.redis.om.spring.util.ObjectUtils.populateRedisKey(entity, redisKey);
+      }
+    }
+    return entity;
   }
 
   /*
@@ -326,8 +338,19 @@ public <T> List<T> getAllOf(String keyspace, Class<T> type, long offset, int row
     query.limit(Math.toIntExact(offset), limit);
     SearchResult searchResult = searchOps.search(query);
 
+    // Get persistent entity once for all results
+    RedisPersistentEntity<?> persistentEntity = converter.getMappingContext().getPersistentEntity(type);
+    boolean useOptimized = persistentEntity instanceof RedisEnhancedPersistentEntity;
+
     return (List<T>) searchResult.getDocuments().stream() //
-        .map(d -> documentToObject(d, type, (MappingRedisOMConverter) converter)) //
+        .map(d -> {
+          Object entity = documentToObject(d, type, (MappingRedisOMConverter) converter);
+          if (useOptimized) {
+            return ((RedisEnhancedPersistentEntity<?>) persistentEntity).populateRedisKey(entity, d.getId());
+          } else {
+            return com.redis.om.spring.util.ObjectUtils.populateRedisKey(entity, d.getId());
+          }
+        }) //
         .toList();
   }
 

redis-om-spring/src/main/java/com/redis/om/spring/RedisJSONKeyValueAdapter.java

@@ -197,7 +197,18 @@ public <T> T get(String key, Class<T> type) {
     @SuppressWarnings(
       "unchecked"
     ) JSONOperations<String> ops = (JSONOperations<String>) redisJSONOperations;
-    return ops.get(key, type);
+    T entity = ops.get(key, type);
+    if (entity != null) {
+      // Use optimized method if we can get the persistent entity
+      RedisPersistentEntity<?> persistentEntity = mappingContext.getPersistentEntity(type);
+      if (persistentEntity instanceof RedisEnhancedPersistentEntity) {
+        ((RedisEnhancedPersistentEntity<?>) persistentEntity).populateRedisKey(entity, key);
+      } else {
+        // Fallback to utility method
+        ObjectUtils.populateRedisKey(entity, key);
+      }
+    }
+    return entity;
   }
 
   /**
@@ -222,8 +233,19 @@ public <T> List<T> getAllOf(String keyspace, Class<T> type, long offset, int row
     query.limit(Math.toIntExact(offset), limit);
     SearchResult searchResult = searchOps.search(query);
     Gson gson = gsonBuilder.create();
-    return searchResult.getDocuments().stream().map(d -> gson.fromJson(SafeEncoder.encode((byte[]) d.get("$")), type)) //
-        .toList();
+
+    // Get persistent entity once for all results
+    RedisPersistentEntity<?> persistentEntity = mappingContext.getPersistentEntity(type);
+    boolean useOptimized = persistentEntity instanceof RedisEnhancedPersistentEntity;
+
+    return searchResult.getDocuments().stream().map(d -> {
+      T entity = gson.fromJson(SafeEncoder.encode((byte[]) d.get("$")), type);
+      if (useOptimized) {
+        return ((RedisEnhancedPersistentEntity<?>) persistentEntity).populateRedisKey(entity, d.getId());
+      } else {
+        return ObjectUtils.populateRedisKey(entity, d.getId());
+      }
+    }).toList();
   }
 
   /**

redis-om-spring/src/main/java/com/redis/om/spring/annotations/RedisKey.java

@@ -0,0 +1,43 @@
+package com.redis.om.spring.annotations;
+
+import java.lang.annotation.*;
+
+/**
+ * Marks a field to be populated with the Redis key during search operations.
+ * <p>
+ * When a field is annotated with @RedisKey, it will automatically be populated
+ * with the full Redis key (including prefix) when the entity is retrieved
+ * through search operations.
+ * </p>
+ *
+ * <p>Example usage:</p>
+ * <pre>
+ * {@code
+ * @Document
+ * public class MyDocument {
+ * 
+ * @Id
+ *     private String id;
+ *
+ * @RedisKey
+ *           private String redisKey;
+ *
+ *           // other fields...
+ *           }
+ *           }
+ *           </pre>
+ *
+ *           <p>Note: Only one field per entity should be annotated with @RedisKey.</p>
+ *
+ * @author Brian Sam-Bodden
+ * @since 1.0.0
+ */
+@Documented
+@Retention(
+  RetentionPolicy.RUNTIME
+)
+@Target(
+  { ElementType.FIELD, ElementType.METHOD }
+)
+public @interface RedisKey {
+}

redis-om-spring/src/main/java/com/redis/om/spring/mapping/RedisEnhancedPersistentEntity.java

@@ -1,7 +1,9 @@
 package com.redis.om.spring.mapping;
 
+import java.lang.reflect.Field;
 import java.util.ArrayList;
 import java.util.List;
+import java.util.Optional;
 
 import org.springframework.data.keyvalue.core.mapping.KeySpaceResolver;
 import org.springframework.data.mapping.MappingException;
@@ -12,6 +14,9 @@
 import org.springframework.data.util.TypeInformation;
 import org.springframework.lang.Nullable;
 
+import com.redis.om.spring.annotations.RedisKey;
+import com.redis.om.spring.util.ObjectUtils;
+
 import jakarta.persistence.IdClass;
 
 /**
@@ -45,6 +50,7 @@ public class RedisEnhancedPersistentEntity<T> extends BasicRedisPersistentEntity
 
   private final List<RedisPersistentProperty> idProperties = new ArrayList<>();
   private final boolean hasIdClass;
+  private volatile Optional<Field> redisKeyField = null; // Lazy initialized
 
   /**
    * Creates a new {@code RedisEnhancedPersistentEntity} for the given type information.
@@ -125,4 +131,62 @@ public List<RedisPersistentProperty> getIdProperties() {
   public boolean isIdClassComposite() {
     return hasIdClass;
   }
+
+  /**
+   * Returns the field annotated with @RedisKey if present.
+   * This method caches the result for performance.
+   * 
+   * @return an {@link Optional} containing the @RedisKey field if present,
+   *         or an empty Optional if no field has the annotation
+   */
+  public Optional<Field> getRedisKeyField() {
+    if (redisKeyField == null) {
+      synchronized (this) {
+        if (redisKeyField == null) {
+          redisKeyField = findRedisKeyField();
+        }
+      }
+    }
+    return redisKeyField;
+  }
+
+  /**
+   * Populates the @RedisKey field of the given entity with the provided Redis key.
+   * This method uses the cached field information for optimal performance.
+   * 
+   * @param entity   the entity to populate
+   * @param redisKey the Redis key to set
+   * @param <E>      the entity type
+   * @return the entity with populated @RedisKey field
+   */
+  @SuppressWarnings(
+    "unchecked"
+  )
+  public <E> E populateRedisKey(E entity, String redisKey) {
+    if (entity == null || redisKey == null) {
+      return entity;
+    }
+
+    Optional<Field> fieldOpt = getRedisKeyField();
+    if (fieldOpt.isPresent()) {
+      try {
+        fieldOpt.get().set(entity, redisKey);
+      } catch (IllegalAccessException e) {
+        throw new RuntimeException("Failed to set @RedisKey field", e);
+      }
+    }
+
+    return entity;
+  }
+
+  private Optional<Field> findRedisKeyField() {
+    List<Field> fields = ObjectUtils.getDeclaredFieldsTransitively(getType());
+    for (Field field : fields) {
+      if (field.isAnnotationPresent(RedisKey.class)) {
+        field.setAccessible(true);
+        return Optional.of(field);
+      }
+    }
+    return Optional.empty();
+  }
 }

redis-om-spring/src/main/java/com/redis/om/spring/repository/query/RediSearchQuery.java

@@ -597,11 +597,13 @@ private Object parseDocumentResult(redis.clients.jedis.search.Document doc) {
 
     Gson gsonInstance = getGson();
 
-    return switch (dialect) {
+    Object entity = switch (dialect) {
       case ONE, TWO -> gsonInstance.fromJson(SafeEncoder.encode((byte[]) doc.get("$")), domainType);
       case THREE -> gsonInstance.fromJson(gsonInstance.fromJson(SafeEncoder.encode((byte[]) doc.get("$")),
           JsonArray.class).get(0), domainType);
     };
+
+    return ObjectUtils.populateRedisKey(entity, doc.getId());
   }
 
   private Object executeDeleteQuery(Object[] parameters) {

redis-om-spring/src/main/java/com/redis/om/spring/repository/query/RedisEnhancedQuery.java

@@ -580,8 +580,10 @@ private Object executeQuery(Object[] parameters) {
     if (queryMethod.getReturnedObjectType() == SearchResult.class) {
       result = searchResult;
     } else if (queryMethod.isPageQuery()) {
-      List<Object> content = searchResult.getDocuments().stream().map(d -> ObjectUtils.documentToObject(d, queryMethod
-          .getReturnedObjectType(), mappingConverter)).collect(Collectors.toList());
+      List<Object> content = searchResult.getDocuments().stream().map(d -> {
+        Object entity = ObjectUtils.documentToObject(d, queryMethod.getReturnedObjectType(), mappingConverter);
+        return ObjectUtils.populateRedisKey(entity, d.getId());
+      }).collect(Collectors.toList());
 
       if (maybePageable.isPresent()) {
         Pageable pageable = maybePageable.get();
@@ -591,14 +593,17 @@ private Object executeQuery(Object[] parameters) {
       }
     } else if (!queryMethod.isCollectionQuery()) {
       if (searchResult.getTotalResults() > 0 && !searchResult.getDocuments().isEmpty()) {
-        result = ObjectUtils.documentToObject(searchResult.getDocuments().get(0), queryMethod.getReturnedObjectType(),
-            mappingConverter);
+        redis.clients.jedis.search.Document doc = searchResult.getDocuments().get(0);
+        Object entity = ObjectUtils.documentToObject(doc, queryMethod.getReturnedObjectType(), mappingConverter);
+        result = ObjectUtils.populateRedisKey(entity, doc.getId());
       } else {
         result = null;
       }
     } else if (queryMethod.isCollectionQuery()) {
-      result = searchResult.getDocuments().stream().map(d -> ObjectUtils.documentToObject(d, queryMethod
-          .getReturnedObjectType(), mappingConverter)).collect(Collectors.toList());
+      result = searchResult.getDocuments().stream().map(d -> {
+        Object entity = ObjectUtils.documentToObject(d, queryMethod.getReturnedObjectType(), mappingConverter);
+        return ObjectUtils.populateRedisKey(entity, d.getId());
+      }).collect(Collectors.toList());
     } else {
       result = null;
     }

redis-om-spring/src/main/java/com/redis/om/spring/repository/support/SimpleRedisDocumentRepository.java

@@ -600,8 +600,10 @@ public Page<T> findAll(Pageable pageable) {
       Gson gson = gsonBuilder.create();
 
       if (searchResult.getTotalResults() > 0) {
-        List<T> content = searchResult.getDocuments().stream().map(d -> gson.fromJson(SafeEncoder.encode((byte[]) d.get(
-            "$")), metadata.getJavaType())).toList();
+        List<T> content = searchResult.getDocuments().stream().map(d -> {
+          T entity = gson.fromJson(SafeEncoder.encode((byte[]) d.get("$")), metadata.getJavaType());
+          return ObjectUtils.populateRedisKey(entity, d.getId());
+        }).toList();
 
         return new PageImpl<>(content, pageable, searchResult.getTotalResults());
       } else {

redis-om-spring/src/main/java/com/redis/om/spring/repository/support/SimpleRedisEnhancedRepository.java

@@ -255,7 +255,10 @@ public Page<T> findAll(Pageable pageable) {
         @SuppressWarnings(
           "unchecked"
         ) List<T> content = (List<T>) searchResult.getDocuments().stream() //
-            .map(d -> ObjectUtils.documentToObject(d, metadata.getJavaType(), mappingConverter)) //
+            .map(d -> {
+              Object entity = ObjectUtils.documentToObject(d, metadata.getJavaType(), mappingConverter);
+              return ObjectUtils.populateRedisKey(entity, d.getId());
+            }) //
             .toList();
         return new PageImpl<>(content, pageable, searchResult.getTotalResults());
       } else {

redis-om-spring/src/main/java/com/redis/om/spring/search/stream/SearchStreamImpl.java

@@ -648,11 +648,15 @@ private List<E> toEntityList(SearchResult searchResult) {
     if (projections.isEmpty()) {
       if (isDocument) {
         Gson g = getGson();
-        return searchResult.getDocuments().stream().map(d -> g.fromJson(SafeEncoder.encode((byte[]) d.get("$")),
-            entityClass)).toList();
+        return searchResult.getDocuments().stream().map(d -> {
+          E entity = g.fromJson(SafeEncoder.encode((byte[]) d.get("$")), entityClass);
+          return ObjectUtils.populateRedisKey(entity, d.getId());
+        }).toList();
       } else {
-        return searchResult.getDocuments().stream().map(d -> (E) ObjectUtils.documentToObject(d, entityClass,
-            mappingConverter)).toList();
+        return searchResult.getDocuments().stream().map(d -> {
+          E entity = (E) ObjectUtils.documentToObject(d, entityClass, mappingConverter);
+          return ObjectUtils.populateRedisKey(entity, d.getId());
+        }).toList();
       }
     } else {
       List<E> projectedEntities = new ArrayList<>();

redis-om-spring/src/main/java/com/redis/om/spring/util/ObjectUtils.java

@@ -14,6 +14,7 @@
 import java.nio.ByteOrder;
 import java.nio.FloatBuffer;
 import java.util.*;
+import java.util.concurrent.ConcurrentHashMap;
 import java.util.function.Function;
 import java.util.stream.Collectors;
 import java.util.stream.Stream;
@@ -60,6 +61,7 @@ public class ObjectUtils {
    */
   public static final Character REPLACEMENT_CHARACTER = '_';
   static final Set<String> JAVA_LITERAL_WORDS = Set.of("true", "false", "null");
+  private static final ConcurrentHashMap<Class<?>, Boolean> HAS_REDIS_KEY_CACHE = new ConcurrentHashMap<>();
   // Java reserved keywords
   static final Set<String> JAVA_RESERVED_WORDS = Collections.unmodifiableSet(Stream.of(
       // Unused
@@ -1033,4 +1035,48 @@ public static Object getPropertyValue(Object object, String propertyName) {
       throw new RuntimeException("Error getting property value", e);
     }
   }
+
+  /**
+   * Populates fields annotated with @RedisKey with the Redis key value.
+   * This method searches for fields annotated with @RedisKey and sets their value
+   * to the provided Redis key.
+   *
+   * @param entity   the entity to populate
+   * @param redisKey the Redis key to set
+   * @param <T>      the entity type
+   * @return the entity with populated @RedisKey field
+   */
+  public static <T> T populateRedisKey(T entity, String redisKey) {
+    if (entity == null || redisKey == null) {
+      return entity;
+    }
+
+    Class<?> clazz = entity.getClass();
+
+    // Quick check: if this class doesn't have @RedisKey fields, return early
+    Boolean hasRedisKey = HAS_REDIS_KEY_CACHE.computeIfAbsent(clazz, cls -> {
+      List<Field> fields = getDeclaredFieldsTransitively(cls);
+      return fields.stream().anyMatch(f -> f.isAnnotationPresent(com.redis.om.spring.annotations.RedisKey.class));
+    });
+
+    if (!hasRedisKey) {
+      return entity;
+    }
+
+    // If we get here, we know there's at least one @RedisKey field
+    List<Field> fields = getDeclaredFieldsTransitively(clazz);
+
+    for (Field field : fields) {
+      if (field.isAnnotationPresent(com.redis.om.spring.annotations.RedisKey.class)) {
+        try {
+          field.setAccessible(true);
+          field.set(entity, redisKey);
+        } catch (IllegalAccessException e) {
+          throw new RuntimeException("Failed to set @RedisKey field", e);
+        }
+      }
+    }
+
+    return entity;
+  }
 }