Use SearchIterator instead of ElasticSource and fromGraph (#2842)

* Use SearchIterator instead of ElasticSource and fromGraph

* a little clarity

* typo

Author: paul-butcher

pipeline/relation_embedder/relation_embedder/docker-compose.yml

@@ -9,7 +9,7 @@ services:
     ports:
       - "4566:4566"
   elasticsearch:
-    image: "docker.elastic.co/elasticsearch/elasticsearch:8.4.0"
+    image: "docker.elastic.co/elasticsearch/elasticsearch:8.8.2"
     ports:
       - "9200:9200"
       - "9300:9300"

pipeline/relation_embedder/relation_embedder/scripts/batches.txt

@@ -1,2 +1,2 @@
 {"rootPath":"GC103","selectors":[{"path":"GC103","type":"Tree"}]}
-
+{"rootPath":"MS7826","selectors":[{"path":"MS7826","type":"Tree"}]}

pipeline/relation_embedder/relation_embedder/src/main/scala/weco/pipeline/relation_embedder/BatchProcessor.scala

@@ -3,7 +3,6 @@ package weco.pipeline.relation_embedder
 import com.sksamuel.elastic4s.Index
 import grizzled.slf4j.Logging
 import org.apache.pekko.NotUsed
-import org.apache.pekko.actor.ActorSystem
 import org.apache.pekko.stream.Materializer
 import org.apache.pekko.stream.scaladsl.{Sink, Source}
 import weco.catalogue.internal_model.work.Work
@@ -79,6 +78,7 @@ class BatchProcessor(
       .getAffectedWorks(batch)
       .map {
         work =>
+          debug(s"transitioning ${work.id}")
           val relations = relationsCache(work)
           work.transition[Denormalised](relations)
       }
@@ -105,8 +105,7 @@ object BatchProcessor {
   def apply(
     config: RelationEmbedderConfig
   )(
-    implicit actorSystem: ActorSystem,
-    ec: ExecutionContext,
+    implicit ec: ExecutionContext,
     materializer: Materializer
   ): BatchProcessor = {
 

pipeline/relation_embedder/relation_embedder/src/main/scala/weco/pipeline/relation_embedder/RelationsService.scala

@@ -1,19 +1,20 @@
 package weco.pipeline.relation_embedder
 
 import org.apache.pekko.NotUsed
-import org.apache.pekko.actor.ActorSystem
 import org.apache.pekko.stream.scaladsl.Source
 import com.sksamuel.elastic4s.ElasticDsl._
 import com.sksamuel.elastic4s.circe._
 import com.sksamuel.elastic4s.{ElasticClient, Index}
 import grizzled.slf4j.Logging
 import weco.json.JsonUtil._
 import weco.catalogue.internal_model.Implicits._
-import com.sksamuel.elastic4s.pekko.streams._
+import com.sksamuel.elastic4s.requests.searches.SearchIterator
 import weco.catalogue.internal_model.work.WorkState.Merged
 import weco.catalogue.internal_model.work.Work
 import weco.pipeline.relation_embedder.models.{Batch, RelationWork}
 
+import scala.concurrent.duration.{Duration, DurationInt}
+
 trait RelationsService {
   def getRelationTree(batch: Batch): Source[RelationWork, NotUsed]
 
@@ -25,8 +26,7 @@ class PathQueryRelationsService(
   index: Index,
   completeTreeScroll: Int = 1000,
   affectedWorksScroll: Int = 250
-)(implicit as: ActorSystem)
-    extends RelationsService
+) extends RelationsService
     with Logging {
 
   private val requestBuilder = RelationsRequestBuilder(index)
@@ -36,34 +36,33 @@ class PathQueryRelationsService(
     debug(
       s"Querying affected works with ES request: ${elasticClient.show(request)}"
     )
-    val sourceSettings = SourceSettings(
-      search = request,
-      maxItems = Long.MaxValue,
-      fetchThreshold = affectedWorksScroll,
-      warm = true
+    // Arbitrary timeout value, it has to exist for SearchIterator,
+    // but it has not been derived either through experimentation or calculation,
+    implicit val timeout: Duration = 5 minutes
+
+    Source.fromIterator(
+      () =>
+        SearchIterator.iterate[Work[Merged]](
+          elasticClient,
+          request
+        )
     )
-    Source
-      .fromGraph(
-        new ElasticSource(elasticClient, sourceSettings)(as.dispatcher)
-      )
-      .map(searchHit => searchHit.safeTo[Work[Merged]].get)
   }
 
   def getRelationTree(batch: Batch): Source[RelationWork, NotUsed] = {
     val request = requestBuilder.completeTree(batch, completeTreeScroll)
     debug(
       s"Querying complete tree with ES request: ${elasticClient.show(request)}"
     )
-    val sourceSettings = SourceSettings(
-      search = request,
-      maxItems = Long.MaxValue,
-      fetchThreshold = affectedWorksScroll,
-      warm = true
+
+    implicit val timeout: Duration = 5 minutes
+
+    Source.fromIterator(
+      () =>
+        SearchIterator.iterate[RelationWork](
+          elasticClient,
+          request
+        )
     )
-    Source
-      .fromGraph(
-        new ElasticSource(elasticClient, sourceSettings)(as.dispatcher)
-      )
-      .map(searchHit => searchHit.safeTo[RelationWork].get)
   }
 }

pipeline/relation_embedder/relation_embedder/src/test/scala/weco/pipeline/relation_embedder/BatchProcessorTest.scala

@@ -56,33 +56,30 @@ class BatchProcessorTest
     )
     withUpstreamIndex(workList) {
       mergedIndex =>
-        withActorSystem {
-          implicit actorSystem =>
-            val relationsService = new PathQueryRelationsService(
-              elasticClient,
-              mergedIndex,
-              10
+        val relationsService = new PathQueryRelationsService(
+          elasticClient,
+          mergedIndex,
+          10
+        )
+        withMaterializer {
+          implicit materializer: Materializer =>
+            val downstream = new MemoryDownstream
+            val processor = new BatchProcessor(
+              relationsService = relationsService,
+              bulkWriter = bulkWriter,
+              downstream = downstream
             )
-            withMaterializer {
-              implicit materializer: Materializer =>
-                val downstream = new MemoryDownstream
-                val processor = new BatchProcessor(
-                  relationsService = relationsService,
-                  bulkWriter = bulkWriter,
-                  downstream = downstream
-                )
-
-                whenReady(processor(batch)) {
-                  _ =>
-                    testWith(
-                      (
-                        downstream.msgSender.messages.map(_.body),
-                        denormalisedIndex
-                      )
-                    )
-                }
 
+            whenReady(processor(batch)) {
+              _ =>
+                testWith(
+                  (
+                    downstream.msgSender.messages.map(_.body),
+                    denormalisedIndex
+                  )
+                )
             }
+
         }
     }
   }

pipeline/relation_embedder/relation_embedder/src/test/scala/weco/pipeline/relation_embedder/RelationsServiceTest.scala

@@ -25,7 +25,7 @@ class RelationsServiceTest
     index: Index,
     completeTreeScroll: Int = 20,
     affectedWorksScroll: Int = 20
-  )(implicit as: ActorSystem) =
+  ) =
     new PathQueryRelationsService(
       elasticClient = elasticClient,
       index = index,