Tests for lineage metrics. (#2894)

* Tests for lineage metrics. * Removing comment. * Remove `DbTest` Signed-off-by: Willy Lulciuc <willy.lulciuc@gmail.com> --------- Signed-off-by: Willy Lulciuc <willy.lulciuc@gmail.com> Co-authored-by: Willy Lulciuc <willy.lulciuc@gmail.com> Co-authored-by: Willy Lulciuc <willy@datakin.com>
MarquezProject · Oct 1, 2024 · 9e8e3d3 · 9e8e3d3
1 parent 9450410
commit 9e8e3d3
Show file tree

Hide file tree

Showing 6 changed files with 175 additions and 66 deletions.
diff --git a/api/src/test/java/marquez/api/BaseResourceIntegrationTest.java b/api/src/test/java/marquez/api/BaseResourceIntegrationTest.java
@@ -16,7 +16,6 @@
 import static marquez.common.models.CommonModelGenerator.newNamespaceName;
 import static marquez.common.models.CommonModelGenerator.newOwnerName;
 import static marquez.common.models.CommonModelGenerator.newSourceName;
-import static marquez.db.DbTest.POSTGRES_14;
 
 import com.google.common.collect.ImmutableList;
 import com.google.common.collect.ImmutableSet;
@@ -59,7 +58,7 @@ abstract class BaseResourceIntegrationTest {
   static final String ERROR_FAIL_IF_NOT_IN = "Expected '%s' in '%s'.";
 
   @Container
-  static final PostgreSQLContainer<?> DB_CONTAINER = new PostgreSQLContainer<>(POSTGRES_14);
+  static final PostgreSQLContainer<?> DB_CONTAINER = new PostgreSQLContainer<>("postgres:14");
 
   static {
     DB_CONTAINER.start();

diff --git a/api/src/test/java/marquez/db/DbRetentionTest.java b/api/src/test/java/marquez/db/DbRetentionTest.java
@@ -40,19 +40,61 @@
 import marquez.db.models.RunRow;
 import marquez.db.models.SourceRow;
 import org.jdbi.v3.core.Handle;
+import org.jdbi.v3.jackson2.Jackson2Plugin;
+import org.jdbi.v3.postgres.PostgresPlugin;
+import org.jdbi.v3.sqlobject.SqlObjectPlugin;
+import org.jdbi.v3.testing.junit5.JdbiExtension;
+import org.jdbi.v3.testing.junit5.tc.JdbiTestcontainersExtension;
+import org.junit.jupiter.api.BeforeAll;
+import org.junit.jupiter.api.Order;
 import org.junit.jupiter.api.Tag;
 import org.junit.jupiter.api.Test;
+import org.junit.jupiter.api.extension.RegisterExtension;
+import org.testcontainers.containers.PostgreSQLContainer;
+import org.testcontainers.junit.jupiter.Container;
+import org.testcontainers.junit.jupiter.Testcontainers;
 import org.testcontainers.shaded.com.google.common.collect.ImmutableSet;
+import org.testcontainers.utility.DockerImageName;
 
 /** The test suite for {@link DbRetention}. */
-@Tag("IntegrationTests")
-public class DbRetentionTest extends DbTest {
+@Tag("DataAccessTests, IntegrationTests")
+@Testcontainers
+public class DbRetentionTest {
   private static final int NUMBER_OF_ROWS_PER_BATCH = 10;
   private static final int RETENTION_DAYS = 30;
   private static final boolean DRY_RUN = true;
   private static final Instant OLDER_THAN_X_DAYS = Instant.now().minus(RETENTION_DAYS + 1, DAYS);
   private static final Instant LAST_X_DAYS = Instant.now().minus(RETENTION_DAYS - 1, DAYS);
 
+  static final DockerImageName POSTGRES_16 = DockerImageName.parse("postgres:16");
+
+  @Container
+  @Order(1)
+  static final PostgreSQLContainer<?> DB_CONTAINER = new PostgreSQLContainer<>(POSTGRES_16);
+
+  // Defined statically to significantly improve overall test execution.
+  @RegisterExtension
+  @Order(2)
+  static final JdbiExtension jdbiExtension =
+      JdbiTestcontainersExtension.instance(DB_CONTAINER)
+          .withPlugin(new SqlObjectPlugin())
+          .withPlugin(new PostgresPlugin())
+          .withPlugin(new Jackson2Plugin())
+          .withInitializer(
+              (source, handle) -> {
+                // Apply migrations.
+                DbMigration.migrateDbOrError(source);
+              });
+
+  // Wraps test database connection.
+  static TestingDb DB;
+
+  @BeforeAll
+  public static void setUpOnce() {
+    // Wrap jdbi configured for running container.
+    DB = TestingDb.newInstance(jdbiExtension.getJdbi());
+  }
+
   @Test
   public void testRetentionOnDbOrErrorWithJobsOlderThanXDays() {
     // (1) Add namespace.

diff --git a/api/src/test/java/marquez/db/DbTest.java b/api/src/test/java/marquez/db/DbTest.java
diff --git a/api/src/test/java/marquez/db/DbTestUtils.java b/api/src/test/java/marquez/db/DbTestUtils.java
@@ -415,4 +415,13 @@ public static boolean olEventsExist(
         .mapTo(Boolean.class)
         .one();
   }
+
+  /**
+   * Materializes all views in the database. lineage_events_by_type_hourly_view
+   * lineage_events_by_type_daily_view
+   */
+  public static void materializeViews(@NonNull final Handle handle) {
+    handle.execute("REFRESH MATERIALIZED VIEW lineage_events_by_type_hourly_view");
+    handle.execute("REFRESH MATERIALIZED VIEW lineage_events_by_type_daily_view");
+  }
 }
diff --git a/api/src/test/java/marquez/db/StatsTest.java b/api/src/test/java/marquez/db/StatsTest.java
@@ -0,0 +1,111 @@
+/*
+ * Copyright 2018-2024 contributors to the Marquez project
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+package marquez.db;
+
+import static marquez.api.models.ApiModelGenerator.newRunEvents;
+import static marquez.common.models.CommonModelGenerator.newJobName;
+import static marquez.common.models.CommonModelGenerator.newNamespaceName;
+import static org.assertj.core.api.Assertions.assertThat;
+import static org.assertj.core.api.Assertions.fail;
+
+import io.openlineage.client.OpenLineage;
+import java.net.URI;
+import java.time.Instant;
+import java.time.temporal.ChronoUnit;
+import java.util.List;
+import java.util.Set;
+import marquez.db.models.LineageMetric;
+import org.jdbi.v3.core.Handle;
+import org.jdbi.v3.jackson2.Jackson2Plugin;
+import org.jdbi.v3.postgres.PostgresPlugin;
+import org.jdbi.v3.sqlobject.SqlObjectPlugin;
+import org.jdbi.v3.testing.junit5.JdbiExtension;
+import org.jdbi.v3.testing.junit5.tc.JdbiTestcontainersExtension;
+import org.junit.jupiter.api.BeforeAll;
+import org.junit.jupiter.api.Order;
+import org.junit.jupiter.api.Tag;
+import org.junit.jupiter.api.Test;
+import org.junit.jupiter.api.extension.RegisterExtension;
+import org.testcontainers.containers.PostgreSQLContainer;
+import org.testcontainers.junit.jupiter.Container;
+import org.testcontainers.junit.jupiter.Testcontainers;
+import org.testcontainers.utility.DockerImageName;
+
+@Tag("DataAccessTests, IntegrationTests")
+@Testcontainers
+public class StatsTest {
+  static final DockerImageName POSTGRES_16 = DockerImageName.parse("postgres:16");
+
+  @Container
+  @Order(1)
+  static final PostgreSQLContainer<?> DB_CONTAINER = new PostgreSQLContainer<>(POSTGRES_16);
+
+  // Defined statically to significantly improve overall test execution.
+  @RegisterExtension
+  @Order(2)
+  static final JdbiExtension jdbiExtension =
+      JdbiTestcontainersExtension.instance(DB_CONTAINER)
+          .withPlugin(new SqlObjectPlugin())
+          .withPlugin(new PostgresPlugin())
+          .withPlugin(new Jackson2Plugin())
+          .withInitializer(
+              (source, handle) -> {
+                // Apply migrations.
+                DbMigration.migrateDbOrError(source);
+              });
+
+  // Wraps test database connection.
+  static TestingDb DB;
+
+  @BeforeAll
+  public static void setUpOnce() {
+    // Wrap jdbi configured for running container.
+    DB = TestingDb.newInstance(jdbiExtension.getJdbi());
+  }
+
+  @Test
+  public void testGetStatsForLastDay() {
+    // (1) Configure OL.
+    final URI olProducer = URI.create("https://test.com/test");
+    final OpenLineage ol = new OpenLineage(olProducer);
+
+    // (2) Add namespace and job for OL events.
+    final String namespaceName = newNamespaceName().getValue();
+    final String jobName = newJobName().getValue();
+
+    // (3) Create some 1 hour old OL events.
+    int hourEvents = 4;
+    final Set<OpenLineage.RunEvent> hourEventSet =
+        newRunEvents(
+            ol, Instant.now().minus(1, ChronoUnit.HOURS), namespaceName, jobName, hourEvents);
+    DB.insertAll(hourEventSet);
+
+    // (4) Create some 2 day old OL events.
+    int dayEvents = 2;
+    final Set<OpenLineage.RunEvent> dayEventSet =
+        newRunEvents(
+            ol, Instant.now().minus(2, ChronoUnit.DAYS), namespaceName, jobName, dayEvents);
+    DB.insertAll(dayEventSet);
+
+    // (4) Materialize views to flush out view data.
+    try (final Handle handle = DB.open()) {
+      DbTestUtils.materializeViews(handle);
+    } catch (Exception e) {
+      fail("failed to apply dry run", e);
+    }
+
+    List<LineageMetric> lastDayLineageMetrics = DB.lastDayLineageMetrics();
+    List<LineageMetric> lastWeekLineageMetrics = DB.lastWeekLineageMetrics();
+
+    assertThat(lastDayLineageMetrics).isNotEmpty();
+    assertThat(lastDayLineageMetrics.get(lastDayLineageMetrics.size() - 1).getComplete())
+        .isEqualTo(hourEvents);
+
+    assertThat(lastWeekLineageMetrics).isNotEmpty();
+    assertThat(lastWeekLineageMetrics.get(lastWeekLineageMetrics.size() - 2).getComplete())
+        .isEqualTo(dayEvents);
+  }
+}
diff --git a/api/src/test/java/marquez/db/TestingDb.java b/api/src/test/java/marquez/db/TestingDb.java
@@ -10,6 +10,7 @@
 import com.google.common.collect.ImmutableSet;
 import io.openlineage.client.OpenLineage;
 import java.time.Instant;
+import java.util.List;
 import java.util.Optional;
 import java.util.Set;
 import java.util.UUID;
@@ -21,6 +22,7 @@
 import marquez.db.models.DatasetVersionRow;
 import marquez.db.models.JobRow;
 import marquez.db.models.JobVersionRow;
+import marquez.db.models.LineageMetric;
 import marquez.db.models.NamespaceRow;
 import marquez.db.models.RunArgsRow;
 import marquez.db.models.RunRow;
@@ -246,4 +248,12 @@ void insert(@NonNull OpenLineage.RunEvent olEvent) {
   Handle open() {
     return delegate.open();
   }
+
+  List<LineageMetric> lastDayLineageMetrics() {
+    return delegate.onDemand(StatsDao.class).getLastDayMetrics();
+  }
+
+  List<LineageMetric> lastWeekLineageMetrics() {
+    return delegate.onDemand(StatsDao.class).getLastWeekMetrics();
+  }
 }