[GOBBLIN-2204] Implement FileSize Data Quality for FileBasedCopy

* Compute data quality and update task states

* Retry for failed data quality check

* Only do data quality check if data quality flag is enabled

* Computing overall dataquality of data set and adding it to dataset summary

* Changes for testing failed data quality changes

* Updating task state and computing overall data quality

* Refactor and checkstyle fix

* Propogate data quality from fork to task state

* Refactoring and unit tests

* Add unit tests for FileAwareInputStreamDataWriterTest

* refactor changes

* Compute task level data quality

* Some refactoring and code cleanup

* Update tests

* Fix checkstyle failures

* Code cleanup

* Adding metrics for overall data quality

* Address PR comments and add metric for no of files evaluated for data quality

* Refactor dq evaluation

* Address PR comments

* Introduce data quality flag, to control and manage job failures due to data quality

* Address PR comments

* Switch back to CounterBuilder approach for metrics

* Evaluate data quality only for commit via CommitActivity

* Remove unused enum

* move getDataQuality to JobState

* Address PR comments

* Fix the policy evaluation logic when both optional and mandatory policies are provided

* Remove codestyle errors and remove flag check at task level

---------

Co-authored-by: Vaibhav Singhal <vaibsing@vaibsing-mn7618.linkedin.biz>

Author: 2 people

gobblin-api/src/main/java/org/apache/gobblin/configuration/ConfigurationKeys.java

@@ -328,6 +328,9 @@ public class ConfigurationKeys {
   public static final String DATASET_URN_KEY = "dataset.urn";
   public static final String GLOBAL_WATERMARK_DATASET_URN = "__globalDatasetWatermark";
   public static final String DEFAULT_DATASET_URN = "";
+  public static final String DATASET_QUALITY_STATUS_KEY = "dataset.quality.status";
+  public static final String ENFORCE_DATA_QUALITY_FAILURE_KEY = "data.quality.enforce.failure";
+  public static final Boolean DEFAULT_ENFORCE_DATA_QUALITY_FAILURE = false;
 
   /**
    * Work unit related configuration properties.
@@ -498,6 +501,7 @@ public class ConfigurationKeys {
    * Configuration properties used by the quality checker.
    */
   public static final String QUALITY_CHECKER_PREFIX = "qualitychecker";
+  public static final String TASK_LEVEL_POLICY_RESULT_KEY = "gobblin.task.level.policy.result";
   public static final String TASK_LEVEL_POLICY_LIST = QUALITY_CHECKER_PREFIX + ".task.policies";
   public static final String TASK_LEVEL_POLICY_LIST_TYPE = QUALITY_CHECKER_PREFIX + ".task.policy.types";
   public static final String ROW_LEVEL_POLICY_LIST = QUALITY_CHECKER_PREFIX + ".row.policies";
@@ -938,7 +942,7 @@ public class ConfigurationKeys {
   public static final Boolean DEFAULT_METRICS_REPORTING_OPENTELEMETRY_ENABLED = false;
 
   public static final String METRICS_REPORTING_OPENTELEMETRY_ENDPOINT = METRICS_REPORTING_OPENTELEMETRY_PREFIX + "endpoint";
-
+  public static final String METRICS_REPORTING_OPENTELEMETRY_FABRIC = METRICS_REPORTING_OPENTELEMETRY_CONFIGS_PREFIX + "fabric";
   // Headers to add to the OpenTelemetry HTTP Exporter, formatted as a JSON String with string keys and values
   public static final String METRICS_REPORTING_OPENTELEMETRY_HEADERS = METRICS_REPORTING_OPENTELEMETRY_PREFIX + "headers";
 

gobblin-api/src/main/java/org/apache/gobblin/qualitychecker/task/TaskLevelPolicy.java

@@ -22,7 +22,7 @@
 
 
 public abstract class TaskLevelPolicy {
-  private final State state;
+  protected final State state;
   private final Type type;
 
   public enum Type {

gobblin-core/src/main/java/org/apache/gobblin/policies/size/FileSizePolicy.java

@@ -0,0 +1,109 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements.  See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License.  You may obtain a copy of the License at
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.apache.gobblin.policies.size;
+
+import java.util.Optional;
+
+import lombok.Getter;
+import lombok.extern.slf4j.Slf4j;
+
+import org.apache.gobblin.configuration.State;
+import org.apache.gobblin.qualitychecker.task.TaskLevelPolicy;
+
+
+/**
+ * A task-level policy that checks if the bytes read matches the bytes written for a file copy operation.
+ */
+@Slf4j
+public class FileSizePolicy extends TaskLevelPolicy {
+
+  public static final String COPY_PREFIX = "gobblin.copy";
+  public static final String BYTES_READ_KEY = COPY_PREFIX + ".bytesRead";
+  public static final String BYTES_WRITTEN_KEY = COPY_PREFIX + ".bytesWritten";
+
+  public FileSizePolicy(State state, TaskLevelPolicy.Type type) {
+    super(state, type);
+  }
+
+  @Override
+  public Result executePolicy() {
+    TransferBytes transferBytes = getBytesReadAndWritten(this.state).orElse(null);
+    if (transferBytes == null) {
+      return Result.FAILED;
+    }
+    Long bytesRead = transferBytes.getBytesRead();
+    Long bytesWritten = transferBytes.getBytesWritten();
+
+    Long sizeDifference = Math.abs(bytesRead - bytesWritten);
+
+    if (sizeDifference == 0) {
+      return Result.PASSED;
+    }
+
+    log.warn("File size check failed - bytes read: {}, bytes written: {}, difference: {}", bytesRead, bytesWritten,
+        sizeDifference);
+    return Result.FAILED;
+  }
+
+  @Override
+  public String toString() {
+    TransferBytes transferBytes = getBytesReadAndWritten(this.state).orElse(null);
+    if (transferBytes != null) {
+      return String.format("FileSizePolicy [bytesRead=%s, bytesWritten=%s]", transferBytes.getBytesRead(),
+          transferBytes.getBytesWritten());
+    } else {
+      return "FileSizePolicy [bytesRead=null, bytesWritten=null]";
+    }
+  }
+
+  /**
+   * Helper class to hold transfer bytes information
+   */
+  @Getter
+  private static class TransferBytes {
+    final long bytesRead;
+    final long bytesWritten;
+
+    TransferBytes(long bytesRead, long bytesWritten) {
+      this.bytesRead = bytesRead;
+      this.bytesWritten = bytesWritten;
+    }
+  }
+
+  /**
+   * Extracts bytesRead and bytesWritten from the given state.
+   * Returns Empty Optional if parsing fails.
+   */
+  private Optional<TransferBytes> getBytesReadAndWritten(State state) {
+    String bytesReadString = state.getProp(BYTES_READ_KEY);
+    String bytesWrittenString = state.getProp(BYTES_WRITTEN_KEY);
+    if (bytesReadString == null || bytesWrittenString == null) {
+      log.error("Missing value(s): bytesReadStr={}, bytesWrittenStr={}", bytesReadString, bytesWrittenString);
+      return Optional.empty();
+    }
+    try {
+      long bytesRead = Long.parseLong(bytesReadString);
+      long bytesWritten = Long.parseLong(bytesWrittenString);
+      return Optional.of(new TransferBytes(bytesRead, bytesWritten));
+    } catch (NumberFormatException e) {
+      log.error("Invalid number format for bytesRead or bytesWritten: bytesRead='{}', bytesWritten='{}'",
+          bytesReadString, bytesWrittenString, e);
+      return Optional.empty();
+    }
+  }
+}

gobblin-core/src/main/java/org/apache/gobblin/publisher/TaskPublisher.java

@@ -17,6 +17,7 @@
 
 package org.apache.gobblin.publisher;
 
+import java.util.Set;
 import org.apache.gobblin.qualitychecker.task.TaskLevelPolicyCheckResults;
 import java.util.Map;
 
@@ -63,8 +64,8 @@ public PublisherState canPublish() throws Exception {
    * Returns true if all tests from the PolicyChecker pass, false otherwise
    */
   public boolean passedAllTests() {
-    for (Map.Entry<TaskLevelPolicy.Result, TaskLevelPolicy.Type> entry : this.results.getPolicyResults().entrySet()) {
-      if (entry.getKey().equals(TaskLevelPolicy.Result.FAILED) && entry.getValue().equals(TaskLevelPolicy.Type.FAIL)) {
+    for (Map.Entry<TaskLevelPolicy.Result, Set<TaskLevelPolicy.Type>> entry : this.results.getPolicyResults().entrySet()) {
+      if (entry.getKey().equals(TaskLevelPolicy.Result.FAILED) && entry.getValue().contains(TaskLevelPolicy.Type.FAIL)) {
         return false;
       }
     }

gobblin-core/src/main/java/org/apache/gobblin/qualitychecker/DataQualityStatus.java

@@ -0,0 +1,49 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements.  See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License.  You may obtain a copy of the License at
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.apache.gobblin.qualitychecker;
+
+import lombok.extern.slf4j.Slf4j;
+
+
+/**
+ * An enumeration for possible statuses for Data quality checks.
+ * Its values will be:
+ * - PASSED: When all data quality checks pass
+ * - FAILED: When any data quality check fails
+ * - NOT_EVALUATED: When data quality check evaluation is not performed
+ * - UNKNOWN: when the data quality status is not recognized or invalid
+ */
+@Slf4j
+public enum DataQualityStatus {
+  PASSED,
+  FAILED,
+  NOT_EVALUATED,
+  UNKNOWN;
+
+  public static DataQualityStatus fromString(String value) {
+    if (value == null) {
+      return NOT_EVALUATED;
+    }
+    try {
+      return DataQualityStatus.valueOf(value.toUpperCase());
+    } catch (IllegalArgumentException e) {
+      log.error("Invalid DataQualityStatus value: {}. Returning UNKNOWN.", value, e);
+      return UNKNOWN;
+    }
+  }
+}

gobblin-core/src/main/java/org/apache/gobblin/qualitychecker/task/TaskLevelPolicyCheckResults.java

@@ -19,19 +19,20 @@
 
 import java.util.HashMap;
 import java.util.Map;
+import java.util.Set;
 
 
 /**
  * Wrapper around a Map of PolicyResults and Policy.Type
  */
 public class TaskLevelPolicyCheckResults {
-  private final Map<TaskLevelPolicy.Result, TaskLevelPolicy.Type> results;
+  private final Map<TaskLevelPolicy.Result, Set<TaskLevelPolicy.Type>> results;
 
   public TaskLevelPolicyCheckResults() {
     this.results = new HashMap<>();
   }
 
-  public Map<TaskLevelPolicy.Result, TaskLevelPolicy.Type> getPolicyResults() {
+  public Map<TaskLevelPolicy.Result, Set<TaskLevelPolicy.Type>> getPolicyResults() {
     return this.results;
   }
 }

gobblin-core/src/main/java/org/apache/gobblin/qualitychecker/task/TaskLevelPolicyChecker.java

@@ -17,8 +17,9 @@
 
 package org.apache.gobblin.qualitychecker.task;
 
+import java.util.EnumSet;
 import java.util.List;
-
+import lombok.Getter;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
 
@@ -28,30 +29,25 @@
  * executes each one, and then stores the output
  * in a PolicyCheckResults object
  */
+@Getter
 public class TaskLevelPolicyChecker {
-  /**
-   * An enumeration for possible statuses for Data quality checks,
-   * its values will be PASSED, FAILED, in case if data quality check
-   * evaluation is not performed for Job, it will be NOT_EVALUATED
-   */
-  public enum DataQualityStatus {
-    PASSED,
-    FAILED,
-    NOT_EVALUATED
-  }
   private final List<TaskLevelPolicy> list;
   private static final Logger LOG = LoggerFactory.getLogger(TaskLevelPolicyChecker.class);
 
+  public static final String TASK_LEVEL_POLICY_RESULT_KEY = "gobblin.task.level.policy.result";
+
   public TaskLevelPolicyChecker(List<TaskLevelPolicy> list) {
     this.list = list;
   }
 
   public TaskLevelPolicyCheckResults executePolicies() {
     TaskLevelPolicyCheckResults results = new TaskLevelPolicyCheckResults();
+
     for (TaskLevelPolicy p : this.list) {
       TaskLevelPolicy.Result result = p.executePolicy();
-      results.getPolicyResults().put(result, p.getType());
-      LOG.info("TaskLevelPolicy " + p + " of type " + p.getType() + " executed with result " + result);
+      results.getPolicyResults().computeIfAbsent(result, r -> EnumSet.noneOf(TaskLevelPolicy.Type.class))
+          .add(p.getType());
+      LOG.info("TaskLevelPolicy {} of type {} executed with result {}", p, p.getType(), result);
     }
     return results;
   }

gobblin-core/src/test/java/org/apache/gobblin/policies/size/FileSizePolicyTest.java

@@ -0,0 +1,72 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements.  See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License.  You may obtain a copy of the License at
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.apache.gobblin.policies.size;
+
+import org.apache.gobblin.configuration.State;
+import org.apache.gobblin.qualitychecker.task.TaskLevelPolicy;
+import org.testng.Assert;
+import org.testng.annotations.Test;
+
+public class FileSizePolicyTest {
+
+  @Test
+  public void testPolicyPass() {
+    State state = new State();
+    state.setProp(FileSizePolicy.BYTES_READ_KEY, 1000L);
+    state.setProp(FileSizePolicy.BYTES_WRITTEN_KEY, 1000L);
+
+    FileSizePolicy policy = new FileSizePolicy(state, TaskLevelPolicy.Type.FAIL);
+    Assert.assertEquals(policy.executePolicy(), TaskLevelPolicy.Result.PASSED);
+  }
+
+  @Test
+  public void testPolicyFail() {
+    State state = new State();
+    state.setProp(FileSizePolicy.BYTES_READ_KEY, 1000L);
+    state.setProp(FileSizePolicy.BYTES_WRITTEN_KEY, 900L);
+
+    FileSizePolicy policy = new FileSizePolicy(state, TaskLevelPolicy.Type.FAIL);
+    Assert.assertEquals(policy.executePolicy(), TaskLevelPolicy.Result.FAILED);
+  }
+
+  @Test
+  public void testMissingProperties() {
+    State state = new State();
+    // No properties set at all
+    FileSizePolicy policy = new FileSizePolicy(state, TaskLevelPolicy.Type.FAIL);
+    Assert.assertEquals(policy.executePolicy(), TaskLevelPolicy.Result.FAILED);
+  }
+
+  @Test
+  public void testPartiallySetProperties() {
+    State state = new State();
+    // Only set bytes read, not bytes written
+    state.setProp(FileSizePolicy.BYTES_READ_KEY, 1000L);
+
+    FileSizePolicy policy = new FileSizePolicy(state, TaskLevelPolicy.Type.FAIL);
+    Assert.assertEquals(policy.executePolicy(), TaskLevelPolicy.Result.FAILED);
+
+    // Reset state and only set bytes written, not bytes read
+    state = new State();
+    state.setProp(FileSizePolicy.BYTES_WRITTEN_KEY, 1000L);
+
+    policy = new FileSizePolicy(state, TaskLevelPolicy.Type.FAIL);
+    Assert.assertEquals(policy.executePolicy(), TaskLevelPolicy.Result.FAILED);
+  }
+
+}

gobblin-core/src/test/java/org/apache/gobblin/qualitychecker/FailingTaskLevelPolicy.java

@@ -0,0 +1,33 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements.  See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License.  You may obtain a copy of the License at
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.apache.gobblin.qualitychecker;
+
+import org.apache.gobblin.configuration.State;
+import org.apache.gobblin.qualitychecker.task.TaskLevelPolicy;
+
+
+public class FailingTaskLevelPolicy extends TaskLevelPolicy {
+  public FailingTaskLevelPolicy(State state, Type type) {
+    super(state, type);
+  }
+
+  @Override
+  public Result executePolicy() {
+    return Result.FAILED;
+  }
+}