Change the user_query field to a flat_object

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
+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

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<String, String> userQuery;
     private final String query;
     private final String application;
     private final Map<String, String> 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<String, String> userQuery, final String clientId, final String query,
                         final String application, final Map<String, String> 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<String, String> 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", "");
     }
 
     /**

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<String, String> 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<String, String> 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());
 

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<String, String> 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<String, String>) 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<String, String> 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<String, String> userQuery, String clientId, String application, String objectIdField, Map<String, String> 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<String, String> userQuery = null;
+        String clientId = null;
+        String application = null;
+        String objectIdField = null;
+        Map<String, String> 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<String, String> 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<String, String> userQuery) {
         this.userQuery = userQuery;
     }
 
+    /**
+     * Set a legacy string user query; wraps to {"text": {@literal <value>}}.
+     * @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.

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 @@
       }
     }
   }
-}
+}

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 }

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<String> 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 {

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());

src/yamlRestTest/resources/rest-api-spec/test/_plugins.ubi/20_queries_with_ubi.yml

@@ -41,23 +41,23 @@
       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 }
 
   - do:
       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 }
 
   - do:
       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 }
 

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": {