Improved testing of ubiQuery() streams to figure out how to deal with…

… queryAttributes map
apache · Nov 18, 2024 · f108585 · f108585
1 parent e72a7ec
commit f108585
Show file tree

Hide file tree

Showing 5 changed files with 318 additions and 324 deletions.
diff --git a/solr/core/src/java/org/apache/solr/handler/LoggingStream.java b/solr/core/src/java/org/apache/solr/handler/LoggingStream.java
@@ -71,6 +71,8 @@ public class LoggingStream extends TupleStream implements Expressible {
    */
   private String filepath;
 
+  private Path filePath;
+
   private int updateBatchSize;
 
   private int batchNumber;
@@ -122,14 +124,14 @@ private void init(String filepath, TupleStream tupleSource) {
     this.tupleSource = tupleSource;
   }
 
-  /** The name of the file being updated */
-  protected String getFilePath() {
-    return filepath;
+  /** The path of the file being logged to */
+  public Path getFilePath() {
+    return filePath;
   }
 
   @Override
   public void open() throws IOException {
-    Path filePath = chroot.resolve(filepath).normalize();
+    filePath = chroot.resolve(filepath).normalize();
     if (!filePath.startsWith(chroot)) {
       throw new SolrException(
           SolrException.ErrorCode.BAD_REQUEST, "file to log to must be under " + chroot);

diff --git a/solr/core/src/java/org/apache/solr/handler/component/UBIComponent.java b/solr/core/src/java/org/apache/solr/handler/component/UBIComponent.java
@@ -230,7 +230,7 @@ public void process(ResponseBuilder rb) throws IOException {
     ubiQuery.setUserQuery(params.get(USER_QUERY));
     ubiQuery.setApplication(params.get(APPLICATION));
 
-    Object queryAttributes = params.get(QUERY_ATTRIBUTES);
+    String queryAttributes = params.get(QUERY_ATTRIBUTES);
 
     if (queryAttributes != null && queryAttributes.toString().startsWith("{")) {
       // Look up the original nested JSON format, typically passed in
@@ -241,8 +241,9 @@ public void process(ResponseBuilder rb) throws IOException {
         @SuppressWarnings("rawtypes")
         Map paramsProperties = (Map) jsonProperties.get("params");
         if (paramsProperties.containsKey(QUERY_ATTRIBUTES)) {
-          queryAttributes = paramsProperties.get(QUERY_ATTRIBUTES);
-          ubiQuery.setQueryAttributes(queryAttributes);
+          @SuppressWarnings("rawtypes")
+          Map queryAttributesAsMap = (Map) paramsProperties.get(QUERY_ATTRIBUTES);
+          ubiQuery.setQueryAttributes(queryAttributesAsMap);
         }
       }
     }

diff --git a/solr/core/src/java/org/apache/solr/handler/component/UBIQuery.java b/solr/core/src/java/org/apache/solr/handler/component/UBIQuery.java
@@ -34,7 +34,10 @@ public class UBIQuery {
   private String application;
   private String queryId;
   private String userQuery;
-  private Object queryAttributes;
+
+  @SuppressWarnings("rawtypes")
+  private Map queryAttributes;
+
   private String docIds;
 
   public UBIQuery(String queryId) {
@@ -69,11 +72,13 @@ public void setUserQuery(String userQuery) {
     this.userQuery = userQuery;
   }
 
-  public Object getQueryAttributes() {
+  @SuppressWarnings("rawtypes")
+  public Map getQueryAttributes() {
     return queryAttributes;
   }
 
-  public void setQueryAttributes(Object queryAttributes) {
+  @SuppressWarnings("rawtypes")
+  public void setQueryAttributes(Map queryAttributes) {
     this.queryAttributes = queryAttributes;
   }
 
@@ -90,9 +95,14 @@ public Map toMap() {
     @SuppressWarnings({"rawtypes", "unchecked"})
     Map map = new HashMap();
     map.put(UBIComponent.QUERY_ID, this.queryId);
-    map.put(UBIComponent.APPLICATION, this.application);
-    map.put(UBIComponent.USER_QUERY, this.userQuery);
+    if (this.application != null) {
+      map.put(UBIComponent.APPLICATION, this.application);
+    }
+    if (this.userQuery != null) {
+      map.put(UBIComponent.USER_QUERY, this.userQuery);
+    }
     if (this.queryAttributes != null) {
+
       ObjectMapper objectMapper = new ObjectMapper();
       try {
         map.put(

diff --git a/solr/core/src/test/org/apache/solr/handler/component/UBIComponentRecordingTest.java b/solr/core/src/test/org/apache/solr/handler/component/UBIComponentRecordingTest.java
@@ -91,6 +91,8 @@ public void testRecordingUBIQueries() throws Exception {
     final QueryResponse queryResponse = cluster.getSolrClient().query(COLLECTION, queryParams);
     final SolrDocumentList documents = queryResponse.getResults();
 
+
+
     //
     //
     //    final ModifiableSolrParams overrideParams = new ModifiableSolrParams();
@@ -102,6 +104,6 @@ public void testRecordingUBIQueries() throws Exception {
     //            .setLimit(0);
     //    QueryResponse queryResponse = req.process(cluster.getSolrClient(), COLLECTION);
     // assertResponseFoundNumDocs(queryResponse, expectedResults);
-    System.out.println(queryResponse);
+    //System.out.println(queryResponse);
   }
 }