diff --git a/Dockerfile b/Dockerfile index ee5781a..b94f72d 100644 --- a/Dockerfile +++ b/Dockerfile @@ -1,4 +1,4 @@ -FROM opensearchstaging/opensearch:3.1.0-SNAPSHOT +FROM opensearchproject/opensearch:3.1.0 ARG UBI_VERSION="3.1.0.0-SNAPSHOT" @@ -7,4 +7,4 @@ COPY ./build/distributions/opensearch-ubi-${UBI_VERSION}.zip /tmp/ # Required for OTel capabilities. #RUN /usr/share/opensearch/bin/opensearch-plugin install --batch telemetry-otel -RUN /usr/share/opensearch/bin/opensearch-plugin install --batch file:/tmp/opensearch-ubi-${UBI_VERSION}.zip \ No newline at end of file +RUN /usr/share/opensearch/bin/opensearch-plugin remove opensearch-ubi || true && /usr/share/opensearch/bin/opensearch-plugin install --batch file:/tmp/opensearch-ubi-${UBI_VERSION}.zip diff --git a/src/main/java/org/opensearch/ubi/QueryRequest.java b/src/main/java/org/opensearch/ubi/QueryRequest.java index 7a5195c..d49959e 100644 --- a/src/main/java/org/opensearch/ubi/QueryRequest.java +++ b/src/main/java/org/opensearch/ubi/QueryRequest.java @@ -26,7 +26,7 @@ public class QueryRequest { private final String timestamp; private final String queryId; private final String clientId; - private final String userQuery; + private final Map userQuery; private final String query; private final String application; private final Map queryAttributes; @@ -43,7 +43,7 @@ public class QueryRequest { * @param queryAttributes An optional map of additional attributes for the query. * @param queryResponse The {@link QueryResponse} for this query request. */ - public QueryRequest(final String queryId, final String userQuery, final String clientId, final String query, + public QueryRequest(final String queryId, final Map userQuery, final String clientId, final String query, final String application, final Map queryAttributes, final QueryResponse queryResponse) { @@ -108,14 +108,22 @@ public String getQueryId() { } /** - * Gets the user query. - * @return The user query. + * Gets the user query map. + * @return The user query map. */ - public String getUserQuery() { - if(userQuery == null) { + public Map getUserQuery() { + return userQuery; + } + + /** + * Gets the user query text (convenience accessor). + * @return The "text" value of user_query or empty string if missing. + */ + public String getUserQueryText() { + if (userQuery == null) { return ""; } - return userQuery; + return userQuery.getOrDefault("text", ""); } /** diff --git a/src/main/java/org/opensearch/ubi/UbiActionFilter.java b/src/main/java/org/opensearch/ubi/UbiActionFilter.java index f01fbf4..127f51c 100644 --- a/src/main/java/org/opensearch/ubi/UbiActionFilter.java +++ b/src/main/java/org/opensearch/ubi/UbiActionFilter.java @@ -129,7 +129,7 @@ private ActionResponse handleSearchRequest(final SearchRequest searchRequest, Ac if (ubiParameters != null) { final String queryId = ubiParameters.getQueryId(); - final String userQuery = ubiParameters.getUserQuery(); + final Map userQuery = ubiParameters.getUserQuery(); final String userId = ubiParameters.getClientId(); final String objectIdField = ubiParameters.getObjectIdField(); final String application = ubiParameters.getApplication(); @@ -263,7 +263,10 @@ private void sendOtelTrace(final Task task, final Tracer tracer, final QueryRequ span.addAttribute("ubi.user_id", queryRequest.getQueryId()); span.addAttribute("ubi.query", queryRequest.getQuery()); - span.addAttribute("ubi.user_query", queryRequest.getUserQuery()); + span.addAttribute("ubi.user_query", queryRequest.getUserQueryText()); + for (final Map.Entry e : queryRequest.getUserQuery().entrySet()) { + span.addAttribute("ubi.user_query." + e.getKey(), e.getValue()); + } span.addAttribute("ubi.client_id", queryRequest.getClientId()); span.addAttribute("ubi.timestamp", queryRequest.getTimestamp()); diff --git a/src/main/java/org/opensearch/ubi/ext/UbiParameters.java b/src/main/java/org/opensearch/ubi/ext/UbiParameters.java index 412b5ee..fe8c039 100644 --- a/src/main/java/org/opensearch/ubi/ext/UbiParameters.java +++ b/src/main/java/org/opensearch/ubi/ext/UbiParameters.java @@ -42,7 +42,6 @@ public class UbiParameters implements Writeable, ToXContentObject { static { PARSER = new ObjectParser<>(UbiParametersExtBuilder.UBI_PARAMETER_NAME, UbiParameters::new); PARSER.declareString(UbiParameters::setQueryId, QUERY_ID); - PARSER.declareString(UbiParameters::setUserQuery, USER_QUERY); PARSER.declareString(UbiParameters::setClientId, CLIENT_ID); PARSER.declareString(UbiParameters::setApplication, APPLICATION); PARSER.declareString(UbiParameters::setObjectIdField, OBJECT_ID_FIELD); @@ -78,7 +77,7 @@ public static UbiParameters getUbiParameters(final SearchRequest request) { } private String queryId; - private String userQuery; + private Map userQuery; private String clientId; private String application; private String objectIdField; @@ -97,7 +96,7 @@ public UbiParameters() {} @SuppressWarnings("unchecked") public UbiParameters(StreamInput input) throws IOException { this.queryId = input.readString(); - this.userQuery = input.readOptionalString(); + this.userQuery = (Map) input.readGenericValue(); this.clientId = input.readOptionalString(); this.application = input.readOptionalString(); this.objectIdField = input.readOptionalString(); @@ -114,6 +113,18 @@ public UbiParameters(StreamInput input) throws IOException { * @param queryAttributes Optional attributes for UBI. */ public UbiParameters(String queryId, String userQuery, String clientId, String application, String objectIdField, Map queryAttributes) { + this.queryId = queryId; + if (userQuery != null) { + this.userQuery = new HashMap<>(); + this.userQuery.put("text", userQuery); + } + this.clientId = clientId; + this.application = application; + this.objectIdField = objectIdField; + this.queryAttributes = queryAttributes; + } + + public UbiParameters(String queryId, Map userQuery, String clientId, String application, String objectIdField, Map queryAttributes) { this.queryId = queryId; this.userQuery = userQuery; this.clientId = clientId; @@ -136,7 +147,7 @@ public XContentBuilder toXContent(XContentBuilder xContentBuilder, Params params @Override public void writeTo(StreamOutput out) throws IOException { out.writeString(getQueryId()); - out.writeOptionalString(userQuery); + out.writeGenericValue(userQuery); out.writeOptionalString(clientId); out.writeOptionalString(application); out.writeOptionalString(objectIdField); @@ -150,7 +161,54 @@ public void writeTo(StreamOutput out) throws IOException { * @throws IOException Thrown if the parameters cannot be read. */ public static UbiParameters parse(XContentParser parser) throws IOException { - return PARSER.parse(parser, null); + String queryId = null; + Map userQuery = null; + String clientId = null; + String application = null; + String objectIdField = null; + Map queryAttributes = null; + + if (parser.currentToken() == null) { + parser.nextToken(); + } + if (parser.currentToken() != XContentParser.Token.START_OBJECT) { + throw new IOException("Expected START_OBJECT for UBI parameters"); + } + while (parser.nextToken() != XContentParser.Token.END_OBJECT) { + String fieldName = parser.currentName(); + XContentParser.Token token = parser.nextToken(); + if (QUERY_ID.match(fieldName, parser.getDeprecationHandler())) { + queryId = parser.text(); + } else if (USER_QUERY.match(fieldName, parser.getDeprecationHandler())) { + if (token == XContentParser.Token.START_OBJECT) { + userQuery = parser.mapStrings(); + } else if (token == XContentParser.Token.VALUE_STRING) { + userQuery = new HashMap<>(); + userQuery.put("text", parser.text()); + } else if (token == XContentParser.Token.VALUE_NULL) { + userQuery = null; + } else { + throw new IOException("Unsupported token for user_query: " + token); + } + } else if (CLIENT_ID.match(fieldName, parser.getDeprecationHandler())) { + clientId = token == XContentParser.Token.VALUE_NULL ? null : parser.text(); + } else if (APPLICATION.match(fieldName, parser.getDeprecationHandler())) { + application = token == XContentParser.Token.VALUE_NULL ? null : parser.text(); + } else if (OBJECT_ID_FIELD.match(fieldName, parser.getDeprecationHandler())) { + objectIdField = token == XContentParser.Token.VALUE_NULL ? null : parser.text(); + } else if (QUERY_ATTRIBUTES.match(fieldName, parser.getDeprecationHandler())) { + if (token == XContentParser.Token.START_OBJECT) { + queryAttributes = parser.mapStrings(); + } else if (token == XContentParser.Token.VALUE_NULL) { + queryAttributes = null; + } else { + throw new IOException("Unsupported token for query_attributes: " + token); + } + } else { + parser.skipChildren(); + } + } + return new UbiParameters(queryId, userQuery, clientId, application, objectIdField, queryAttributes); } @Override @@ -244,21 +302,49 @@ public void setObjectIdField(String objectIdField) { } /** - * Get the user query. - * @return The user query. + * Get the user query map. + * @return The user query map. */ - public String getUserQuery() { + public Map getUserQuery() { + if (userQuery == null) { + userQuery = new HashMap<>(); + } return userQuery; } /** - * Set the user query. - * @param userQuery The user query. + * Convenience accessor for the user query text value. + * @return The user query "text" value or empty string if not present. + */ + public String getUserQueryText() { + return userQuery != null ? userQuery.getOrDefault("text", "") : ""; + } + + /** + * Set the user query map. + * @param userQuery The user query map. */ - public void setUserQuery(String userQuery) { + public void setUserQuery(Map userQuery) { this.userQuery = userQuery; } + /** + * Set a legacy string user query; wraps to {"text": {@literal }}. + * @param userQuery The legacy user query string. + */ + public void setLegacyUserQuery(String userQuery) { + if (userQuery == null) { + this.userQuery = null; + } else { + if (this.userQuery == null) { + this.userQuery = new HashMap<>(); + } else { + this.userQuery.clear(); + } + this.userQuery.put("text", userQuery); + } + } + /** * Get the attributes. * @return A map of attributes. diff --git a/src/main/resources/events-mapping.json b/src/main/resources/events-mapping.json index 0bfe27b..e9dd2f7 100644 --- a/src/main/resources/events-mapping.json +++ b/src/main/resources/events-mapping.json @@ -6,7 +6,7 @@ "query_id": { "type": "keyword", "ignore_above": 100 }, "message": { "type": "keyword", "ignore_above": 1024 }, "message_type": { "type": "keyword", "ignore_above": 100 }, - "user_query": { "type": "keyword" }, + "user_query": { "type": "flat_object" }, "timestamp": { "type": "date", "format":"strict_date_time", @@ -44,4 +44,4 @@ } } } -} \ No newline at end of file +} diff --git a/src/main/resources/queries-mapping.json b/src/main/resources/queries-mapping.json index 3cf851d..a05383e 100644 --- a/src/main/resources/queries-mapping.json +++ b/src/main/resources/queries-mapping.json @@ -6,7 +6,7 @@ "query": { "type": "text" }, "query_response_id": { "type": "keyword", "ignore_above": 100 }, "query_response_hit_ids": { "type": "keyword" }, - "user_query": { "type": "keyword" }, + "user_query": { "type": "flat_object" }, "query_attributes": { "type": "flat_object" }, "client_id": { "type": "keyword", "ignore_above": 100 }, "application": { "type": "keyword", "ignore_above": 100 } diff --git a/src/test/java/org/opensearch/ubi/UbIParametersTests.java b/src/test/java/org/opensearch/ubi/UbIParametersTests.java index 47e6fbe..333e871 100644 --- a/src/test/java/org/opensearch/ubi/UbIParametersTests.java +++ b/src/test/java/org/opensearch/ubi/UbIParametersTests.java @@ -10,6 +10,7 @@ import org.opensearch.action.search.SearchRequest; import org.opensearch.core.common.io.stream.StreamOutput; +import org.opensearch.common.io.stream.BytesStreamOutput; import org.opensearch.core.xcontent.XContent; import org.opensearch.core.xcontent.XContentBuilder; import org.opensearch.core.xcontent.XContentGenerator; @@ -91,14 +92,14 @@ public void testUbiParameters() { public void testWriteTo() throws IOException { final UbiParameters params = new UbiParameters("query_id", "user_query", "client_id", "app", "object_id", Collections.emptyMap()); - StreamOutput output = new DummyStreamOutput(); - params.writeTo(output); - List actual = ((DummyStreamOutput) output).getList(); - assertEquals("query_id", actual.get(0)); - assertEquals("user_query", actual.get(1)); - assertEquals("client_id", actual.get(2)); - assertEquals("app", actual.get(3)); - assertEquals("object_id", actual.get(4)); + BytesStreamOutput bso = new BytesStreamOutput(); + params.writeTo(bso); + UbiParameters roundTrip = new UbiParameters(bso.bytes().streamInput()); + assertEquals("query_id", roundTrip.getQueryId()); + assertEquals("user_query", roundTrip.getUserQueryText()); + assertEquals("client_id", roundTrip.getClientId()); + assertEquals("app", roundTrip.getApplication()); + assertEquals("object_id", roundTrip.getObjectIdField()); } public void testToXContent() throws IOException { diff --git a/src/test/java/org/opensearch/ubi/UbiParametersExtBuilderTests.java b/src/test/java/org/opensearch/ubi/UbiParametersExtBuilderTests.java index 8626f59..36b585b 100644 --- a/src/test/java/org/opensearch/ubi/UbiParametersExtBuilderTests.java +++ b/src/test/java/org/opensearch/ubi/UbiParametersExtBuilderTests.java @@ -41,11 +41,13 @@ public void testCtor() { } public void testParse() throws IOException { - XContentParser xcParser = mock(XContentParser.class); - when(xcParser.nextToken()).thenReturn(XContentParser.Token.START_OBJECT).thenReturn(XContentParser.Token.END_OBJECT); - UbiParametersExtBuilder builder = UbiParametersExtBuilder.parse(xcParser); - assertNotNull(builder); - assertNotNull(builder.getParams()); + XContentType xContentType = randomFrom(XContentType.values()); + org.opensearch.core.xcontent.XContentBuilder builder = org.opensearch.core.xcontent.XContentBuilder.builder(xContentType.xContent()); + builder.startObject().endObject(); + XContentParser parser = createParser(xContentType.xContent(), BytesReference.bytes(builder)); + UbiParametersExtBuilder ext = UbiParametersExtBuilder.parse(parser); + assertNotNull(ext); + assertNotNull(ext.getParams()); } public void testXContentRoundTrip() throws IOException { @@ -59,7 +61,7 @@ public void testXContentRoundTrip() throws IOException { assertEquals(extBuilder, deserialized); UbiParameters parameters = deserialized.getParams(); assertEquals("query_id", parameters.getQueryId()); - assertEquals("user_query", parameters.getUserQuery()); + assertEquals("user_query", parameters.getUserQueryText()); assertEquals("client_id", parameters.getClientId()); assertEquals("app", parameters.getApplication()); assertEquals("object_id_field", parameters.getObjectIdField()); @@ -86,7 +88,7 @@ public void testStreamRoundTrip() throws IOException { assertEquals(extBuilder, deserialized); UbiParameters parameters = deserialized.getParams(); assertEquals("query_id", parameters.getQueryId()); - assertEquals("user_query", parameters.getUserQuery()); + assertEquals("user_query", parameters.getUserQueryText()); assertEquals("client_id", parameters.getClientId()); assertEquals("app", parameters.getApplication()); assertEquals("object_id_field", parameters.getObjectIdField()); diff --git a/src/yamlRestTest/resources/rest-api-spec/test/_plugins.ubi/20_queries_with_ubi.yml b/src/yamlRestTest/resources/rest-api-spec/test/_plugins.ubi/20_queries_with_ubi.yml index 8e41336..13e79f4 100644 --- a/src/yamlRestTest/resources/rest-api-spec/test/_plugins.ubi/20_queries_with_ubi.yml +++ b/src/yamlRestTest/resources/rest-api-spec/test/_plugins.ubi/20_queries_with_ubi.yml @@ -41,7 +41,7 @@ search: rest_total_hits_as_int: true index: ecommerce - body: "{\"query\": {\"match\": {\"category\": \"notebook\"}}, \"ext\": {\"ubi\": {\"query_id\": \"wertwert\", \"client_id\": \"abcabc\", \"user_query\": \"notebook\", \"query_attributes\": {\"experiment\": \"number_1\"}}}}" + body: "{\"query\": {\"match\": {\"category\": \"notebook\"}}, \"ext\": {\"ubi\": {\"query_id\": \"wertwert\", \"client_id\": \"abcabc\", \"user_query\": {\"text\": \"notebook\"}, \"query_attributes\": {\"experiment\": \"number_1\"}}}}" - gte: { hits.total: 1 } @@ -49,7 +49,7 @@ search: rest_total_hits_as_int: true index: ecommerce - body: "{\"query\": {\"match\": {\"category\": \"notebook\"}}, \"ext\": {\"ubi\": {\"query_id\": \"1234512345\", \"client_id\": \"abcabc\", \"user_query\": \"notebook\"}}}" + body: "{\"query\": {\"match\": {\"category\": \"notebook\"}}, \"ext\": {\"ubi\": {\"query_id\": \"1234512345\", \"client_id\": \"abcabc\", \"user_query\": {\"text\": \"notebook\"}}}}" - gte: { hits.total: 1 } @@ -57,7 +57,7 @@ search: rest_total_hits_as_int: true index: ecommerce - body: "{\"query\": {\"match\": {\"category\": \"notebook\"}}, \"ext\": {\"ubi\": {\"query_id\": \"abcdef\", \"client_id\": \"abcabc\", \"user_query\": \"notebook\", \"application\": \"app1\"}}}" + body: "{\"query\": {\"match\": {\"category\": \"notebook\"}}, \"ext\": {\"ubi\": {\"query_id\": \"abcdef\", \"client_id\": \"abcabc\", \"user_query\": {\"text\": \"notebook\"}, \"application\": \"app1\"}}}" - gte: { hits.total: 1 } diff --git a/ubi-data-generator/ubi_data_generator.py b/ubi-data-generator/ubi_data_generator.py index f50cff2..7b00ff4 100644 --- a/ubi-data-generator/ubi_data_generator.py +++ b/ubi-data-generator/ubi_data_generator.py @@ -310,13 +310,16 @@ def convert_to_ndjson(gen_config, queries, events): return data +def wrap_user_query(uq): + return {"text": uq} if isinstance(uq, str) else uq + def make_query_event(gen_config, row): response_id = str(uuid.uuid4()) query_event = { "application": gen_config.application, "query_id": row["query_id"], "client_id": row["client_id"], - "user_query": row["user_query"], + "user_query": wrap_user_query(row["user_query"]), "query_attributes": {}, "timestamp": row["timestamp"], } @@ -331,7 +334,7 @@ def make_ubi_event(gen_config, row): "session_id": row["session_id"], "client_id": row["client_id"], "timestamp": row["timestamp"], - "user_query": row["user_query"], + "user_query": wrap_user_query(row["user_query"]), "message_type": None, "message": None, "event_attributes": {