Model and store column lineage in Marquez DB

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

@@ -50,4 +50,7 @@ public interface BaseDao extends SqlObject {
 
   @CreateSqlObject
   OpenLineageDao createOpenLineageDao();
+
+  @CreateSqlObject
+  ColumnLevelLineageDao createColumnLevelLineageDao();
 }

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

@@ -0,0 +1,79 @@
+package marquez.db;
+
+/*
+ * Copyright 2018-2022 contributors to the Marquez project
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+import java.time.Instant;
+import java.util.Optional;
+import java.util.UUID;
+import marquez.db.mappers.ColumnLevelLineageRowMapper;
+import marquez.db.models.ColumnLevelLineageRow;
+import org.jdbi.v3.sqlobject.config.RegisterRowMapper;
+import org.jdbi.v3.sqlobject.statement.SqlQuery;
+import org.jdbi.v3.sqlobject.statement.SqlUpdate;
+
+@RegisterRowMapper(ColumnLevelLineageRowMapper.class)
+public interface ColumnLevelLineageDao extends BaseDao {
+
+  default ColumnLevelLineageRow upsertColumnLevelLineageRow(
+      UUID uuid,
+      UUID dataset_version_uuid,
+      String output_column_name,
+      String input_field,
+      String transformation_description,
+      String transformation_type,
+      Instant now) {
+    doUpsertColumnLevelLineageRow(
+        uuid,
+        dataset_version_uuid,
+        output_column_name,
+        input_field,
+        transformation_description,
+        transformation_type,
+        now);
+    return findColumnLevelLineageByDatasetVersionColumnAndInput(
+            dataset_version_uuid, output_column_name, input_field)
+        .orElseThrow();
+  }
+
+  @SqlQuery("SELECT * FROM column_level_lineage WHERE dataset_version_uuid = :datasetVersionUuid")
+  Optional<ColumnLevelLineageRow> findColumnLevelLineageByDatasetVersionColumnAndInput(
+      UUID datasetVersionUuid, String outputColumnName, String inputField);
+
+  @SqlUpdate(
+      "INSERT INTO column_level_lineage ("
+          + "uuid, "
+          + "dataset_version_uuid, "
+          + "output_column_name, "
+          + "input_field, "
+          + "transformation_description, "
+          + "transformation_type, "
+          + "created_at, "
+          + "updated_at"
+          + ") VALUES ( "
+          + ":uuid, "
+          + ":dataset_version_uuid, "
+          + ":output_column_name, "
+          + ":input_field, "
+          + ":transformation_description, "
+          + ":transformation_type, "
+          + ":now, "
+          + ":now) "
+          + "ON CONFLICT (dataset_version_uuid, output_column_name, input_field) "
+          + "DO UPDATE SET "
+          + "input_field = EXCLUDED.input_field, "
+          + "transformation_description = EXCLUDED.transformation_description, "
+          + "transformation_type = EXCLUDED.transformation_type, "
+          + "updated_at = EXCLUDED.updated_at "
+          + "RETURNING *")
+  void doUpsertColumnLevelLineageRow(
+      UUID uuid,
+      UUID dataset_version_uuid,
+      String output_column_name,
+      String input_field,
+      String transformation_description,
+      String transformation_type,
+      Instant now);
+}

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

@@ -126,6 +126,12 @@ private Columns() {}
   public static final String RUN_UUID = "run_uuid";
   public static final String STATE = "state";
 
+  /* COLUMN LEVEL LINEAGE ROW COLUMNS */
+  public static final String OUTPUT_COLUMN_NAME = "output_column_name";
+  public static final String INPUT_FIELD = "input_field";
+  public static final String TRANSFORMATION_DESCRIPTION = "transformation_description";
+  public static final String TRANSFORMATION_TYPE = "transformation_type";
+
   public static UUID uuidOrNull(final ResultSet results, final String column) throws SQLException {
     if (results.getObject(column) == null) {
       return null;

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

@@ -30,6 +30,7 @@
 import marquez.db.DatasetFieldDao.DatasetFieldMapping;
 import marquez.db.JobVersionDao.BagOfJobVersionInfo;
 import marquez.db.mappers.LineageEventMapper;
+import marquez.db.models.ColumnLevelLineageRow;
 import marquez.db.models.DatasetFieldRow;
 import marquez.db.models.DatasetRow;
 import marquez.db.models.DatasetVersionRow;
@@ -108,6 +109,7 @@ default UpdateLineageRow updateBaseMarquezModel(LineageEvent event, ObjectMapper
     RunDao runDao = createRunDao();
     RunArgsDao runArgsDao = createRunArgsDao();
     RunStateDao runStateDao = createRunStateDao();
+    ColumnLevelLineageDao columnLevelLineageDao = createColumnLevelLineageDao();
 
     Instant now = event.getEventTime().withZoneSameInstant(ZoneId.of("UTC")).toInstant();
 
@@ -349,7 +351,8 @@ default UpdateLineageRow updateBaseMarquezModel(LineageEvent event, ObjectMapper
                 datasetDao,
                 datasetVersionDao,
                 datasetFieldDao,
-                runDao);
+                runDao,
+                columnLevelLineageDao);
         datasetInputs.add(record);
       }
     }
@@ -370,7 +373,8 @@ default UpdateLineageRow updateBaseMarquezModel(LineageEvent event, ObjectMapper
                 datasetDao,
                 datasetVersionDao,
                 datasetFieldDao,
-                runDao);
+                runDao,
+                columnLevelLineageDao);
         datasetOutputs.add(record);
       }
     }
@@ -435,7 +439,8 @@ default DatasetRecord upsertLineageDataset(
       DatasetDao datasetDao,
       DatasetVersionDao datasetVersionDao,
       DatasetFieldDao datasetFieldDao,
-      RunDao runDao) {
+      RunDao runDao,
+      ColumnLevelLineageDao columnLevelLineageDao) {
     NamespaceRow dsNamespace =
         namespaceDao.upsertNamespaceRow(
             UUID.randomUUID(), now, ds.getNamespace(), DEFAULT_NAMESPACE_OWNER);
@@ -543,6 +548,37 @@ default DatasetRecord upsertLineageDataset(
     }
     datasetFieldDao.updateFieldMapping(datasetFieldMappings);
 
+    List<ColumnLevelLineageRow> columnLineageRows = null;
+    if (ds.getFacets().getColumnLineage() != null) {
+      columnLineageRows = new ArrayList<>();
+      List<LineageEvent.ColumnLineageOutputColumn> columnLevelLineageOutputColumnsList =
+          Optional.ofNullable(ds.getFacets())
+              .map(DatasetFacets::getColumnLineage)
+              .map(LineageEvent.ColumnLineageFacet::getOutputColumnsList)
+              .orElse(null);
+
+      if (columnLevelLineageOutputColumnsList != null) {
+        for (LineageEvent.ColumnLineageOutputColumn outputColumn :
+            columnLevelLineageOutputColumnsList) {
+          for (LineageEvent.ColumnLineageInputField inputField : outputColumn.getInputFields()) {
+            columnLineageRows.add(
+                columnLevelLineageDao.upsertColumnLevelLineageRow(
+                    UUID.randomUUID(),
+                    datasetVersionRow.getUuid(),
+                    outputColumn.getName(),
+                    String.format(
+                        "%s.%s.%s",
+                        inputField.getDatasetNamespace(),
+                        inputField.getDatasetName(),
+                        inputField.getFieldName()),
+                    outputColumn.getTransformationDescription(),
+                    outputColumn.getTransformationType(),
+                    now));
+          }
+        }
+      }
+    }
+
     if (isInput) {
       runDao.updateInputMapping(runUuid, datasetVersionRow.getUuid());
 
@@ -555,7 +591,7 @@ default DatasetRecord upsertLineageDataset(
       }
     }
 
-    return new DatasetRecord(datasetRow, datasetVersionRow, datasetNamespace);
+    return new DatasetRecord(datasetRow, datasetVersionRow, datasetNamespace, columnLineageRows);
   }
 
   default String formatDatasetName(String name) {

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

@@ -0,0 +1,33 @@
+package marquez.db.mappers;
+
+import static marquez.db.Columns.INPUT_FIELD;
+import static marquez.db.Columns.TRANSFORMATION_DESCRIPTION;
+import static marquez.db.Columns.TRANSFORMATION_TYPE;
+import static marquez.db.Columns.stringOrThrow;
+import static marquez.db.Columns.timestampOrThrow;
+import static marquez.db.Columns.uuidOrThrow;
+
+import java.sql.ResultSet;
+import java.sql.SQLException;
+import lombok.NonNull;
+import marquez.db.Columns;
+import marquez.db.models.ColumnLevelLineageRow;
+import org.jdbi.v3.core.mapper.RowMapper;
+import org.jdbi.v3.core.statement.StatementContext;
+
+public class ColumnLevelLineageRowMapper implements RowMapper<ColumnLevelLineageRow> {
+
+  @Override
+  public ColumnLevelLineageRow map(@NonNull ResultSet results, @NonNull StatementContext context)
+      throws SQLException {
+    return new ColumnLevelLineageRow(
+        uuidOrThrow(results, Columns.ROW_UUID),
+        uuidOrThrow(results, Columns.DATASET_VERSION_UUID),
+        stringOrThrow(results, Columns.OUTPUT_COLUMN_NAME),
+        stringOrThrow(results, INPUT_FIELD),
+        stringOrThrow(results, TRANSFORMATION_DESCRIPTION),
+        stringOrThrow(results, TRANSFORMATION_TYPE),
+        timestampOrThrow(results, Columns.CREATED_AT),
+        timestampOrThrow(results, Columns.UPDATED_AT));
+  }
+}

api/src/main/java/marquez/db/models/ColumnLevelLineageRow.java

@@ -0,0 +1,23 @@
+package marquez.db.models;
+
+import java.time.Instant;
+import java.util.UUID;
+import lombok.AllArgsConstructor;
+import lombok.EqualsAndHashCode;
+import lombok.Getter;
+import lombok.NonNull;
+import lombok.ToString;
+
+@AllArgsConstructor
+@EqualsAndHashCode
+@ToString
+public class ColumnLevelLineageRow {
+  @Getter @NonNull private final UUID uuid;
+  @Getter @NonNull private final UUID datasetVersionUuid;
+  @Getter @NonNull private final String outputColumnName;
+  @Getter @NonNull private final String inputField;
+  @Getter @NonNull private final String transformationDescription;
+  @Getter @NonNull private final String transformationType;
+  @Getter @NonNull private final Instant createdAt;
+  @Getter @NonNull private Instant updatedAt;
+}

api/src/main/java/marquez/db/models/UpdateLineageRow.java

@@ -30,5 +30,6 @@ public static class DatasetRecord {
     DatasetRow datasetRow;
     DatasetVersionRow datasetVersionRow;
     NamespaceRow namespaceRow;
+    List<ColumnLevelLineageRow> columnLineageRows;
   }
 }

api/src/main/java/marquez/service/models/LineageEvent.java

@@ -317,14 +317,16 @@ public static class Dataset extends BaseJsonModel {
     "schema",
     "dataSource",
     "description",
-    "lifecycleStateChange"
+    "lifecycleStateChange",
+    "columnLineage",
   })
   public static class DatasetFacets {
 
     @Valid private DocumentationDatasetFacet documentation;
     @Valid private SchemaDatasetFacet schema;
     @Valid private LifecycleStateChangeFacet lifecycleStateChange;
     @Valid private DatasourceDatasetFacet dataSource;
+    @Valid private ColumnLineageFacet columnLineage;
     private String description;
     @Builder.Default @JsonIgnore private Map<String, Object> additional = new LinkedHashMap<>();
 
@@ -354,6 +356,10 @@ public DatasourceDatasetFacet getDataSource() {
       return dataSource;
     }
 
+    public ColumnLineageFacet getColumnLineage() {
+      return columnLineage;
+    }
+
     public String getDescription() {
       return description;
     }
@@ -442,4 +448,52 @@ public LifecycleStateChangeFacet(
       this.lifecycleStateChange = lifecycleStateChange;
     }
   }
+
+  @NoArgsConstructor
+  @Getter
+  @Setter
+  @Valid
+  @ToString
+  public static class ColumnLineageFacet extends BaseFacet {
+
+    private List<ColumnLineageOutputColumn> outputColumnsList;
+
+    @Builder
+    public ColumnLineageFacet(
+        @NotNull URI _producer,
+        @NotNull URI _schemaURL,
+        List<ColumnLineageOutputColumn> outputColumnsList) {
+      super(_producer, _schemaURL);
+      this.outputColumnsList = outputColumnsList;
+    }
+  }
+
+  @Builder
+  @AllArgsConstructor
+  @NoArgsConstructor
+  @Setter
+  @Getter
+  @Valid
+  @ToString
+  public static class ColumnLineageOutputColumn extends BaseJsonModel {
+
+    @NotNull private String name;
+    @NotNull private List<ColumnLineageInputField> inputFields;
+    @NotNull private String transformationDescription;
+    @NotNull private String transformationType;
+  }
+
+  @Builder
+  @AllArgsConstructor
+  @NoArgsConstructor
+  @Setter
+  @Getter
+  @Valid
+  @ToString
+  public static class ColumnLineageInputField extends BaseJsonModel {
+
+    @NotNull private String datasetNamespace;
+    @NotNull private String datasetName;
+    @NotNull private String fieldName;
+  }
 }

api/src/main/resources/marquez/db/migration/V47__column_level_lineage.sql

@@ -0,0 +1,25 @@
+/* SPDX-License-Identifier: Apache-2.0 */
+
+-- DROP TABLE column_level_lineage;
+
+CREATE TABLE column_level_lineage (
+                                      uuid                          uuid primary key,
+                                      dataset_version_uuid          uuid REFERENCES dataset_versions(uuid),
+                                      output_column_name            VARCHAR(255) NOT NULL,
+                                      input_field                   VARCHAR(255) NOT NULL, -- reference dataset_fields.uuid
+                                      transformation_description    VARCHAR(255) NOT NULL,
+                                      transformation_type           VARCHAR(255) NOT NULL,
+                                      created_at                    TIMESTAMP NOT NULL,
+                                      updated_at                    TIMESTAMP NOT NULL,
+                                      UNIQUE (dataset_version_uuid, output_column_name, input_field)
+);
+
+-- INSERT INTO column_level_lineage (uuid, dataset_version_uuid, output_column_name, input_field,
+--                                   transformation_description, transformation_type, created_at,
+--                                   updated_at)
+-- VALUES (md5('whatever')::uuid, md5('dataset_version_uuid_example')::uuid,  'column_a', 'input_field_a', 'Identity transformation', 'IDENTITY', current_timestamp, current_timestamp);
+--
+-- INSERT INTO column_level_lineage (uuid, dataset_version_uuid, output_column_name, input_field,
+--                                   transformation_description, transformation_type, created_at,
+--                                   updated_at)
+-- VALUES (md5('whatever')::uuid, md5('dataset_version_uuid_example')::uuid,  'column_a', 'input_field_b', 'Identity transformation', 'IDENTITY', current_timestamp, current_timestamp);