From 81c9f32cdc0b1fa72a6f356a319226b7cbaf9d91 Mon Sep 17 00:00:00 2001 From: M Sazzadul Hoque <7600764+sazzad16@users.noreply.github.com> Date: Tue, 27 Aug 2024 16:58:43 +0600 Subject: [PATCH] Decoding FT.SEARCH reply can be disabled at field level (#3926) * Decoding FT.SEARCH reply can be disabled at field level * added javadoc in Document * make the pr not breaking --- .../redis/clients/jedis/CommandObjects.java | 32 +++--- .../redis/clients/jedis/search/Document.java | 101 +++++++++++++----- .../clients/jedis/search/FTSearchParams.java | 62 ++++++----- .../clients/jedis/search/SearchResult.java | 35 +++++- .../modules/search/SearchWithParamsTest.java | 39 ++++++- 5 files changed, 200 insertions(+), 69 deletions(-) diff --git a/src/main/java/redis/clients/jedis/CommandObjects.java b/src/main/java/redis/clients/jedis/CommandObjects.java index b89f06f62e..89eb3e16ca 100644 --- a/src/main/java/redis/clients/jedis/CommandObjects.java +++ b/src/main/java/redis/clients/jedis/CommandObjects.java @@ -3384,19 +3384,20 @@ public final CommandObject ftDropIndexDD(String indexName) { public final CommandObject ftSearch(String indexName, String query) { return new CommandObject<>(checkAndRoundRobinSearchCommand(SearchCommand.SEARCH, indexName).add(query), - getSearchResultBuilder(() -> new SearchResultBuilder(true, false, true))); + getSearchResultBuilder(null, () -> new SearchResultBuilder(true, false, true))); } public final CommandObject ftSearch(String indexName, String query, FTSearchParams params) { return new CommandObject<>(checkAndRoundRobinSearchCommand(SearchCommand.SEARCH, indexName) .add(query).addParams(params.dialectOptional(searchDialect.get())), - getSearchResultBuilder(() -> new SearchResultBuilder(!params.getNoContent(), params.getWithScores(), true))); + getSearchResultBuilder(params.getReturnFieldDecodeMap(), () -> new SearchResultBuilder( + !params.getNoContent(), params.getWithScores(), true, params.getReturnFieldDecodeMap()))); } public final CommandObject ftSearch(String indexName, Query query) { return new CommandObject<>(checkAndRoundRobinSearchCommand(SearchCommand.SEARCH, indexName) - .addParams(query.dialectOptional(searchDialect.get())), getSearchResultBuilder(() -> - new SearchResultBuilder(!query.getNoContent(), query.getWithScores(), true))); + .addParams(query.dialectOptional(searchDialect.get())), getSearchResultBuilder(null, + () -> new SearchResultBuilder(!query.getNoContent(), query.getWithScores(), true))); } @Deprecated @@ -3405,8 +3406,8 @@ public final CommandObject ftSearch(byte[] indexName, Query query) throw new UnsupportedOperationException("binary ft.search is not implemented with resp3."); } return new CommandObject<>(checkAndRoundRobinSearchCommand(commandArguments(SearchCommand.SEARCH), indexName) - .addParams(query.dialectOptional(searchDialect.get())), getSearchResultBuilder(() -> - new SearchResultBuilder(!query.getNoContent(), query.getWithScores(), false))); + .addParams(query.dialectOptional(searchDialect.get())), getSearchResultBuilder(null, + () -> new SearchResultBuilder(!query.getNoContent(), query.getWithScores(), false))); } public final CommandObject ftExplain(String indexName, Query query) { @@ -3449,20 +3450,27 @@ public final CommandObject>> ftProfi String indexName, FTProfileParams profileParams, Query query) { return new CommandObject<>(checkAndRoundRobinSearchCommand(SearchCommand.PROFILE, indexName) .add(SearchKeyword.SEARCH).addParams(profileParams).add(SearchKeyword.QUERY) - .addParams(query.dialectOptional(searchDialect.get())), new SearchProfileResponseBuilder<>( - getSearchResultBuilder(() -> new SearchResultBuilder(!query.getNoContent(), query.getWithScores(), true)))); + .addParams(query.dialectOptional(searchDialect.get())), + new SearchProfileResponseBuilder<>(getSearchResultBuilder(null, + () -> new SearchResultBuilder(!query.getNoContent(), query.getWithScores(), true)))); } public final CommandObject>> ftProfileSearch( String indexName, FTProfileParams profileParams, String query, FTSearchParams searchParams) { return new CommandObject<>(checkAndRoundRobinSearchCommand(SearchCommand.PROFILE, indexName) .add(SearchKeyword.SEARCH).addParams(profileParams).add(SearchKeyword.QUERY).add(query) - .addParams(searchParams.dialectOptional(searchDialect.get())), new SearchProfileResponseBuilder<>( - getSearchResultBuilder(() -> new SearchResultBuilder(!searchParams.getNoContent(), searchParams.getWithScores(), true)))); + .addParams(searchParams.dialectOptional(searchDialect.get())), + new SearchProfileResponseBuilder<>(getSearchResultBuilder(searchParams.getReturnFieldDecodeMap(), + () -> new SearchResultBuilder(!searchParams.getNoContent(), searchParams.getWithScores(), true, + searchParams.getReturnFieldDecodeMap())))); } - private Builder getSearchResultBuilder(Supplier> resp2) { - if (protocol == RedisProtocol.RESP3) return SearchResult.SEARCH_RESULT_BUILDER; + private Builder getSearchResultBuilder( + Map isReturnFieldDecode, Supplier> resp2) { + if (protocol == RedisProtocol.RESP3) { + return isReturnFieldDecode == null ? SearchResult.SEARCH_RESULT_BUILDER + : new SearchResult.PerFieldDecoderSearchResultBuilder(isReturnFieldDecode); + } return resp2.get(); } diff --git a/src/main/java/redis/clients/jedis/search/Document.java b/src/main/java/redis/clients/jedis/search/Document.java index 20149e581a..4c7d569023 100644 --- a/src/main/java/redis/clients/jedis/search/Document.java +++ b/src/main/java/redis/clients/jedis/search/Document.java @@ -7,6 +7,7 @@ import java.util.HashMap; import java.util.List; import java.util.Map; + import redis.clients.jedis.Builder; import redis.clients.jedis.BuilderFactory; import redis.clients.jedis.util.KeyValue; @@ -50,24 +51,6 @@ public Iterable> getProperties() { return fields.entrySet(); } - public static Document load(String id, double score, byte[] payload, List fields) { - return Document.load(id, score, fields, true); - } - - public static Document load(String id, double score, List fields, boolean decode) { - Document ret = new Document(id, score); - if (fields != null) { - for (int i = 0; i < fields.size(); i += 2) { - byte[] rawKey = fields.get(i); - byte[] rawValue = fields.get(i + 1); - String key = SafeEncoder.encode(rawKey); - Object value = rawValue == null ? null : decode ? SafeEncoder.encode(rawValue) : rawValue; - ret.set(key, value); - } - } - return ret; - } - /** * @return the document's id */ @@ -102,16 +85,22 @@ public Object get(String key) { */ public String getString(String key) { Object value = fields.get(key); - if (value instanceof String) { + if (value == null) { + return null; + } else if (value instanceof String) { return (String) value; + } else if (value instanceof byte[]) { + return SafeEncoder.encode((byte[]) value); + } else { + return String.valueOf(value); } - return value instanceof byte[] ? SafeEncoder.encode((byte[]) value) : value.toString(); } public boolean hasProperty(String key) { return fields.containsKey(key); } + // TODO: private ?? public Document set(String key, Object value) { fields.put(key, value); return this; @@ -122,7 +111,9 @@ public Document set(String key, Object value) { * * @param score new score to set * @return the document itself + * @deprecated */ + @Deprecated public Document setScore(float score) { this.score = (double) score; return this; @@ -134,13 +125,58 @@ public String toString() { ", properties:" + this.getProperties(); } - static Builder SEARCH_DOCUMENT = new Builder() { + /// RESP2 --> + public static Document load(String id, double score, byte[] payload, List fields) { + return Document.load(id, score, fields, true); + } + + public static Document load(String id, double score, List fields, boolean decode) { + return load(id, score, fields, decode, null); + } + + /** + * Parse document object from FT.SEARCH reply. + * @param id + * @param score + * @param fields + * @param decode + * @param isFieldDecode checked only if {@code decode=true} + * @return document + */ + public static Document load(String id, double score, List fields, boolean decode, + Map isFieldDecode) { + Document ret = new Document(id, score); + if (fields != null) { + for (int i = 0; i < fields.size(); i += 2) { + byte[] rawKey = fields.get(i); + byte[] rawValue = fields.get(i + 1); + String key = SafeEncoder.encode(rawKey); + Object value = rawValue == null ? null + : (decode && (isFieldDecode == null || !Boolean.FALSE.equals(isFieldDecode.get(key)))) + ? SafeEncoder.encode(rawValue) : rawValue; + ret.set(key, value); + } + } + return ret; + } + /// <-- RESP2 + + /// RESP3 --> + // TODO: final + static Builder SEARCH_DOCUMENT = new PerFieldDecoderDocumentBuilder((Map) null); + + static final class PerFieldDecoderDocumentBuilder extends Builder { private static final String ID_STR = "id"; private static final String SCORE_STR = "score"; - // private static final String FIELDS_STR = "fields"; private static final String FIELDS_STR = "extra_attributes"; + private final Map isFieldDecode; + + public PerFieldDecoderDocumentBuilder(Map isFieldDecode) { + this.isFieldDecode = isFieldDecode != null ? isFieldDecode : Collections.emptyMap(); + } + @Override public Document build(Object data) { List list = (List) data; @@ -157,13 +193,28 @@ public Document build(Object data) { score = BuilderFactory.DOUBLE.build(kv.getValue()); break; case FIELDS_STR: - fields = BuilderFactory.ENCODED_OBJECT_MAP.build(kv.getValue()); + fields = makeFieldsMap(isFieldDecode, kv.getValue()); break; } } -// assert id != null; -// if (fields == null) fields = Collections.emptyMap(); return new Document(id, score, fields); } }; + + private static Map makeFieldsMap(Map isDecode, Object data) { + if (data == null) return null; + + final List list = (List) data; + + Map map = new HashMap<>(list.size(), 1f); + list.stream().filter((kv) -> (kv != null && kv.getKey() != null && kv.getValue() != null)) + .forEach((kv) -> { + String key = BuilderFactory.STRING.build(kv.getKey()); + map.put(key, + (Boolean.FALSE.equals(isDecode.get(key)) ? BuilderFactory.RAW_OBJECT + : BuilderFactory.AGGRESSIVE_ENCODED_OBJECT).build(kv.getValue())); + }); + return map; + } + /// <-- RESP3 } diff --git a/src/main/java/redis/clients/jedis/search/FTSearchParams.java b/src/main/java/redis/clients/jedis/search/FTSearchParams.java index 2d4c6e0057..03c905edb3 100644 --- a/src/main/java/redis/clients/jedis/search/FTSearchParams.java +++ b/src/main/java/redis/clients/jedis/search/FTSearchParams.java @@ -24,8 +24,7 @@ public class FTSearchParams implements IParams { private final List filters = new LinkedList<>(); private Collection inKeys; private Collection inFields; - private Collection returnFields; - private Collection returnFieldNames; + private Collection returnFieldsNames; private boolean summarize; private SummarizeParams summarizeParams; private boolean highlight; @@ -43,6 +42,9 @@ public class FTSearchParams implements IParams { private Map params; private Integer dialect; + /// non command parameters + private Map returnFieldDecodeMap = null; + public FTSearchParams() { } @@ -78,17 +80,15 @@ public void addParams(CommandArguments args) { args.add(INFIELDS).add(inFields.size()).addObjects(inFields); } - if (returnFieldNames != null && !returnFieldNames.isEmpty()) { + if (returnFieldsNames != null && !returnFieldsNames.isEmpty()) { args.add(RETURN); LazyRawable returnCountObject = new LazyRawable(); args.add(returnCountObject); // holding a place for setting the total count later. int returnCount = 0; - for (FieldName fn : returnFieldNames) { + for (FieldName fn : returnFieldsNames) { returnCount += fn.addCommandArguments(args); } returnCountObject.setRaw(Protocol.toByteArray(returnCount)); - } else if (returnFields != null && !returnFields.isEmpty()) { - args.add(RETURN).add(returnFields.size()).addObjects(returnFields); } if (summarizeParams != null) { @@ -256,21 +256,15 @@ public FTSearchParams inFields(Collection fields) { * @return the query object itself */ public FTSearchParams returnFields(String... fields) { - if (returnFieldNames != null) { - Arrays.stream(fields).forEach(f -> returnFieldNames.add(FieldName.of(f))); - } else { - if (returnFields == null) { - returnFields = new ArrayList<>(); - } - Arrays.stream(fields).forEach(f -> returnFields.add(f)); + if (returnFieldsNames == null) { + returnFieldsNames = new ArrayList<>(); } + Arrays.stream(fields).forEach(f -> returnFieldsNames.add(FieldName.of(f))); return this; } public FTSearchParams returnField(FieldName field) { - initReturnFieldNames(); - returnFieldNames.add(field); - return this; + return returnFields(Collections.singleton(field)); } public FTSearchParams returnFields(FieldName... fields) { @@ -278,19 +272,30 @@ public FTSearchParams returnFields(FieldName... fields) { } public FTSearchParams returnFields(Collection fields) { - initReturnFieldNames(); - returnFieldNames.addAll(fields); + if (returnFieldsNames == null) { + returnFieldsNames = new ArrayList<>(); + } + returnFieldsNames.addAll(fields); return this; } - private void initReturnFieldNames() { - if (returnFieldNames == null) { - returnFieldNames = new ArrayList<>(); - } - if (returnFields != null) { - returnFields.forEach(f -> returnFieldNames.add(FieldName.of(f))); - returnFields = null; + public FTSearchParams returnField(String field, boolean decode) { + returnFields(field); + addReturnFieldDecode(field, decode); + return this; + } + + public FTSearchParams returnField(FieldName field, boolean decode) { + returnFields(field); + addReturnFieldDecode(field.getAttribute() != null ? field.getAttribute() : field.getName(), decode); + return this; + } + + private void addReturnFieldDecode(String returnName, boolean decode) { + if (returnFieldDecodeMap == null) { + returnFieldDecodeMap = new HashMap<>(); } + returnFieldDecodeMap.put(returnName, decode); } public FTSearchParams summarize() { @@ -436,14 +441,21 @@ public FTSearchParams dialectOptional(int dialect) { return this; } + @Internal public boolean getNoContent() { return noContent; } + @Internal public boolean getWithScores() { return withScores; } + @Internal + public Map getReturnFieldDecodeMap() { + return returnFieldDecodeMap; + } + /** * NumericFilter wraps a range filter on a numeric field. It can be inclusive or exclusive */ diff --git a/src/main/java/redis/clients/jedis/search/SearchResult.java b/src/main/java/redis/clients/jedis/search/SearchResult.java index cc28cc942a..55afbe0b24 100644 --- a/src/main/java/redis/clients/jedis/search/SearchResult.java +++ b/src/main/java/redis/clients/jedis/search/SearchResult.java @@ -3,9 +3,13 @@ import java.util.ArrayList; import java.util.Collections; import java.util.List; +import java.util.Map; +import java.util.Objects; import java.util.stream.Collectors; + import redis.clients.jedis.Builder; import redis.clients.jedis.BuilderFactory; +import redis.clients.jedis.annots.Internal; import redis.clients.jedis.util.KeyValue; /** @@ -43,10 +47,18 @@ public static class SearchResultBuilder extends Builder { private final boolean hasScores; private final boolean decode; + private final Map isFieldDecode; + public SearchResultBuilder(boolean hasContent, boolean hasScores, boolean decode) { + this(hasContent, hasScores, decode, null); + } + + public SearchResultBuilder(boolean hasContent, boolean hasScores, boolean decode, + Map isFieldDecode) { this.hasContent = hasContent; this.hasScores = hasScores; this.decode = decode; + this.isFieldDecode = isFieldDecode; } @Override @@ -75,18 +87,34 @@ public SearchResult build(Object data) { double score = hasScores ? BuilderFactory.DOUBLE.build(resp.get(i + scoreOffset)) : 1.0; List fields = hasContent ? (List) resp.get(i + contentOffset) : null; - documents.add(Document.load(id, score, fields, decode)); + documents.add(Document.load(id, score, fields, decode, isFieldDecode)); } return new SearchResult(totalResults, documents); } } - public static Builder SEARCH_RESULT_BUILDER = new Builder() { + /// RESP3 --> + // TODO: final + public static Builder SEARCH_RESULT_BUILDER + = new PerFieldDecoderSearchResultBuilder(Document.SEARCH_DOCUMENT); + + @Internal + public static final class PerFieldDecoderSearchResultBuilder extends Builder { private static final String TOTAL_RESULTS_STR = "total_results"; private static final String RESULTS_STR = "results"; + private final Builder documentBuilder; + + public PerFieldDecoderSearchResultBuilder(Map isFieldDecode) { + this(new Document.PerFieldDecoderDocumentBuilder(isFieldDecode)); + } + + private PerFieldDecoderSearchResultBuilder(Builder builder) { + this.documentBuilder = Objects.requireNonNull(builder); + } + @Override public SearchResult build(Object data) { List list = (List) data; @@ -100,7 +128,7 @@ public SearchResult build(Object data) { break; case RESULTS_STR: results = ((List) kv.getValue()).stream() - .map(Document.SEARCH_DOCUMENT::build) + .map(documentBuilder::build) .collect(Collectors.toList()); break; } @@ -108,4 +136,5 @@ public SearchResult build(Object data) { return new SearchResult(totalResults, results); } }; + /// <-- RESP3 } diff --git a/src/test/java/redis/clients/jedis/modules/search/SearchWithParamsTest.java b/src/test/java/redis/clients/jedis/modules/search/SearchWithParamsTest.java index f4c584d507..9f190a06ae 100644 --- a/src/test/java/redis/clients/jedis/modules/search/SearchWithParamsTest.java +++ b/src/test/java/redis/clients/jedis/modules/search/SearchWithParamsTest.java @@ -960,7 +960,7 @@ public void tagFieldParams() { } @Test - public void testReturnFields() { + public void returnFields() { assertOK(client.ftCreate(index, TextField.of("field1"), TextField.of("field2"))); Map doc = new HashMap<>(); @@ -975,10 +975,16 @@ public void testReturnFields() { Document ret = res.getDocuments().get(0); assertEquals("value1", ret.get("field1")); assertNull(ret.get("field2")); + + res = client.ftSearch(index, "*", FTSearchParams.searchParams().returnField("field1", true)); + assertEquals("value1", res.getDocuments().get(0).get("field1")); + + res = client.ftSearch(index, "*", FTSearchParams.searchParams().returnField("field1", false)); + assertArrayEquals("value1".getBytes(), (byte[]) res.getDocuments().get(0).get("field1")); } @Test - public void returnWithFieldNames() { + public void returnFieldsNames() { assertOK(client.ftCreate(index, TextField.of("a"), TextField.of("b"), TextField.of("c"))); Map map = new HashMap<>(); @@ -989,14 +995,39 @@ public void returnWithFieldNames() { // Query SearchResult res = client.ftSearch(index, "*", - FTSearchParams.searchParams().returnFields( - FieldName.of("a"), FieldName.of("b").as("d"))); + FTSearchParams.searchParams() + .returnFields(FieldName.of("a"), + FieldName.of("b").as("d"))); assertEquals(1, res.getTotalResults()); Document doc = res.getDocuments().get(0); assertEquals("value1", doc.get("a")); assertNull(doc.get("b")); assertEquals("value2", doc.get("d")); assertNull(doc.get("c")); + + res = client.ftSearch(index, "*", + FTSearchParams.searchParams() + .returnField(FieldName.of("a")) + .returnField(FieldName.of("b").as("d"))); + assertEquals(1, res.getTotalResults()); + assertEquals("value1", res.getDocuments().get(0).get("a")); + assertEquals("value2", res.getDocuments().get(0).get("d")); + + res = client.ftSearch(index, "*", + FTSearchParams.searchParams() + .returnField(FieldName.of("a"), true) + .returnField(FieldName.of("b").as("d"), true)); + assertEquals(1, res.getTotalResults()); + assertEquals("value1", res.getDocuments().get(0).get("a")); + assertEquals("value2", res.getDocuments().get(0).get("d")); + + res = client.ftSearch(index, "*", + FTSearchParams.searchParams() + .returnField(FieldName.of("a"), false) + .returnField(FieldName.of("b").as("d"), false)); + assertEquals(1, res.getTotalResults()); + assertArrayEquals("value1".getBytes(), (byte[]) res.getDocuments().get(0).get("a")); + assertArrayEquals("value2".getBytes(), (byte[]) res.getDocuments().get(0).get("d")); } @Test