gov.nasa.pds.api.registry.search.HitIterator Maven / Gradle / Ivy

Go to download

Show more of this group Show more artifacts with this name
Show all versions of registry-api-service Show documentation

Registry API Service contributing to the PDS Federated Search API

There is a newer version: 4.5.6

package gov.nasa.pds.api.registry.search;

import java.io.IOException;
import java.util.Iterator;
import java.util.Map;

import org.opensearch.action.search.SearchRequest;
import org.opensearch.client.RequestOptions;
import org.opensearch.client.RestHighLevelClient;
import org.opensearch.search.SearchHit;
import org.opensearch.search.SearchHits;

public class HitIterator implements Iterable>, Iterator> {
  private int size = 500; // define size to use here to prevent page skipping if opensearch default
                          // size
                          // ever changes
  private int at = 0, page = 0;
  private SearchHits currentBatch;
  private RestHighLevelClient client;
  private SearchRequest request;

  public HitIterator(RestHighLevelClient client, SearchRequest request) throws IOException {
    super();
    this.client = client;
    this.request = request;
    this.currentBatch = this.fetch();
  }

  public HitIterator(int size, RestHighLevelClient client, SearchRequest request)
      throws IOException {
    super();
    this.client = client;
    this.request = request;
    this.size = size;
    this.currentBatch = this.fetch();
  }

  private SearchHits fetch() throws IOException {
    this.request.source().from(this.page * this.size);
    this.request.source().size(this.size);
    return this.client.search(this.request, RequestOptions.DEFAULT).getHits();
  }

  private SearchHit getAt() throws IOException {
    if (this.size <= this.at) {
      this.page++;
      this.at = 0;
      this.currentBatch = this.fetch();
    }

    return this.currentBatch.getAt(this.at);
  }

  public String getCurrentId() {
    try {
      return this.getAt().getId();
    } catch (IOException ioe) {
      return null;
    }
  }

  @Override
  public boolean hasNext() {
    return this.currentBatch == null ? false
        : (this.at + this.page * this.size) < this.currentBatch.getTotalHits().value;
  }

  @Override
  public Iterator> iterator() {
    return this;
  }

  @Override
  public Map next() {
    if (this.hasNext()) {
      try {
        SearchHit hit = this.getAt();
        at++;
        return hit.getSourceAsMap();
      } catch (IOException ioe) {
        return null;
      }
    } else {
      return null;
    }
  }
}