Skip to content

Commit

Permalink
opensearch-project#40 Working on capturing application.
Browse files Browse the repository at this point in the history 
Signed-off-by: jzonthemtn <[email protected]>
  • Loading branch information
@jzonthemtn
jzonthemtn committed Oct 23, 2024
1 parent ef0a928 commit eabd769
Show file tree
Hide file tree
Showing 6 changed files with 51 additions and 14 deletions.
13 changes: 12 additions & 1 deletion src/main/java/org/opensearch/ubi/QueryRequest.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -28,6 +28,7 @@ public class QueryRequest {
private final String clientId;
private final String userQuery;
private final String query;
private final String application;
private final Map<String, String> queryAttributes;
private final QueryResponse queryResponse;

Expand All @@ -43,13 +44,15 @@ public class QueryRequest {
* @param queryResponse The {@link QueryResponse} for this query request.
*/
public QueryRequest(final String queryId, final String userQuery, final String clientId, final String query,
final Map<String, String> queryAttributes, final QueryResponse queryResponse) {
final String application, final Map<String, String> queryAttributes,
final QueryResponse queryResponse) {

this.timestamp = sdf.format(new Date());
this.queryId = queryId;
this.clientId = clientId;
this.userQuery = userQuery;
this.query = query;
this.application = application;
this.queryAttributes = queryAttributes;
this.queryResponse = queryResponse;

Expand Down Expand Up @@ -88,6 +91,14 @@ public String getTimestamp() {
return timestamp;
}

/**
* Gets the application.
* @return The application.
*/
public String getApplication() {
return application;
}

/**
* Gets the query ID.
* @return The query ID.
Expand Down
3 changes: 2 additions & 1 deletion src/main/java/org/opensearch/ubi/UbiActionFilter.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -147,6 +147,7 @@ private ActionResponse handleSearchRequest(final SearchRequest searchRequest, Ac
final String userQuery = ubiParameters.getUserQuery();
final String userId = ubiParameters.getClientId();
final String objectIdField = ubiParameters.getObjectIdField();
final String application = ubiParameters.getApplication();
final Map<String, String> queryAttributes = ubiParameters.getQueryAttributes();

final String query = searchRequest.source().toString();
Expand All @@ -167,7 +168,7 @@ private ActionResponse handleSearchRequest(final SearchRequest searchRequest, Ac

final String queryResponseId = UUID.randomUUID().toString();
final QueryResponse queryResponse = new QueryResponse(queryId, queryResponseId, queryResponseHitIds);
final QueryRequest queryRequest = new QueryRequest(queryId, userQuery, userId, query, queryAttributes, queryResponse);
final QueryRequest queryRequest = new QueryRequest(queryId, userQuery, userId, query, application, queryAttributes, queryResponse);

final String dataPrepperUrl = environment.settings().get(UbiSettings.DATA_PREPPER_URL);
if (dataPrepperUrl != null) {
Expand Down
28 changes: 26 additions & 2 deletions src/main/java/org/opensearch/ubi/ext/UbiParameters.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -13,7 +13,6 @@
import org.opensearch.core.common.io.stream.StreamInput;
import org.opensearch.core.common.io.stream.StreamOutput;
import org.opensearch.core.common.io.stream.Writeable;
import org.opensearch.core.xcontent.ConstructingObjectParser;
import org.opensearch.core.xcontent.ObjectParser;
import org.opensearch.core.xcontent.ToXContentObject;
import org.opensearch.core.xcontent.XContentBuilder;
Expand All @@ -36,6 +35,7 @@ public class UbiParameters implements Writeable, ToXContentObject {
private static final ParseField QUERY_ID = new ParseField("query_id");
private static final ParseField USER_QUERY = new ParseField("user_query");
private static final ParseField CLIENT_ID = new ParseField("client_id");
private static final ParseField APPLICATION = new ParseField("application");
private static final ParseField OBJECT_ID_FIELD = new ParseField("object_id_field");
private static final ParseField QUERY_ATTRIBUTES = new ParseField("query_attributes");

Expand All @@ -44,6 +44,7 @@ public class UbiParameters implements Writeable, ToXContentObject {
PARSER.declareString(UbiParameters::setQueryId, QUERY_ID);
PARSER.declareString(UbiParameters::setUserQuery, USER_QUERY);
PARSER.declareString(UbiParameters::setClientId, CLIENT_ID);
PARSER.declareString(UbiParameters::setObjectIdField, APPLICATION);
PARSER.declareString(UbiParameters::setObjectIdField, OBJECT_ID_FIELD);
PARSER.declareObject(UbiParameters::setQueryAttributes, (p, c) -> p.mapStrings(), QUERY_ATTRIBUTES);
}
Expand Down Expand Up @@ -79,6 +80,7 @@ public static UbiParameters getUbiParameters(final SearchRequest request) {
private String queryId;
private String userQuery;
private String clientId;
private String application;
private String objectIdField;
private Map<String, String> queryAttributes;

Expand All @@ -97,6 +99,7 @@ public UbiParameters(StreamInput input) throws IOException {
this.queryId = input.readString();
this.userQuery = input.readOptionalString();
this.clientId = input.readOptionalString();
this.application = input.readOptionalString();
this.objectIdField = input.readOptionalString();
this.queryAttributes = (Map<String, String>) input.readGenericValue();
}
Expand All @@ -106,13 +109,15 @@ public UbiParameters(StreamInput input) throws IOException {
* @param queryId The query ID.
* @param userQuery The user-entered search query.
* @param clientId The client ID.
* @param application The name of the application making the query.
* @param objectIdField The object ID field.
* @param queryAttributes Optional attributes for UBI.
*/
public UbiParameters(String queryId, String userQuery, String clientId, String objectIdField, Map<String, String> queryAttributes) {
public UbiParameters(String queryId, String userQuery, String clientId, String application, String objectIdField, Map<String, String> queryAttributes) {
this.queryId = queryId;
this.userQuery = userQuery;
this.clientId = clientId;
this.application = application;
this.objectIdField = objectIdField;
this.queryAttributes = queryAttributes;
}
Expand All @@ -123,6 +128,7 @@ public XContentBuilder toXContent(XContentBuilder xContentBuilder, Params params
.field(QUERY_ID.getPreferredName(), this.queryId)
.field(USER_QUERY.getPreferredName(), this.userQuery)
.field(CLIENT_ID.getPreferredName(), this.clientId)
.field(APPLICATION.getPreferredName(), this.application)
.field(OBJECT_ID_FIELD.getPreferredName(), this.objectIdField)
.field(QUERY_ATTRIBUTES.getPreferredName(), this.queryAttributes);
}
Expand All @@ -132,6 +138,7 @@ public void writeTo(StreamOutput out) throws IOException {
out.writeString(getQueryId());
out.writeOptionalString(userQuery);
out.writeOptionalString(clientId);
out.writeOptionalString(application);
out.writeOptionalString(objectIdField);
out.writeGenericValue(queryAttributes);
}
Expand Down Expand Up @@ -159,6 +166,7 @@ public boolean equals(Object o) {
return Objects.equals(this.queryId, other.getQueryId())
&& Objects.equals(this.userQuery, other.getUserQuery())
&& Objects.equals(this.clientId, other.getClientId())
&& Objects.equals(this.application, other.getApplication())
&& Objects.equals(this.objectIdField, other.getObjectIdField())
&& Objects.equals(this.queryAttributes, other.getQueryAttributes());
}
Expand Down Expand Up @@ -203,6 +211,22 @@ public void setClientId(String clientId) {
this.clientId = clientId;
}

/**
* Gets the application.
* @return The application.
*/
public String getApplication() {
return application;
}

/**
* Set the application.
* @param application The application.
*/
public void setApplication(String application) {
this.application = application;
}

/**
* Get the object ID field.
* @return The object ID field.
Expand Down
3 changes: 2 additions & 1 deletion src/main/resources/queries-mapping.json
View file
Original file line number Diff line number Diff line change
Expand Up @@ -10,6 +10,7 @@
"fields": { "keyword": { "type": "keyword", "ignore_above": 256 } }
},
"query_attributes": { "type": "flat_object" },
"client_id": { "type": "keyword", "ignore_above": 100 }
"client_id": { "type": "keyword", "ignore_above": 100 },
"application": { "type": "keyword", "ignore_above": 100 }
}
}
8 changes: 4 additions & 4 deletions src/test/java/org/opensearch/ubi/UbIParametersTests.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -80,7 +80,7 @@ public void reset() throws IOException {
}

public void testUbiParameters() {
final UbiParameters params = new UbiParameters("query_id", "user_query", "client_id", "object_id", Collections.emptyMap());
final UbiParameters params = new UbiParameters("query_id", "user_query", "client_id", "app", "object_id", Collections.emptyMap());
UbiParametersExtBuilder extBuilder = new UbiParametersExtBuilder();
extBuilder.setParams(params);
SearchSourceBuilder srcBulder = SearchSourceBuilder.searchSource().ext(List.of(extBuilder));
Expand All @@ -90,7 +90,7 @@ public void testUbiParameters() {
}

public void testWriteTo() throws IOException {
final UbiParameters params = new UbiParameters("query_id", "user_query", "client_id", "object_id", Collections.emptyMap());
final UbiParameters params = new UbiParameters("query_id", "user_query", "client_id", "app", "object_id", Collections.emptyMap());
StreamOutput output = new DummyStreamOutput();
params.writeTo(output);
List<String> actual = ((DummyStreamOutput) output).getList();
Expand All @@ -101,7 +101,7 @@ public void testWriteTo() throws IOException {
}

public void testToXContent() throws IOException {
final UbiParameters params = new UbiParameters("query_id", "user_query", "client_id", "object_id", Collections.emptyMap());
final UbiParameters params = new UbiParameters("query_id", "user_query", "client_id", "app", "object_id", Collections.emptyMap());
XContent xc = mock(XContent.class);
OutputStream os = mock(OutputStream.class);
XContentGenerator generator = mock(XContentGenerator.class);
Expand All @@ -111,7 +111,7 @@ public void testToXContent() throws IOException {
}

public void testToXContentAllOptionalParameters() throws IOException {
final UbiParameters params = new UbiParameters("query_id", "user_query", "client_id", "object_id", Collections.emptyMap());
final UbiParameters params = new UbiParameters("query_id", "user_query", "client_id", "app", "object_id", Collections.emptyMap());
XContent xc = mock(XContent.class);
OutputStream os = mock(OutputStream.class);
XContentGenerator generator = mock(XContentGenerator.class);
Expand Down
10 changes: 5 additions & 5 deletions src/test/java/org/opensearch/ubi/UbiParametersExtBuilderTests.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -34,7 +34,7 @@ public void testCtor() {
final Map<String, String> queryAttributes = new HashMap<>();

final UbiParametersExtBuilder builder = new UbiParametersExtBuilder();
final UbiParameters parameters = new UbiParameters("query_id", "user_query", "client_id", "object_id_field", queryAttributes);
final UbiParameters parameters = new UbiParameters("query_id", "user_query", "client_id", "app", "object_id_field", queryAttributes);
builder.setParams(parameters);
assertEquals(parameters, builder.getParams());

Expand All @@ -49,7 +49,7 @@ public void testParse() throws IOException {
}

public void testXContentRoundTrip() throws IOException {
UbiParameters param1 = new UbiParameters("query_id", "user_query", "client_id", "object_id_field", Collections.emptyMap());
UbiParameters param1 = new UbiParameters("query_id", "user_query", "client_id", "app", "object_id_field", Collections.emptyMap());
UbiParametersExtBuilder extBuilder = new UbiParametersExtBuilder();
extBuilder.setParams(param1);
XContentType xContentType = randomFrom(XContentType.values());
Expand All @@ -65,7 +65,7 @@ public void testXContentRoundTrip() throws IOException {
}

public void testXContentRoundTripAllValues() throws IOException {
UbiParameters param1 = new UbiParameters("query_id", "user_query", "client_id", "object_id_field", Collections.emptyMap());
UbiParameters param1 = new UbiParameters("query_id", "user_query", "client_id", "app","object_id_field", Collections.emptyMap());
UbiParametersExtBuilder extBuilder = new UbiParametersExtBuilder();
extBuilder.setParams(param1);
XContentType xContentType = randomFrom(XContentType.values());
Expand All @@ -76,7 +76,7 @@ public void testXContentRoundTripAllValues() throws IOException {
}

public void testStreamRoundTrip() throws IOException {
UbiParameters param1 = new UbiParameters("query_id", "user_query", "client_id", "object_id_field", Collections.emptyMap());
UbiParameters param1 = new UbiParameters("query_id", "user_query", "client_id", "app","object_id_field", Collections.emptyMap());
UbiParametersExtBuilder extBuilder = new UbiParametersExtBuilder();
extBuilder.setParams(param1);
BytesStreamOutput bso = new BytesStreamOutput();
Expand All @@ -91,7 +91,7 @@ public void testStreamRoundTrip() throws IOException {
}

public void testStreamRoundTripAllValues() throws IOException {
UbiParameters param1 = new UbiParameters("query_id", "user_query", "client_id", "object_id_field", Collections.emptyMap());
UbiParameters param1 = new UbiParameters("query_id", "user_query", "client_id", "app","object_id_field", Collections.emptyMap());
UbiParametersExtBuilder extBuilder = new UbiParametersExtBuilder();
extBuilder.setParams(param1);
BytesStreamOutput bso = new BytesStreamOutput();
Expand Down

0 comments on commit eabd769

Please sign in to comment.