Remove need to query JobRow on run completion

Signed-off-by: Michael Collado <collado.mike@gmail.com>
MarquezProject · Jul 20, 2022 · 5c88d5b · 5c88d5b
1 parent 7e0f47a
commit 5c88d5b
Show file tree

Hide file tree

Showing 5 changed files with 13 additions and 29 deletions.
diff --git a/api/src/main/java/marquez/db/JobVersionDao.java b/api/src/main/java/marquez/db/JobVersionDao.java
@@ -286,22 +286,19 @@ default List<UUID> findOutputDatasetsFor(UUID jobVersionUuid) {
    * code location, and context. A version for a given job is created <i>only</i> when a {@link Run}
    * transitions into a {@code COMPLETED}, {@code ABORTED}, or {@code FAILED} state.
    *
-   * @param namespaceName The namespace for the job version.
-   * @param jobName The name of the job.
+   * @param jobRow The job.
    * @param runUuid The unique ID of the run associated with the job version.
    * @param runState The current run state.
    * @param transitionedAt The timestamp of the run state transition.
    * @return A {@link BagOfJobVersionInfo} object.
    */
   default BagOfJobVersionInfo upsertJobVersionOnRunTransition(
-      @NonNull String namespaceName,
-      @NonNull String jobName,
+      @NonNull JobRow jobRow,
       @NonNull UUID runUuid,
       @NonNull RunState runState,
       @NonNull Instant transitionedAt) {
     // Get the job.
     final JobDao jobDao = createJobDao();
-    final JobRow jobRow = jobDao.findJobByNameAsRow(namespaceName, jobName).get();
 
     // Get the job context.
     final UUID jobContextUuid = jobRow.getJobContextUuid().get();

diff --git a/api/src/main/java/marquez/db/OpenLineageDao.java b/api/src/main/java/marquez/db/OpenLineageDao.java
@@ -390,8 +390,7 @@ default void updateMarquezOnComplete(
     BagOfJobVersionInfo bagOfJobVersionInfo =
         createJobVersionDao()
             .upsertJobVersionOnRunTransition(
-                updateLineageRow.getRun().getNamespaceName(),
-                updateLineageRow.getRun().getJobName(),
+                updateLineageRow.getJob(),
                 updateLineageRow.getRun().getUuid(),
                 runState,
                 event.getEventTime().toInstant());

diff --git a/api/src/main/java/marquez/service/RunService.java b/api/src/main/java/marquez/service/RunService.java
@@ -27,6 +27,7 @@
 import marquez.common.models.RunId;
 import marquez.common.models.RunState;
 import marquez.db.BaseDao;
+import marquez.db.JobDao;
 import marquez.db.JobVersionDao;
 import marquez.db.JobVersionDao.BagOfJobVersionInfo;
 import marquez.db.RunStateDao;
@@ -48,13 +49,15 @@ public class RunService extends DelegatingDaos.DelegatingRunDao {
   private final JobVersionDao jobVersionDao;
   private final RunStateDao runStateDao;
   private final Collection<RunTransitionListener> runTransitionListeners;
+  private final JobDao jobDao;
 
   public RunService(
       @NonNull BaseDao baseDao, Collection<RunTransitionListener> runTransitionListeners) {
     super(baseDao.createRunDao());
     this.jobVersionDao = baseDao.createJobVersionDao();
     this.runStateDao = baseDao.createRunStateDao();
     this.runTransitionListeners = runTransitionListeners;
+    this.jobDao = baseDao.createJobDao();
   }
 
   /**
@@ -83,8 +86,9 @@ public void markRunAs(
     if (runState.isDone()) {
       BagOfJobVersionInfo bagOfJobVersionInfo =
           jobVersionDao.upsertJobVersionOnRunTransition(
-              runRow.getNamespaceName(),
-              runRow.getJobName(),
+              jobDao
+                  .findJobByNameAsRow(runRow.getNamespaceName(), runRow.getJobName())
+                  .orElseThrow(),
               runRow.getUuid(),
               runState,
               transitionedAt);

diff --git a/api/src/test/java/marquez/db/JobVersionDaoTest.java b/api/src/test/java/marquez/db/JobVersionDaoTest.java
@@ -215,11 +215,7 @@ public void testGetJobVersions() {
             jdbiForTesting, runRow.getUuid(), RunState.COMPLETED, jobMeta.getOutputs());
 
     jobVersionDao.upsertJobVersionOnRunTransition(
-        jobRow.getNamespaceName(),
-        jobRow.getName(),
-        runRow.getUuid(),
-        RunState.COMPLETED,
-        Instant.now());
+        jobRow, runRow.getUuid(), RunState.COMPLETED, Instant.now());
 
     List<JobVersion> jobVersions =
         jobVersionDao.findAllJobVersions(namespaceRow.getName(), jobRow.getName(), 10, 0);
@@ -288,11 +284,7 @@ public void testUpsertJobVersionOnRunTransition() {
     // (6) Add a new job version on the run state transition to COMPLETED.
     final BagOfJobVersionInfo bagOfJobVersionInfo =
         jobVersionDao.upsertJobVersionOnRunTransition(
-            jobRow.getNamespaceName(),
-            jobRow.getName(),
-            runRow.getUuid(),
-            RunState.COMPLETED,
-            newTimestamp());
+            jobRow, runRow.getUuid(), RunState.COMPLETED, newTimestamp());
 
     // Ensure the job version is associated with the latest run.
     final RunRow latestRunRowForJobVersion = runDao.findRunByUuidAsRow(runRow.getUuid()).get();

diff --git a/api/src/test/java/marquez/db/RunDaoTest.java b/api/src/test/java/marquez/db/RunDaoTest.java
@@ -97,11 +97,7 @@ public void getRun() {
         jdbi, runRow.getUuid(), RunState.COMPLETED, jobMeta.getOutputs());
 
     jobVersionDao.upsertJobVersionOnRunTransition(
-        jobRow.getNamespaceName(),
-        jobRow.getName(),
-        runRow.getUuid(),
-        RunState.COMPLETED,
-        Instant.now());
+        jobRow, runRow.getUuid(), RunState.COMPLETED, Instant.now());
 
     Optional<Run> run = runDao.findRunByUuid(runRow.getUuid());
     assertThat(run)
@@ -211,11 +207,7 @@ private Stream<RunRow> createRunsForJob(
                   jdbi, runRow.getUuid(), RunState.COMPLETED, outputs);
 
               jobVersionDao.upsertJobVersionOnRunTransition(
-                  jobRow.getNamespaceName(),
-                  jobRow.getName(),
-                  runRow.getUuid(),
-                  RunState.COMPLETED,
-                  Instant.now());
+                  jobRow, runRow.getUuid(), RunState.COMPLETED, Instant.now());
               return runRow;
             });
   }