Skip to content

Commit d9d1387

Browse files
ruai0511ruai0511
committed
modify based on comments
Signed-off-by: Ruirui Zhang <[email protected]>
1 parent 634bce0 commit d9d1387

File tree

16 files changed

+588
-199
lines changed

16 files changed

+588
-199
lines changed

‎libs/autotagging-commons/src/main/java/org/opensearch/rule/service/IndexStoredRulePersistenceService.java

Lines changed: 23 additions & 14 deletions
Original file line numberDiff line numberDiff line change
@@ -20,10 +20,10 @@
2020
import org.opensearch.core.action.ActionListener;
2121
import org.opensearch.core.rest.RestStatus;
2222
import org.opensearch.index.query.BoolQueryBuilder;
23-
import org.opensearch.rule.Utils.IndexStoredRuleParser;
24-
import org.opensearch.rule.Utils.IndexStoredRuleUtils;
2523
import org.opensearch.rule.action.GetRuleRequest;
2624
import org.opensearch.rule.action.GetRuleResponse;
25+
import org.opensearch.rule.utils.IndexStoredRuleParser;
26+
import org.opensearch.rule.utils.IndexStoredRuleUtils;
2727
import org.opensearch.search.SearchHit;
2828
import org.opensearch.search.sort.SortOrder;
2929
import org.opensearch.transport.client.Client;
@@ -47,27 +47,35 @@ public class IndexStoredRulePersistenceService implements RulePersistenceService
4747
private final String indexName;
4848
private final Client client;
4949
private final FeatureType featureType;
50-
private final int maxReturnSizeAllowedPerGetRequest;
50+
private final int maxRulesPerGetRequest;
5151
private static final Logger logger = LogManager.getLogger(IndexStoredRulePersistenceService.class);
5252

53-
public IndexStoredRulePersistenceService(
54-
String indexName,
55-
Client client,
56-
FeatureType featureType,
57-
int maxReturnSizeAllowedPerGetRequest
58-
) {
53+
/**
54+
* Constructs an instance of {@link IndexStoredRulePersistenceService} with the specified parameters.
55+
* This service handles persistence and retrieval of stored rules within an OpenSearch index.
56+
* @param indexName - The name of the OpenSearch index where the rules are stored.
57+
* @param client - The OpenSearch client used to interact with the OpenSearch cluster.
58+
* @param featureType - The feature type associated with the stored rules.
59+
* @param maxRulesPerGetRequest - The maximum number of rules that can be returned in a single get request.
60+
*/
61+
public IndexStoredRulePersistenceService(String indexName, Client client, FeatureType featureType, int maxRulesPerGetRequest) {
5962
this.indexName = indexName;
6063
this.client = client;
6164
this.featureType = featureType;
62-
this.maxReturnSizeAllowedPerGetRequest = maxReturnSizeAllowedPerGetRequest;
65+
this.maxRulesPerGetRequest = maxRulesPerGetRequest;
6366
}
6467

68+
/**
69+
* Entry point for the get rule api logic in persistence service.
70+
* @param getRuleRequest the getRuleRequest to process.
71+
* @param listener the listener for GetRuleResponse.
72+
*/
6573
public void getRule(GetRuleRequest getRuleRequest, ActionListener<GetRuleResponse> listener) {
6674
getRuleFromIndex(getRuleRequest.getId(), getRuleRequest.getAttributeFilters(), getRuleRequest.getSearchAfter(), listener);
6775
}
6876

6977
/**
70-
* Entry point for the get rule api logic in persistence service. If id is provided, we only get a single rule.
78+
* Get rules from index. If id is provided, we only get a single rule.
7179
* Otherwise, we get all rules that satisfy the attributeFilters.
7280
* @param id - The id of the rule to get.
7381
* @param attributeFilters - A map containing the attributes that user want to filter on
@@ -85,9 +93,7 @@ public void getRuleFromIndex(
8593
// actions within this block are trusted and executed with system-level privileges.
8694
try (ThreadContext.StoredContext context = getContext()) {
8795
BoolQueryBuilder boolQuery = IndexStoredRuleUtils.buildGetRuleQuery(id, attributeFilters, featureType);
88-
SearchRequestBuilder searchRequest = client.prepareSearch(indexName)
89-
.setQuery(boolQuery)
90-
.setSize(maxReturnSizeAllowedPerGetRequest);
96+
SearchRequestBuilder searchRequest = client.prepareSearch(indexName).setQuery(boolQuery).setSize(maxRulesPerGetRequest);
9197
if (searchAfter != null) {
9298
searchRequest.addSort(_ID_STRING, SortOrder.ASC).searchAfter(new Object[] { searchAfter });
9399
}
@@ -120,6 +126,9 @@ private ThreadContext.StoredContext getContext() {
120126
return client.threadPool().getThreadContext().stashContext();
121127
}
122128

129+
/**
130+
* client getter
131+
*/
123132
public Client getClient() {
124133
return client;
125134
}

‎libs/autotagging-commons/src/main/java/org/opensearch/rule/service/RulePersistenceService.java

Lines changed: 10 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -12,6 +12,16 @@
1212
import org.opensearch.rule.action.GetRuleRequest;
1313
import org.opensearch.rule.action.GetRuleResponse;
1414

15+
/**
16+
* Interface for a service that handles rule persistence CRUD operations.
17+
* @opensearch.experimental
18+
*/
1519
public interface RulePersistenceService {
20+
21+
/**
22+
* Get rules based on the provided request.
23+
* @param request The request containing the details for retrieving the rule.
24+
* @param listener The listener that will handle the response or failure.
25+
*/
1626
void getRule(GetRuleRequest request, ActionListener<GetRuleResponse> listener);
1727
}

‎libs/autotagging-commons/src/main/java/org/opensearch/rule/Utils/IndexStoredRuleParser.java renamed to ‎libs/autotagging-commons/src/main/java/org/opensearch/rule/utils/IndexStoredRuleParser.java

Lines changed: 12 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -6,7 +6,7 @@
66
* compatible open source license.
77
*/
88

9-
package org.opensearch.rule.Utils;
9+
package org.opensearch.rule.utils;
1010

1111
import org.apache.logging.log4j.LogManager;
1212
import org.apache.logging.log4j.Logger;
@@ -19,12 +19,23 @@
1919

2020
import java.io.IOException;
2121

22+
/**
23+
* Utility class for parsing index stored rules into Rule objects.
24+
* @opensearch.experimental
25+
*/
2226
public class IndexStoredRuleParser {
27+
28+
/**
29+
* constructor for IndexStoredRuleParser
30+
*/
31+
public IndexStoredRuleParser() {}
32+
2333
private static final Logger logger = LogManager.getLogger(IndexStoredRuleParser.class);
2434

2535
/**
2636
* Parses a source string into a Rule object
2737
* @param source - The raw source string representing the rule to be parsed
38+
* @param featureType - The feature type to associate with the parsed rule
2839
*/
2940
public static Rule parseRule(String source, FeatureType featureType) {
3041
try (

‎libs/autotagging-commons/src/main/java/org/opensearch/rule/Utils/IndexStoredRuleUtils.java renamed to ‎libs/autotagging-commons/src/main/java/org/opensearch/rule/utils/IndexStoredRuleUtils.java

Lines changed: 17 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -6,7 +6,7 @@
66
* compatible open source license.
77
*/
88

9-
package org.opensearch.rule.Utils;
9+
package org.opensearch.rule.utils;
1010

1111
import org.opensearch.autotagging.Attribute;
1212
import org.opensearch.autotagging.FeatureType;
@@ -18,7 +18,23 @@
1818

1919
import static org.opensearch.autotagging.Rule._ID_STRING;
2020

21+
/**
22+
* Utility class that provides methods for the lifecycle of rules.
23+
* @opensearch.experimental
24+
*/
2125
public class IndexStoredRuleUtils {
26+
27+
/**
28+
* constructor for IndexStoredRuleUtils
29+
*/
30+
public IndexStoredRuleUtils() {}
31+
32+
/**
33+
* Builds a Boolean query to retrieve a rule by its ID or attribute filters.
34+
* @param id The ID of the rule to search for. If null, no ID-based filtering is applied.
35+
* @param attributeFilters A map of attributes and their corresponding filter values. This allows filtering by specific attribute values.
36+
* @param featureType The feature type that is required in the query.
37+
*/
2238
public static BoolQueryBuilder buildGetRuleQuery(String id, Map<Attribute, Set<String>> attributeFilters, FeatureType featureType) {
2339
BoolQueryBuilder boolQuery = QueryBuilders.boolQuery();
2440
if (id != null) {

‎libs/autotagging-commons/src/main/java/org/opensearch/rule/Utils/package-info.java renamed to ‎libs/autotagging-commons/src/main/java/org/opensearch/rule/utils/package-info.java

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -9,4 +9,4 @@
99
/**
1010
* This package contains utility classes for rules
1111
*/
12-
package org.opensearch.rule.Utils;
12+
package org.opensearch.rule.utils;

‎libs/autotagging-commons/src/test/java/org/opensearch/rule/RuleTestUtils.java

Lines changed: 157 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,157 @@
1+
/*
2+
* SPDX-License-Identifier: Apache-2.0
3+
*
4+
* The OpenSearch Contributors require contributions made to
5+
* this file be licensed under the Apache-2.0 license or a
6+
* compatible open source license.
7+
*/
8+
9+
package org.opensearch.rule;
10+
11+
import org.opensearch.autotagging.Attribute;
12+
import org.opensearch.autotagging.AutoTaggingRegistry;
13+
import org.opensearch.autotagging.FeatureType;
14+
import org.opensearch.autotagging.Rule;
15+
import org.opensearch.cluster.ClusterState;
16+
import org.opensearch.cluster.metadata.Metadata;
17+
import org.opensearch.cluster.metadata.QueryGroup;
18+
import org.opensearch.cluster.service.ClusterService;
19+
import org.opensearch.common.settings.Settings;
20+
import org.opensearch.common.util.concurrent.ThreadContext;
21+
import org.opensearch.rule.service.IndexStoredRulePersistenceService;
22+
import org.opensearch.threadpool.ThreadPool;
23+
import org.opensearch.transport.client.Client;
24+
25+
import java.util.Map;
26+
import java.util.Set;
27+
28+
import static org.junit.Assert.assertEquals;
29+
import static org.junit.Assert.assertTrue;
30+
import static org.mockito.Mockito.mock;
31+
import static org.mockito.Mockito.when;
32+
33+
public class RuleTestUtils {
34+
public static final String _ID_ONE = "AgfUO5Ja9yfvhdONlYi3TQ==";
35+
public static final String _ID_TWO = "G5iIq84j7eK1qIAAAAIH53=1";
36+
public static final String FEATURE_VALUE_ONE = "feature_value_one";
37+
public static final String FEATURE_VALUE_TWO = "feature_value_two";
38+
public static final String ATTRIBUTE_VALUE_ONE = "mock_attribute_one";
39+
public static final String ATTRIBUTE_VALUE_TWO = "mock_attribute_two";
40+
public static final String DESCRIPTION_ONE = "description_1";
41+
public static final String DESCRIPTION_TWO = "description_2";
42+
public static final String SEARCH_AFTER = "search_after_id";
43+
public static final String TIMESTAMP_ONE = "2024-01-26T08:58:57.558Z";
44+
public static final String TIMESTAMP_TWO = "2023-01-26T08:58:57.558Z";
45+
public static final String FEATURE_TYPE_NAME = "mock_feature_type";
46+
public static final String TEST_INDEX_NAME = ".test_index_for_rule";
47+
public static final Map<Attribute, Set<String>> ATTRIBUTE_MAP = Map.of(
48+
MockRuleAttributes.MOCK_RULE_ATTRIBUTE_ONE,
49+
Set.of(ATTRIBUTE_VALUE_ONE)
50+
);
51+
public static final Rule ruleOne = Rule.builder()
52+
.description(DESCRIPTION_ONE)
53+
.featureType(MockRuleFeatureType.INSTANCE)
54+
.featureValue(FEATURE_VALUE_ONE)
55+
.attributeMap(ATTRIBUTE_MAP)
56+
.updatedAt(TIMESTAMP_ONE)
57+
.build();
58+
59+
public static final Rule ruleTwo = Rule.builder()
60+
.description(DESCRIPTION_TWO)
61+
.featureType(MockRuleFeatureType.INSTANCE)
62+
.featureValue(FEATURE_VALUE_TWO)
63+
.attributeMap(Map.of(MockRuleAttributes.MOCK_RULE_ATTRIBUTE_TWO, Set.of(ATTRIBUTE_VALUE_TWO)))
64+
.updatedAt(TIMESTAMP_TWO)
65+
.build();
66+
67+
public static Map<String, Rule> ruleMap() {
68+
return Map.of(_ID_ONE, ruleOne, _ID_TWO, ruleTwo);
69+
}
70+
71+
public static IndexStoredRulePersistenceService setUpIndexStoredRulePersistenceService(Map<String, QueryGroup> queryGroupMap) {
72+
Client client = mock(Client.class);
73+
ClusterService clusterService = mock(ClusterService.class);
74+
ClusterState clusterState = mock(ClusterState.class);
75+
Metadata metadata = mock(Metadata.class);
76+
ThreadPool threadPool = mock(ThreadPool.class);
77+
78+
ThreadContext threadContext = new ThreadContext(Settings.EMPTY);
79+
when(client.threadPool()).thenReturn(threadPool);
80+
when(threadPool.getThreadContext()).thenReturn(threadContext);
81+
when(clusterService.state()).thenReturn(clusterState);
82+
when(clusterState.metadata()).thenReturn(metadata);
83+
when(metadata.queryGroups()).thenReturn(queryGroupMap);
84+
return new IndexStoredRulePersistenceService(TEST_INDEX_NAME, client, MockRuleFeatureType.INSTANCE, 50);
85+
}
86+
87+
public static void assertEqualRules(Map<String, Rule> mapOne, Map<String, Rule> mapTwo, boolean ruleUpdated) {
88+
assertEquals(mapOne.size(), mapTwo.size());
89+
for (Map.Entry<String, Rule> entry : mapOne.entrySet()) {
90+
String id = entry.getKey();
91+
assertTrue(mapTwo.containsKey(id));
92+
Rule one = mapOne.get(id);
93+
Rule two = mapTwo.get(id);
94+
assertEqualRule(one, two, ruleUpdated);
95+
}
96+
}
97+
98+
public static void assertEqualRule(Rule one, Rule two, boolean ruleUpdated) {
99+
if (ruleUpdated) {
100+
assertEquals(one.getDescription(), two.getDescription());
101+
assertEquals(one.getFeatureType(), two.getFeatureType());
102+
assertEquals(one.getFeatureValue(), two.getFeatureValue());
103+
assertEquals(one.getAttributeMap(), two.getAttributeMap());
104+
assertEquals(one.getAttributeMap(), two.getAttributeMap());
105+
} else {
106+
assertEquals(one, two);
107+
}
108+
}
109+
110+
public static class MockRuleFeatureType implements FeatureType {
111+
112+
public static final MockRuleFeatureType INSTANCE = new MockRuleFeatureType();
113+
114+
private MockRuleFeatureType() {}
115+
116+
static {
117+
INSTANCE.registerFeatureType();
118+
}
119+
120+
@Override
121+
public String getName() {
122+
return FEATURE_TYPE_NAME;
123+
}
124+
125+
@Override
126+
public Map<String, Attribute> getAllowedAttributesRegistry() {
127+
return Map.of(
128+
ATTRIBUTE_VALUE_ONE,
129+
MockRuleAttributes.MOCK_RULE_ATTRIBUTE_ONE,
130+
ATTRIBUTE_VALUE_TWO,
131+
MockRuleAttributes.MOCK_RULE_ATTRIBUTE_TWO
132+
);
133+
}
134+
135+
@Override
136+
public void registerFeatureType() {
137+
AutoTaggingRegistry.registerFeatureType(INSTANCE);
138+
}
139+
}
140+
141+
public enum MockRuleAttributes implements Attribute {
142+
MOCK_RULE_ATTRIBUTE_ONE(ATTRIBUTE_VALUE_ONE),
143+
MOCK_RULE_ATTRIBUTE_TWO(ATTRIBUTE_VALUE_TWO);
144+
;
145+
146+
private final String name;
147+
148+
MockRuleAttributes(String name) {
149+
this.name = name;
150+
}
151+
152+
@Override
153+
public String getName() {
154+
return name;
155+
}
156+
}
157+
}

‎libs/autotagging-commons/src/test/java/org/opensearch/rule/action/GetRuleRequestTests.java

Lines changed: 53 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,53 @@
1+
/*
2+
* SPDX-License-Identifier: Apache-2.0
3+
*
4+
* The OpenSearch Contributors require contributions made to
5+
* this file be licensed under the Apache-2.0 license or a
6+
* compatible open source license.
7+
*/
8+
9+
package org.opensearch.rule.action;
10+
11+
import org.opensearch.common.io.stream.BytesStreamOutput;
12+
import org.opensearch.core.common.io.stream.StreamInput;
13+
import org.opensearch.rule.RuleTestUtils;
14+
import org.opensearch.test.OpenSearchTestCase;
15+
16+
import java.io.IOException;
17+
import java.util.HashMap;
18+
19+
import static org.opensearch.rule.RuleTestUtils.ATTRIBUTE_MAP;
20+
import static org.opensearch.rule.RuleTestUtils.MockRuleFeatureType;
21+
import static org.opensearch.rule.RuleTestUtils.SEARCH_AFTER;
22+
import static org.opensearch.rule.RuleTestUtils._ID_ONE;
23+
24+
public class GetRuleRequestTests extends OpenSearchTestCase {
25+
26+
/**
27+
* Test case to verify the serialization and deserialization of GetRuleRequest
28+
*/
29+
public void testSerialization() throws IOException {
30+
GetRuleRequest request = new GetRuleRequest(_ID_ONE, ATTRIBUTE_MAP, null, RuleTestUtils.MockRuleFeatureType.INSTANCE);
31+
assertEquals(_ID_ONE, request.getId());
32+
BytesStreamOutput out = new BytesStreamOutput();
33+
request.writeTo(out);
34+
StreamInput streamInput = out.bytes().streamInput();
35+
GetRuleRequest otherRequest = new GetRuleRequest(streamInput);
36+
assertEquals(request.getId(), otherRequest.getId());
37+
assertEquals(request.getAttributeFilters(), otherRequest.getAttributeFilters());
38+
}
39+
40+
/**
41+
* Test case to verify the serialization and deserialization of GetRuleRequest when name is null
42+
*/
43+
public void testSerializationWithNull() throws IOException {
44+
GetRuleRequest request = new GetRuleRequest((String) null, new HashMap<>(), SEARCH_AFTER, MockRuleFeatureType.INSTANCE);
45+
assertNull(request.getId());
46+
BytesStreamOutput out = new BytesStreamOutput();
47+
request.writeTo(out);
48+
StreamInput streamInput = out.bytes().streamInput();
49+
GetRuleRequest otherRequest = new GetRuleRequest(streamInput);
50+
assertEquals(request.getId(), otherRequest.getId());
51+
assertEquals(request.getAttributeFilters(), otherRequest.getAttributeFilters());
52+
}
53+
}

0 commit comments

Comments
 (0)