Make it configurable for OpenLineageService async executor. (MarquezP…

…roject#1866) Signed-off-by: Minkyu Park <[email protected]>
renwayle · Feb 7, 2022 · a1106a0 · a1106a0
1 parent d31149d
commit a1106a0
Showing 1 changed file with 10 additions and 3 deletions.
diff --git a/api/src/main/java/marquez/service/OpenLineageService.java b/api/src/main/java/marquez/service/OpenLineageService.java
@@ -14,6 +14,7 @@
 import java.util.Optional;
 import java.util.UUID;
 import java.util.concurrent.CompletableFuture;
+import java.util.concurrent.Executor;
 import java.util.concurrent.ForkJoinPool;
 import java.util.stream.Collectors;
 import lombok.extern.slf4j.Slf4j;
@@ -44,17 +45,23 @@ public class OpenLineageService extends DelegatingDaos.DelegatingOpenLineageDao
   private final DatasetVersionDao datasetVersionDao;
   private final ObjectMapper mapper = Utils.newObjectMapper();
 
+  private final Executor executor;
+
   public OpenLineageService(BaseDao baseDao, RunService runService) {
+    this(baseDao, runService, ForkJoinPool.commonPool());
+  }
+
+  public OpenLineageService(BaseDao baseDao, RunService runService, Executor executor) {
     super(baseDao.createOpenLineageDao());
     this.runService = runService;
     this.datasetVersionDao = baseDao.createDatasetVersionDao();
+    this.executor = executor;
   }
 
   public CompletableFuture<Void> createAsync(LineageEvent event) {
     CompletableFuture<Void> marquez =
         CompletableFuture.supplyAsync(
-                withSentry(withMdc(() -> updateMarquezModel(event, mapper))),
-                ForkJoinPool.commonPool())
+                withSentry(withMdc(() -> updateMarquezModel(event, mapper))), executor)
             .thenAccept(
                 (update) -> {
                   if (event.getEventType() != null) {
@@ -79,7 +86,7 @@ public CompletableFuture<Void> createAsync(LineageEvent event) {
                             event.getJob().getNamespace(),
                             createJsonArray(event, mapper),
                             event.getProducer()))),
-            ForkJoinPool.commonPool());
+            executor);
 
     return CompletableFuture.allOf(marquez, openLineage);
   }