add operation column

Signed-off-by: Pawel Leszczynski <leszczynski.pawel@gmail.com>

Author: pawel-big-lebowski

api/src/main/java/marquez/common/Utils.java

@@ -216,6 +216,7 @@ public static Version newJobVersionFor(
    * @param sourceName The source name of the dataset.
    * @param physicalName The physical name of the dataset.
    * @param datasetName The dataset name.
+   * @param stateChange The dataset stateChange like CREATE, DROP, TRUNCATE.
    * @param fields The fields of the dataset.
    * @param runId The UUID of the run linked to the dataset.
    * @return A {@link Version} object based on the specified job meta.
@@ -225,6 +226,7 @@ public static Version newDatasetVersionFor(
       String sourceName,
       String physicalName,
       String datasetName,
+      String stateChange,
       List<LineageEvent.SchemaField> fields,
       UUID runId) {
     DatasetVersionData data =
@@ -233,6 +235,7 @@ public static Version newDatasetVersionFor(
             .sourceName(sourceName)
             .physicalName(physicalName)
             .datasetName(datasetName)
+            .stateChange(stateChange)
             .schemaFields(fields)
             .runId(runId)
             .build();
@@ -271,6 +274,7 @@ private static Version newDatasetVersionFor(DatasetVersionData data) {
                 data.getPhysicalName(),
                 data.getSchemaLocation(),
                 data.getFields().stream().map(Utils::joinField).collect(joining(VERSION_DELIM)),
+                data.getStateChange(),
                 data.getRunId())
             .getBytes(UTF_8);
     return Version.of(UUID.nameUUIDFromBytes(bytes));
@@ -287,6 +291,7 @@ private static class DatasetVersionData {
     private String sourceName;
     private String physicalName;
     private String datasetName;
+    private String stateChange;
     private String schemaLocation;
     private Set<Triple<String, String, String>> fields;
     private UUID runId;

api/src/main/java/marquez/db/Columns.java

@@ -84,6 +84,7 @@ private Columns() {}
 
   /* DATASET VERSION ROW COLUMNS */
   public static final String FIELD_UUIDS = "field_uuids";
+  public static final String STATE_CHANGE = "state_change";
 
   /* STREAM VERSION ROW COLUMNS */
   public static final String SCHEMA_LOCATION = "schema_location";

api/src/main/java/marquez/db/DatasetDao.java

@@ -84,15 +84,15 @@ void updateLastModifiedAt(
           + "    WHERE d.namespace_name = :namespaceName\n"
           + "    AND d.name = :datasetName\n"
           + "), dataset_runs AS (\n"
-          + "    SELECT d.uuid, d.name, d.namespace_name, dv.run_uuid, event_time, event\n"
+          + "    SELECT d.uuid, d.name, d.namespace_name, dv.run_uuid, dv.state_change, event_time, event\n"
           + "    FROM selected_datasets d\n"
           + "    INNER JOIN dataset_versions dv ON dv.uuid = d.current_version_uuid\n"
           + "    LEFT JOIN LATERAL (\n"
           + "        SELECT run_uuid, event_time, event FROM lineage_events\n"
           + "        WHERE run_uuid = dv.run_uuid\n"
           + "    ) e ON e.run_uuid = dv.run_uuid\n"
           + "    UNION\n"
-          + "    SELECT d.uuid, d.name, d.namespace_name, rim.run_uuid, event_time, event\n"
+          + "    SELECT d.uuid, d.name, d.namespace_name, rim.run_uuid, state_change, event_time, event\n"
           + "    FROM selected_datasets d\n"
           + "    INNER JOIN dataset_versions dv ON dv.uuid = d.current_version_uuid\n"
           + "    LEFT JOIN runs_input_mapping rim ON dv.uuid = rim.dataset_version_uuid\n"
@@ -101,7 +101,7 @@ void updateLastModifiedAt(
           + "        WHERE run_uuid = rim.run_uuid\n"
           + "    ) e ON e.run_uuid = rim.run_uuid\n"
           + ")\n"
-          + "SELECT d.*, dv.fields, sv.schema_location, t.tags, facets\n"
+          + "SELECT d.*, dv.fields, dv.state_change, sv.schema_location, t.tags, facets\n"
           + "FROM selected_datasets d\n"
           + "LEFT JOIN dataset_versions dv ON d.current_version_uuid = dv.uuid\n"
           + "LEFT JOIN stream_versions AS sv ON sv.dataset_version_uuid = dv.uuid\n"
@@ -154,15 +154,15 @@ default void setFields(Dataset ds) {
           + "    ORDER BY d.name\n"
           + "    LIMIT :limit OFFSET :offset\n"
           + "), dataset_runs AS (\n"
-          + "    SELECT d.uuid, d.name, d.namespace_name, dv.run_uuid, event_time, event\n"
+          + "    SELECT d.uuid, d.name, d.namespace_name, dv.run_uuid, dv.state_change, event_time, event\n"
           + "    FROM selected_datasets d\n"
           + "    INNER JOIN dataset_versions dv ON dv.uuid = d.current_version_uuid\n"
           + "    LEFT JOIN LATERAL (\n"
           + "        SELECT run_uuid, event_time, event FROM lineage_events\n"
           + "        WHERE run_uuid = dv.run_uuid\n"
           + "    ) e ON e.run_uuid = dv.run_uuid\n"
           + "    UNION\n"
-          + "    SELECT d.uuid, d.name, d.namespace_name, rim.run_uuid, event_time, event\n"
+          + "    SELECT d.uuid, d.name, d.namespace_name, rim.run_uuid, state_change, event_time, event\n"
           + "    FROM selected_datasets d\n"
           + "    INNER JOIN dataset_versions dv ON dv.uuid = d.current_version_uuid\n"
           + "    LEFT JOIN runs_input_mapping rim ON dv.uuid = rim.dataset_version_uuid\n"
@@ -171,7 +171,7 @@ default void setFields(Dataset ds) {
           + "        WHERE run_uuid = rim.run_uuid\n"
           + "    ) e ON e.run_uuid = rim.run_uuid\n"
           + ")\n"
-          + "SELECT d.*, dv.fields, sv.schema_location, t.tags, facets\n"
+          + "SELECT d.*, dv.fields, dv.state_change, sv.schema_location, t.tags, facets\n"
           + "FROM selected_datasets d\n"
           + "LEFT JOIN dataset_versions dv ON d.current_version_uuid = dv.uuid\n"
           + "LEFT JOIN stream_versions AS sv ON sv.dataset_version_uuid = dv.uuid\n"
@@ -352,6 +352,7 @@ default Dataset upsertDatasetMeta(
                 now,
                 namespaceName.getValue(),
                 datasetName.getValue(),
+                null,
                 datasetMeta);
 
     return findWithTags(namespaceName.getValue(), datasetName.getValue()).get();

api/src/main/java/marquez/db/DatasetVersionDao.java

@@ -60,6 +60,7 @@ default DatasetVersionRow upsertDatasetVersion(
       Instant now,
       String namespaceName,
       String datasetName,
+      String stateChange,
       DatasetMeta datasetMeta) {
     TagDao tagDao = createTagDao();
     DatasetFieldDao datasetFieldDao = createDatasetFieldDao();
@@ -75,7 +76,8 @@ default DatasetVersionRow upsertDatasetVersion(
             datasetMeta.getRunId().map(RunId::getValue).orElse(null),
             toPgObjectFields(datasetMeta.getFields()),
             namespaceName,
-            datasetName);
+            datasetName,
+            stateChange);
     updateDatasetVersionMetric(
         namespaceName,
         datasetMeta.getType().toString(),
@@ -179,7 +181,7 @@ default void updateDatasetVersionMetric(
           + "    FROM selected_dataset_version_runs dv\n"
           + "    LEFT JOIN lineage_events le ON le.run_uuid = dv.run_uuid\n"
           + ")\n"
-          + "SELECT d.type, d.name, d.physical_name, d.namespace_name, d.source_name, d.description,\n"
+          + "SELECT d.type, d.name, d.physical_name, d.namespace_name, d.source_name, d.description, dv.state_change, \n"
           + "    dv.created_at, dv.version, dv.fields, dv.run_uuid AS createdByRunUuid, sv.schema_location,\n"
           + "    t.tags, f.facets\n"
           + "FROM selected_dataset_versions dv\n"
@@ -221,7 +223,7 @@ default void updateDatasetVersionMetric(
           + "    FROM selected_dataset_version_runs dv\n"
           + "    LEFT JOIN lineage_events le ON le.run_uuid = dv.run_uuid\n"
           + ")\n"
-          + "SELECT d.type, d.name, d.physical_name, d.namespace_name, d.source_name, d.description,\n"
+          + "SELECT d.type, d.name, d.physical_name, d.namespace_name, d.source_name, d.description, dv.state_change, \n"
           + "    dv.created_at, dv.version, dv.fields, dv.run_uuid AS createdByRunUuid, sv.schema_location,\n"
           + "    t.tags, f.facets\n"
           + "FROM selected_dataset_versions dv\n"
@@ -292,7 +294,7 @@ default Optional<DatasetVersion> findByWithRun(UUID version) {
           + "    FROM selected_dataset_version_runs dv\n"
           + "    LEFT JOIN lineage_events le ON le.run_uuid = dv.run_uuid\n"
           + ")\n"
-          + "SELECT d.type, d.name, d.physical_name, d.namespace_name, d.source_name, d.description,\n"
+          + "SELECT d.type, d.name, d.physical_name, d.namespace_name, d.source_name, d.description, dv.state_change,\n"
           + "    dv.created_at, dv.version, dv.fields, dv.run_uuid AS createdByRunUuid, sv.schema_location,\n"
           + "    t.tags, f.facets\n"
           + "FROM selected_dataset_versions dv\n"
@@ -336,9 +338,9 @@ default List<DatasetVersion> findAllWithRun(
 
   @SqlQuery(
       "INSERT INTO dataset_versions "
-          + "(uuid, created_at, dataset_uuid, version, run_uuid, fields, namespace_name, dataset_name) "
+          + "(uuid, created_at, dataset_uuid, version, run_uuid, fields, namespace_name, dataset_name, state_change) "
           + "VALUES "
-          + "(:uuid, :now, :datasetUuid, :version, :runUuid, :fields, :namespaceName, :datasetName) "
+          + "(:uuid, :now, :datasetUuid, :version, :runUuid, :fields, :namespaceName, :datasetName, :stateChange) "
           + "ON CONFLICT(version) "
           + "DO UPDATE SET "
           + "run_uuid = EXCLUDED.run_uuid "
@@ -351,7 +353,8 @@ DatasetVersionRow upsert(
       UUID runUuid,
       PGobject fields,
       String namespaceName,
-      String datasetName);
+      String datasetName,
+      String stateChange);
 
   @SqlUpdate("UPDATE dataset_versions SET fields = :fields WHERE uuid = :uuid")
   void updateFields(UUID uuid, PGobject fields);

api/src/main/java/marquez/db/LineageDao.java

@@ -79,7 +79,7 @@ public interface LineageDao {
   Optional<UUID> getJobUuid(String jobName, String namespace);
 
   @SqlQuery(
-      "SELECT ds.*, dv.fields\n"
+      "SELECT ds.*, dv.fields, dv.state_change\n"
           + "FROM datasets ds\n"
           + "LEFT JOIN dataset_versions dv on dv.uuid = ds.current_version_uuid\n"
           + "WHERE ds.uuid IN (<dsUuids>);")

api/src/main/java/marquez/db/OpenLineageDao.java

@@ -367,6 +367,12 @@ default DatasetRecord upsertLineageDataset(
             .map(SchemaDatasetFacet::getFields)
             .orElse(null);
 
+    String dsStateChange =
+        Optional.ofNullable(ds.getFacets())
+            .map(DatasetFacets::getTableStateChange)
+            .map(LineageEvent.TableStateChangeFacet::getStateChange)
+            .orElse(null);
+
     final DatasetRow dsRow = datasetRow;
     DatasetVersionRow datasetVersionRow =
         datasetRow
@@ -383,6 +389,7 @@ default DatasetRecord upsertLineageDataset(
                               source.getName(),
                               dsRow.getPhysicalName(),
                               dsRow.getName(),
+                              dsStateChange,
                               fields,
                               runUuid)
                           .getValue();
@@ -395,8 +402,8 @@ default DatasetRecord upsertLineageDataset(
                           isInput ? null : runUuid,
                           datasetVersionDao.toPgObjectSchemaFields(fields),
                           dsNamespace.getName(),
-                          ds.getName());
-
+                          ds.getName(),
+                          dsStateChange);
                   return row;
                 });
     List<DatasetFieldMapping> datasetFieldMappings = new ArrayList<>();

api/src/main/java/marquez/db/RunDao.java

@@ -289,6 +289,7 @@ default void upsertOutputDatasetsFor(UUID runUuid, ImmutableSet<DatasetId> runOu
                           d.getSourceName().getValue(),
                           d.getPhysicalName().getValue(),
                           d.getName().getValue(),
+                          null,
                           toSchemaFields(d.getFields()),
                           runUuid)
                       .getValue();
@@ -300,7 +301,8 @@ default void upsertOutputDatasetsFor(UUID runUuid, ImmutableSet<DatasetId> runOu
                   runUuid,
                   datasetVersionDao.toPgObjectFields(d.getFields()),
                   d.getNamespace().getValue(),
-                  d.getName().getValue());
+                  d.getName().getValue(),
+                  null);
             });
       }
     }

api/src/main/java/marquez/db/mappers/DatasetDataMapper.java

@@ -50,7 +50,8 @@ public DatasetData map(@NonNull ResultSet results, @NonNull StatementContext con
         toFields(results, "fields"),
         ImmutableSet.of(),
         timestampOrNull(results, Columns.LAST_MODIFIED_AT),
-        stringOrNull(results, Columns.DESCRIPTION));
+        stringOrNull(results, Columns.DESCRIPTION),
+        stringOrNull(results, Columns.STATE_CHANGE));
   }
 
   public static ImmutableList<Field> toFields(ResultSet results, String column)

api/src/main/java/marquez/db/mappers/DatasetMapper.java

@@ -74,6 +74,7 @@ public Dataset map(@NonNull ResultSet results, @NonNull StatementContext context
           toFields(results, "fields"),
           toTags(results, "tags"),
           timestampOrNull(results, Columns.LAST_MODIFIED_AT),
+          stringOrNull(results, Columns.STATE_CHANGE),
           stringOrNull(results, Columns.DESCRIPTION),
           uuidOrNull(results, Columns.CURRENT_VERSION_UUID),
           toFacetsOrNull(results, Columns.FACETS));
@@ -91,6 +92,7 @@ public Dataset map(@NonNull ResultSet results, @NonNull StatementContext context
           toFields(results, "fields"),
           toTags(results, "tags"),
           timestampOrNull(results, Columns.LAST_MODIFIED_AT),
+          stringOrNull(results, Columns.STATE_CHANGE),
           stringOrNull(results, Columns.DESCRIPTION),
           uuidOrNull(results, Columns.CURRENT_VERSION_UUID),
           toFacetsOrNull(results, Columns.FACETS));

api/src/main/java/marquez/db/mappers/DatasetVersionMapper.java

@@ -66,6 +66,7 @@ public DatasetVersion map(@NonNull ResultSet results, @NonNull StatementContext
               toFields(results, "fields"),
               columnNames.contains("tags") ? toTags(results, "tags") : null,
               stringOrNull(results, Columns.DESCRIPTION),
+              stringOrNull(results, Columns.STATE_CHANGE),
               null,
               toFacetsOrNull(results, Columns.FACETS));
     } else {
@@ -83,6 +84,7 @@ public DatasetVersion map(@NonNull ResultSet results, @NonNull StatementContext
               toFields(results, "fields"),
               columnNames.contains("tags") ? toTags(results, "tags") : null,
               stringOrNull(results, Columns.DESCRIPTION),
+              stringOrNull(results, Columns.STATE_CHANGE),
               null,
               toFacetsOrNull(results, Columns.FACETS));
     }