add pagination to coloc, filter study types with opensearch

formatting
opentargets · Oct 9, 2024 · fdbec7b · fdbec7b
1 parent 7fab2c9
commit fdbec7b
Show file tree

Hide file tree

Showing 4 changed files with 45 additions and 21 deletions.
diff --git 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
@@ -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
@@ -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
@@ -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))
       }
     )
   )