From fdbec7b6a9c8e0e090f07bf522c504ff263ec872 Mon Sep 17 00:00:00 2001 From: James Hayhurst Date: Wed, 9 Oct 2024 11:00:48 +0000 Subject: [PATCH] add pagination to coloc, filter study types with opensearch formatting --- app/models/Backend.scala | 27 ++++++++---------- app/models/ElasticRetriever.scala | 28 +++++++++++++++++-- .../ElasticRetrieverQueryBuilders.scala | 5 ++++ app/models/entities/CredibleSet.scala | 6 ++-- 4 files changed, 45 insertions(+), 21 deletions(-) diff --git a/app/models/Backend.scala b/app/models/Backend.scala index a421848a..02ac2a69 100644 --- a/app/models/Backend.scala +++ b/app/models/Backend.scala @@ -192,26 +192,21 @@ class Backend @Inject() (implicit } def getColocalisation(studyLocusId: String, - studyTypes: Option[Seq[StudyTypeEnum.Value]] + studyTypes: Option[Seq[StudyTypeEnum.Value]], + pagination: Option[Pagination] ): Future[IndexedSeq[Colocalisation]] = { val indexName = getIndexOrDefault("colocalisation") + val pag = pagination.getOrElse(Pagination.mkDefault) + val terms = Map("leftStudyLocusId.keyword" -> Seq(studyLocusId), + "rightStudyLocusId.keyword" -> Seq(studyLocusId) + ).filter(_._2.nonEmpty) + val filter = Seq( + termsQuery("rightStudyType.keyword", studyTypes.getOrElse(StudyTypeEnum.values)) + ) val colocs = esRetriever - .getByIndexedQueryShould(indexName, - Map("leftStudyLocusId.keyword" -> studyLocusId, - "rightStudyLocusId.keyword" -> studyLocusId - ), - Pagination.mkDefault, - fromJsValue[Colocalisation] - ) + .getByIndexedTermsShould(indexName, terms, pag, fromJsValue[Colocalisation], filter = filter) .map(_._1) - val colocsFilteredByStudyType = if (studyTypes.isDefined) { - colocs.map( - _.filter(coloc => studyTypes.get.contains(StudyTypeEnum.withName(coloc.rightStudyType))) - ) - } else { - colocs - } - colocsFilteredByStudyType.map(_.map { coloc => + colocs.map(_.map { coloc => val otherStudyLocusId: String = if (coloc.leftStudyLocusId != studyLocusId) { coloc.leftStudyLocusId } else { diff --git a/app/models/ElasticRetriever.scala b/app/models/ElasticRetriever.scala index 3c904aec..c1e83adf 100644 --- a/app/models/ElasticRetriever.scala +++ b/app/models/ElasticRetriever.scala @@ -169,11 +169,13 @@ class ElasticRetriever @Inject() ( buildF: JsValue => Option[A], aggs: Iterable[AbstractAggregation] = Iterable.empty, sortByField: Option[sort.FieldSort] = None, - excludedFields: Seq[String] = Seq.empty + excludedFields: Seq[String] = Seq.empty, + filter: Seq[Query] = Seq.empty ): Future[(IndexedSeq[A], JsValue)] = { // just log and execute the query val indexQuery: IndexQuery[V] = IndexQuery(esIndex = esIndex, kv = kv, + filters = filter, pagination = pagination, aggs = aggs, excludedFields = excludedFields @@ -182,6 +184,28 @@ class ElasticRetriever @Inject() ( getByIndexedQuery(searchRequest, sortByField, buildF) } + def getByIndexedTermsShould[A, V]( + esIndex: String, + kv: Map[String, V], + pagination: Pagination, + buildF: JsValue => Option[A], + aggs: Iterable[AbstractAggregation] = Iterable.empty, + sortByField: Option[sort.FieldSort] = None, + excludedFields: Seq[String] = Seq.empty, + filter: Seq[Query] = Seq.empty + ): Future[(IndexedSeq[A], JsValue)] = { + // just log and execute the query + val indexQuery: IndexQuery[V] = IndexQuery(esIndex = esIndex, + kv = kv, + filters = filter, + pagination = pagination, + aggs = aggs, + excludedFields = excludedFields + ) + val searchRequest: SearchRequest = IndexTermsShould(indexQuery) + getByIndexedQuery(searchRequest, sortByField, buildF) + } + /** This fn represents a query where each kv from the map is used in * a bool 'should'. Based on the query asked by `getByIndexedQuery` and aggregation is applied */ @@ -229,7 +253,7 @@ class ElasticRetriever @Inject() ( case None => searchRequest } - logger.debug(s"Elasticsearch query: ${client.show(sortedSearchRequest)}") + logger.info(s"Elasticsearch query: ${client.show(sortedSearchRequest)}") sortedSearchRequest } diff --git a/app/models/ElasticRetrieverQueryBuilders.scala b/app/models/ElasticRetrieverQueryBuilders.scala index 4e500668..b29c16c3 100644 --- a/app/models/ElasticRetrieverQueryBuilders.scala +++ b/app/models/ElasticRetrieverQueryBuilders.scala @@ -42,6 +42,11 @@ trait ElasticRetrieverQueryBuilders extends QueryApi with Logging { ): SearchRequest = getByIndexTermsBuilder(indexQuery, must) + def IndexTermsShould[V]( + indexQuery: IndexQuery[V] + ): SearchRequest = + getByIndexTermsBuilder(indexQuery, should) + def getByIndexQueryBuilder[V]( indexQuery: IndexQuery[V], f: Iterable[Query] => BoolQuery diff --git a/app/models/entities/CredibleSet.scala b/app/models/entities/CredibleSet.scala index c24e2820..fb9bd503 100644 --- a/app/models/entities/CredibleSet.scala +++ b/app/models/entities/CredibleSet.scala @@ -18,7 +18,7 @@ import sangria.schema.{ StringType, fields } -import models.gql.Arguments.studyTypes +import models.gql.Arguments.{studyTypes, pageArg} case class Locus( variantId: Option[String], @@ -279,10 +279,10 @@ object CredibleSet extends Logging { "colocalisation", OptionType(ListType(colocalisationImp)), description = None, - arguments = studyTypes :: Nil, + arguments = studyTypes :: pageArg :: Nil, resolve = js => { val id = (js.value \ "studyLocusId").as[String] - js.ctx.getColocalisation(id, js.arg(studyTypes)) + js.ctx.getColocalisation(id, js.arg(studyTypes), js.arg(pageArg)) } ) )