one more useful method

Author: johanandren

akka-persistence-query/src/main/scala/akka/persistence/query/QueryCorrelationId.scala

@@ -16,8 +16,6 @@ object QueryCorrelationId {
   /**
    * Expected to be used "around" calls to plugin query method, will clear the correlation id from thread local
    * to make sure there is no leak between logic executed on shared threads.
-   *
-   * @param correlationId
    */
   def withCorrelationId[T](correlationId: String)(block: () => T): T = {
     threadLocal.set(correlationId)
@@ -28,6 +26,17 @@ object QueryCorrelationId {
     }
   }
 
+  /**
+   * Expected to be used "around" calls to plugin query method to pass along a prevously extracted optional correlation id,
+   * will clear the correlation id from thread local to make sure there is no leak between logic executed on shared threads.
+   */
+  def withCorrelationId[T](correlationId: Option[String])(block: () => T): T = {
+    correlationId match {
+      case None           => block()
+      case Some(actualId) => withCorrelationId(actualId)(block)
+    }
+  }
+
   /**
    * @return Expected to be called directly after receiving a query call, before starting any asynchronous tasks,
    *         returns and clears out the correlation id to make sure there is no leak between tasks. Further passing

akka-persistence-query/src/test/scala/akka/persistence/query/QueryCorrelationIdSpec.scala

@@ -29,6 +29,26 @@ class QueryCorrelationIdSpec extends AnyWordSpecLike with Matchers {
       QueryCorrelationId.get() shouldBe None
     }
 
+    "pass along and clear correlation id if present" in {
+      val uuid = UUID.randomUUID().toString
+      val observed =
+        QueryCorrelationId.withCorrelationId(Some(uuid)) { () =>
+          pretendQueryMethod()
+        }
+      observed shouldEqual Some(uuid)
+
+      // cleared after returning
+      QueryCorrelationId.get() shouldBe None
+    }
+
+    "just invoke the block if correlation id not present" in {
+      val observed =
+        QueryCorrelationId.withCorrelationId(None) { () =>
+          pretendQueryMethod()
+        }
+      observed shouldEqual None
+    }
+
     "clear correlation id when call fails" in {
       val uuid = UUID.randomUUID().toString
       intercept[TestException] {