Merge pull request #175 from o19s/2.14.0-query_attributes

Adding query_attributes to ubi search request
o19s · May 28, 2024 · ba0ad7c · ba0ad7c
2 parents 1b94080 + 9db8a08
commit ba0ad7c
Show file tree

Hide file tree

Showing 7 changed files with 79 additions and 47 deletions.
diff --git a/docker-compose.yaml b/docker-compose.yaml
@@ -34,20 +34,20 @@ services:
     networks:
       - ubi-dev-os-net
 
-  ubi-dev-os-dashboards:
-    image: opensearchproject/opensearch-dashboards:2.12.0
-    container_name: ubi-dev-os-dashboards
-    ports:
-      - 5601:5601
-    expose:
-      - 5601
-    environment:
-      OPENSEARCH_HOSTS: '["http://ubi-dev-os:9200"]'
-      DISABLE_SECURITY_DASHBOARDS_PLUGIN: "true"
-    depends_on:
-      - ubi-dev-os
-    networks:
-      - ubi-dev-os-net      
+#  ubi-dev-os-dashboards:
+#    image: opensearchproject/opensearch-dashboards:2.12.0
+#    container_name: ubi-dev-os-dashboards
+#    ports:
+#      - 5601:5601
+#    expose:
+#      - 5601
+#    environment:
+#      OPENSEARCH_HOSTS: '["http://ubi-dev-os:9200"]'
+#      DISABLE_SECURITY_DASHBOARDS_PLUGIN: "true"
+#    depends_on:
+#      - ubi-dev-os
+#    networks:
+#      - ubi-dev-os-net
 
 networks:
   ubi-dev-os-net:

diff --git a/search.sh b/search.sh
@@ -5,7 +5,11 @@ curl -s http://localhost:9200/ecommerce/_search -H "Content-Type: application/js
   "ext": {
    "ubi": {
      "query_id": "12300d16cb-b6f1-4012-93ebcc49cac90426",
-     "user_query": "toner"
+     "user_query": "toner",
+     "query_attributes": {
+       "experiment": "exp_1",
+       "system": "my_system"
+     }
     }
    },
    "query": {

diff --git a/src/main/java/org/opensearch/ubi/QueryRequest.java b/src/main/java/org/opensearch/ubi/QueryRequest.java
@@ -8,6 +8,8 @@
 
 package org.opensearch.ubi;
 
+import java.util.Map;
+
 /**
  * A query received by OpenSearch.
  */
@@ -18,6 +20,7 @@ public class QueryRequest {
     private final String userId;
     private final String userQuery;
     private final String query;
+    private final Map<String, String> queryAttributes;
     private final QueryResponse queryResponse;
 
     /**
@@ -26,17 +29,28 @@ public class QueryRequest {
      * @param userQuery The user-entered query.
      * @param userId The ID of the user that initiated the query.
      * @param query The raw query.
+     * @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 userId, final String query, final QueryResponse queryResponse) {
+    public QueryRequest(final String queryId, final String userQuery, final String userId, final String query,
+                        final Map<String, String> queryAttributes, final QueryResponse queryResponse) {
         this.timestamp = System.currentTimeMillis();
         this.queryId = queryId;
         this.userId = userId;
         this.userQuery = userQuery;
         this.query = query;
+        this.queryAttributes = queryAttributes;
         this.queryResponse = queryResponse;
     }
 
+    /**
+     * Gets the query attributes.
+     * @return The query attributes.
+     */
+    public Map<String, String> getQueryAttributes() {
+        return queryAttributes;
+    }
+
     /**
      * Gets the timestamp.
      * @return The timestamp.

diff --git a/src/main/java/org/opensearch/ubi/UbiActionFilter.java b/src/main/java/org/opensearch/ubi/UbiActionFilter.java
@@ -101,7 +101,7 @@ public void onResponse(Response response) {
                         final String userQuery = ubiParameters.getUserQuery();
                         final String userId = ubiParameters.getClientId();
                         final String objectId = ubiParameters.getObjectId();
-
+                        final Map<String, String> queryAttributes = ubiParameters.getQueryAttributes();
                         final String query = searchRequest.source().toString();
 
                         final List<String> queryResponseHitIds = new LinkedList<>();
@@ -120,7 +120,7 @@ public void onResponse(Response response) {
 
                         final String queryResponseId = UUID.randomUUID().toString();
                         final QueryResponse queryResponse = new QueryResponse(queryId, queryResponseId, queryResponseHitIds);
-                        final QueryRequest queryRequest = new QueryRequest(queryId, userQuery, userId, query, queryResponse);
+                        final QueryRequest queryRequest = new QueryRequest(queryId, userQuery, userId, query, queryAttributes, queryResponse);
 
                         SearchResponse searchResponse = (SearchResponse) response;
 
@@ -202,6 +202,7 @@ public void onResponse(IndicesExistsResponse indicesExistsResponse) {
                 source.put("query_response_object_ids", queryRequest.getQueryResponse().getQueryResponseObjectIds());
                 source.put("user_id", queryRequest.getUserId());
                 source.put("user_query", queryRequest.getUserQuery());
+                source.put("query_attributes", queryRequest.getQueryAttributes());
 
                 // The query can be null for some types of queries.
                 if(queryRequest.getQuery() != null) {

diff --git a/src/main/java/org/opensearch/ubi/ext/UbiParameters.java b/src/main/java/org/opensearch/ubi/ext/UbiParameters.java
@@ -13,6 +13,7 @@
 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;
@@ -35,23 +36,17 @@ public class UbiParameters implements Writeable, ToXContentObject {
     private static final ParseField USER_QUERY = new ParseField("user_query");
     private static final ParseField CLIENT_ID = new ParseField("client_id");
     private static final ParseField OBJECT_ID = new ParseField("object_id");
-    //private static final ParseField ATTRIBUTES = new ParseField("attributes");
+    private static final ParseField QUERY_ATTRIBUTES = new ParseField("query_attributes");
 
     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::setObjectId, OBJECT_ID);
-        //PARSER.declareNamedObject(UbiParameters::setAttributes, (p, c, n) -> p.namedObject(Map.class, n, c), ATTRIBUTES);
+        PARSER.declareObject(UbiParameters::setQueryAttributes, (p, c) -> p.mapStrings(), QUERY_ATTRIBUTES);
     }
 
-    private String queryId;
-    private String userQuery;
-    private String clientId;
-    private String objectId;
-    //private Map<String, Object> attributes;
-
     /**
      * Get the {@link UbiParameters} from a {@link SearchRequest}.
      * @param request A {@link SearchRequest},
@@ -63,10 +58,10 @@ public static UbiParameters getUbiParameters(final SearchRequest request) {
 
         if (request.source() != null && request.source().ext() != null && !request.source().ext().isEmpty()) {
             final Optional<SearchExtBuilder> b = request.source()
-                .ext()
-                .stream()
-                .filter(bldr -> UbiParametersExtBuilder.UBI_PARAMETER_NAME.equals(bldr.getWriteableName()))
-                .findFirst();
+                    .ext()
+                    .stream()
+                    .filter(bldr -> UbiParametersExtBuilder.UBI_PARAMETER_NAME.equals(bldr.getWriteableName()))
+                    .findFirst();
             if (b.isPresent()) {
                 builder = (UbiParametersExtBuilder) b.get();
             }
@@ -80,6 +75,12 @@ public static UbiParameters getUbiParameters(final SearchRequest request) {
 
     }
 
+    private String queryId;
+    private String userQuery;
+    private String clientId;
+    private String objectId;
+    private Map<String, String> queryAttributes;
+
     /**
      * Creates a new instance.
      */
@@ -90,12 +91,13 @@ public UbiParameters() {}
      * @param input The {@link StreamInput} to read parameters from.
      * @throws IOException Thrown if the parameters cannot be read.
      */
+    @SuppressWarnings("unchecked")
     public UbiParameters(StreamInput input) throws IOException {
         this.queryId = input.readString();
         this.userQuery = input.readOptionalString();
         this.clientId = input.readOptionalString();
         this.objectId = input.readOptionalString();
-        //this.attributes = input.readMap();
+        this.queryAttributes = (Map<String, String>) input.readGenericValue();
     }
 
     /**
@@ -104,14 +106,14 @@ public UbiParameters(StreamInput input) throws IOException {
      * @param userQuery The user-entered search query.
      * @param clientId The client ID.
      * @param objectId The object ID.
-  //   * @param attributes Optional attributes for UBI.
+     * @param queryAttributes Optional attributes for UBI.
      */
-    public UbiParameters(String queryId, String userQuery, String clientId, String objectId) {
+    public UbiParameters(String queryId, String userQuery, String clientId, String objectId, Map<String, String> queryAttributes) {
         this.queryId = queryId;
         this.userQuery = userQuery;
         this.clientId = clientId;
         this.objectId = objectId;
-        //this.attributes = attributes;
+        this.queryAttributes = queryAttributes;
     }
 
     @Override
@@ -120,8 +122,8 @@ 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(OBJECT_ID.getPreferredName(), this.objectId);
-                //.field(ATTRIBUTES.getPreferredName(), this.attributes);
+                .field(OBJECT_ID.getPreferredName(), this.objectId)
+                .field(QUERY_ATTRIBUTES.getPreferredName(), this.queryAttributes);
     }
 
     @Override
@@ -130,7 +132,7 @@ public void writeTo(StreamOutput out) throws IOException {
         out.writeOptionalString(userQuery);
         out.writeOptionalString(clientId);
         out.writeOptionalString(objectId);
-        //out.writeMap(attributes);
+        out.writeGenericValue(queryAttributes);
     }
 
     /**
@@ -152,12 +154,12 @@ public boolean equals(Object o) {
             return false;
         }
 
-        UbiParameters other = (UbiParameters) o;
+        final UbiParameters other = (UbiParameters) o;
         return Objects.equals(this.queryId, other.getQueryId())
                 && Objects.equals(this.userQuery, other.getUserQuery())
                 && Objects.equals(this.clientId, other.getClientId())
-                && Objects.equals(this.objectId, other.getObjectId());
-               // && Objects.equals(this.attributes, other.getAttributes());
+                && Objects.equals(this.objectId, other.getObjectId())
+                && Objects.equals(this.queryAttributes, other.getQueryAttributes());
     }
 
     @Override
@@ -236,16 +238,16 @@ public void setUserQuery(String userQuery) {
      * Get the attributes.
      * @return A map of attributes.
      */
-//    public Map<String, Object> getAttributes() {
-//        return attributes;
-//    }
+    public Map<String, String> getQueryAttributes() {
+        return queryAttributes;
+    }
 
     /**
      * Sets the attributes.
-     * @param attributes A map of attributes.
+     * @param queryAttributes A map of attributes.
      */
-//    public void setAttributes(Map<String, Object> attributes) {
-//        this.attributes = attributes;
-//    }
+    public void setQueryAttributes(Map<String, String> queryAttributes) {
+        this.queryAttributes = queryAttributes;
+    }
 
 }
diff --git a/src/main/resources/queries-mapping.json b/src/main/resources/queries-mapping.json
@@ -26,6 +26,9 @@
           "type": "keyword"
         }
       }
+    },
+    "query_attributes": {
+      "type": "flat_object"
     }
   }
 }
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
@@ -23,6 +23,14 @@
       indices.refresh:
         index: [ "ecommerce" ]
 
+  - 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\", \"query_attributes\": {\"experiment\": \"number_1\"}}}}"
+
+  - gte: { hits.total: 1 }
+
   - do:
       search:
         rest_total_hits_as_int: true