unit test fix

data-integrations · Jan 27, 2025 · c682a83 · c682a83
1 parent 4975d27
commit c682a83
Show file tree

Hide file tree

Showing 3 changed files with 10 additions and 7 deletions.
diff --git a/src/main/java/io/cdap/plugin/gcp/bigquery/action/BigQueryArgumentSetter.java b/src/main/java/io/cdap/plugin/gcp/bigquery/action/BigQueryArgumentSetter.java
@@ -116,8 +116,9 @@ public void run(ActionContext context) {
     // Check for errors
     if (queryJob.getStatus().getError() != null) {
       String error = queryJob.getStatus().getExecutionErrors().toString();
-      throw BigQueryErrorUtil.getProgramFailureException(error,
-          queryJob.getStatus().getError().getReason(), null);
+      ErrorType type = BigQueryErrorUtil.getErrorType(queryJob.getStatus().getError().getReason());
+      throw ErrorUtils.getProgramFailureException(new ErrorCategory(ErrorCategory.ErrorCategoryEnum.PLUGIN),
+          error, error, type, true, null);
     }
     TableResult queryResults;
     try {

diff --git a/src/main/java/io/cdap/plugin/gcp/bigquery/action/BigQueryExecute.java b/src/main/java/io/cdap/plugin/gcp/bigquery/action/BigQueryExecute.java
@@ -190,8 +190,8 @@ public void run(ActionContext context) {
     }
   }
 
-  private void executeQueryWithExponentialBackoff(BigQuery bigQuery,
-                                          QueryJobConfiguration queryConfig, ActionContext context) {
+  protected void executeQueryWithExponentialBackoff(BigQuery bigQuery,
+      QueryJobConfiguration queryConfig, ActionContext context) {
     try {
       Failsafe.with(getRetryPolicy()).run(() -> executeQuery(bigQuery, queryConfig, context));
     } catch (FailsafeException e) {
@@ -262,8 +262,10 @@ private void executeQuery(BigQuery bigQuery, QueryJobConfiguration queryConfig,
       String error = String.format("Failed to execute query with reason: %s and message: %s",
           queryJob.getStatus().getError().getReason(),
           queryJob.getStatus().getError().getMessage());
-      throw BigQueryErrorUtil.getProgramFailureException(error,
-          queryJob.getStatus().getError().getReason(), null);
+      ErrorType type = BigQueryErrorUtil.getErrorType(queryJob.getStatus().getError().getReason());
+      throw ErrorUtils.getProgramFailureException(
+          new ErrorCategory(ErrorCategory.ErrorCategoryEnum.PLUGIN), error, error, type, true,
+          null);
     }
 
     TableResult queryResults;

diff --git a/src/main/java/io/cdap/plugin/gcp/bigquery/common/BigQueryErrorUtil.java b/src/main/java/io/cdap/plugin/gcp/bigquery/common/BigQueryErrorUtil.java
@@ -70,7 +70,7 @@ public class BigQueryErrorUtil {
    * @param errorReason the error reason to classify
    * @return the corresponding ErrorType (USER, SYSTEM, UNKNOWN)
    */
-  private static ErrorType getErrorType(String errorReason) {
+  public static ErrorType getErrorType(String errorReason) {
     if (errorReason != null && ERROR_REASON_TO_ERROR_TYPE.containsKey(errorReason)) {
       return ERROR_REASON_TO_ERROR_TYPE.get(errorReason);
     }