feat: Add ExportDataStats to QueryStatistics (#3244)

PhongChuong · web-flow · commit e91be80ebdd3 · 2024-04-17T16:04:35.000-04:00
diff --git a/google-cloud-bigquery/src/main/java/com/google/cloud/bigquery/JobStatistics.java b/google-cloud-bigquery/src/main/java/com/google/cloud/bigquery/JobStatistics.java
@@ -17,12 +17,14 @@
 package com.google.cloud.bigquery;
 
 import com.google.api.core.ApiFunction;
+import com.google.api.services.bigquery.model.ExportDataStatistics;
 import com.google.api.services.bigquery.model.JobConfiguration;
 import com.google.api.services.bigquery.model.JobStatistics2;
 import com.google.api.services.bigquery.model.JobStatistics3;
 import com.google.api.services.bigquery.model.JobStatistics4;
 import com.google.api.services.bigquery.model.JobStatistics5;
 import com.google.api.services.bigquery.model.QueryParameter;
+import com.google.auto.value.AutoValue;
 import com.google.cloud.StringEnumType;
 import com.google.cloud.StringEnumValue;
 import com.google.common.base.Function;
@@ -32,6 +34,7 @@
 import java.io.Serializable;
 import java.util.List;
 import java.util.Objects;
+import javax.annotation.Nullable;
 import org.checkerframework.checker.nullness.compatqual.NullableDecl;
 
 /** A Google BigQuery Job statistics. */
@@ -398,6 +401,7 @@ public static class QueryStatistics extends JobStatistics {
     private final Long estimatedBytesProcessed;
     private final Long numDmlAffectedRows;
     private final DmlStats dmlStats;
+    private final ExportDataStats exportDataStats;
     private final List<TableId> referencedTables;
     private final StatementType statementType;
     private final Long totalBytesBilled;
@@ -472,6 +476,80 @@ public static StatementType[] values() {
       }
     }
 
+    /**
+     * Statistics for the EXPORT DATA statement as part of Query Job. EXTRACT JOB statistics are
+     * populated in ExtractStatistics.
+     */
+    @AutoValue
+    public abstract static class ExportDataStats implements Serializable {
+      private static final long serialVersionUID = 1L;
+
+      /**
+       * Returns number of destination files generated in case of EXPORT DATA statement only.
+       *
+       * @return value or {@code null} for none
+       */
+      @Nullable
+      public abstract Long getFileCount();
+
+      /**
+       * Returns number of destination rows generated in case of EXPORT DATA statement only.
+       *
+       * @return value or {@code null} for none
+       */
+      @Nullable
+      public abstract Long getRowCount();
+
+      public abstract Builder toBuilder();
+
+      public static Builder newBuilder() {
+        return new AutoValue_JobStatistics_QueryStatistics_ExportDataStats.Builder();
+      }
+
+      static ExportDataStats fromPb(ExportDataStatistics exportDataStatisticsPb) {
+        Builder builder = newBuilder();
+        if (exportDataStatisticsPb.getFileCount() != null) {
+          builder.setFileCount(exportDataStatisticsPb.getFileCount());
+        }
+        if (exportDataStatisticsPb.getRowCount() != null) {
+          builder.setRowCount(exportDataStatisticsPb.getRowCount());
+        }
+        return builder.build();
+      }
+
+      ExportDataStatistics toPb() {
+        ExportDataStatistics exportDataStatisticsPb = new ExportDataStatistics();
+        if (getFileCount() != null) {
+          exportDataStatisticsPb.setFileCount(getFileCount());
+        }
+        if (getRowCount() != null) {
+          exportDataStatisticsPb.setRowCount(getRowCount());
+        }
+        return exportDataStatisticsPb;
+      }
+
+      @AutoValue.Builder
+      public abstract static class Builder {
+
+        /**
+         * Number of destination files generated in case of EXPORT DATA statement only.
+         *
+         * @param fileCount fileCount or {@code null} for none
+         */
+        public abstract Builder setFileCount(Long fileCount);
+
+        /**
+         * Number of destination rows generated in case of EXPORT DATA statement only.
+         *
+         * @param rowCount rowCount or {@code null} for none
+         */
+        public abstract Builder setRowCount(Long rowCount);
+
+        /** Creates a {@code ExportDataStats} object. */
+        public abstract ExportDataStats build();
+      }
+    }
+
     static final class Builder extends JobStatistics.Builder<QueryStatistics, Builder> {
 
       private BiEngineStats biEngineStats;
@@ -483,6 +561,7 @@ static final class Builder extends JobStatistics.Builder<QueryStatistics, Builde
       private Long estimatedBytesProcessed;
       private Long numDmlAffectedRows;
       private DmlStats dmlStats;
+      private ExportDataStats exportDataStats;
       private List<TableId> referencedTables;
       private StatementType statementType;
       private Long totalBytesBilled;
@@ -553,6 +632,10 @@ private Builder(com.google.api.services.bigquery.model.JobStatistics statisticsP
           if (statisticsPb.getQuery().getDmlStats() != null) {
             this.dmlStats = DmlStats.fromPb(statisticsPb.getQuery().getDmlStats());
           }
+          if (statisticsPb.getQuery().getExportDataStatistics() != null) {
+            this.exportDataStats =
+                ExportDataStats.fromPb(statisticsPb.getQuery().getExportDataStatistics());
+          }
         }
       }
 
@@ -601,6 +684,11 @@ Builder setDmlStats(DmlStats dmlStats) {
         return self();
       }
 
+      Builder setExportDataStats(ExportDataStats exportDataStats) {
+        this.exportDataStats = exportDataStats;
+        return self();
+      }
+
       Builder setReferenceTables(List<TableId> referencedTables) {
         this.referencedTables = referencedTables;
         return self();
@@ -683,6 +771,7 @@ private QueryStatistics(Builder builder) {
       this.estimatedBytesProcessed = builder.estimatedBytesProcessed;
       this.numDmlAffectedRows = builder.numDmlAffectedRows;
       this.dmlStats = builder.dmlStats;
+      this.exportDataStats = builder.exportDataStats;
       this.referencedTables = builder.referencedTables;
       this.statementType = builder.statementType;
       this.totalBytesBilled = builder.totalBytesBilled;
@@ -749,6 +838,11 @@ public DmlStats getDmlStats() {
       return dmlStats;
     }
 
+    /** Detailed statistics for EXPORT DATA statement. */
+    public ExportDataStats getExportDataStats() {
+      return exportDataStats;
+    }
+
     /**
      * Referenced tables for the job. Queries that reference more than 50 tables will not have a
      * complete list.
@@ -900,6 +994,9 @@ com.google.api.services.bigquery.model.JobStatistics toPb() {
       if (dmlStats != null) {
         queryStatisticsPb.setDmlStats(dmlStats.toPb());
       }
+      if (exportDataStats != null) {
+        queryStatisticsPb.setExportDataStatistics(exportDataStats.toPb());
+      }
       if (referencedTables != null) {
         queryStatisticsPb.setReferencedTables(
             Lists.transform(referencedTables, TableId.TO_PB_FUNCTION));
diff --git a/google-cloud-bigquery/src/test/java/com/google/cloud/bigquery/JobStatisticsTest.java b/google-cloud-bigquery/src/test/java/com/google/cloud/bigquery/JobStatisticsTest.java
@@ -23,6 +23,7 @@
 import com.google.cloud.bigquery.JobStatistics.ExtractStatistics;
 import com.google.cloud.bigquery.JobStatistics.LoadStatistics;
 import com.google.cloud.bigquery.JobStatistics.QueryStatistics;
+import com.google.cloud.bigquery.JobStatistics.QueryStatistics.ExportDataStats;
 import com.google.cloud.bigquery.JobStatistics.ReservationUsage;
 import com.google.cloud.bigquery.JobStatistics.ScriptStatistics;
 import com.google.cloud.bigquery.JobStatistics.ScriptStatistics.ScriptStackFrame;
@@ -64,6 +65,13 @@ public class JobStatisticsTest {
           .setInsertedRowCount(INSERTED_ROW_COUNT)
           .setUpdatedRowCount(UPDATED_ROW_COUNT)
           .build();
+  private static final Long EXPORT_DATA_STATS_ROW_COUNT = 3L;
+  private static final Long EXPORT_DATA_STATS_FILE_COUNT = 2L;
+  private static final ExportDataStats EXPORT_DATA_STATS =
+      ExportDataStats.newBuilder()
+          .setRowCount(EXPORT_DATA_STATS_ROW_COUNT)
+          .setFileCount(EXPORT_DATA_STATS_FILE_COUNT)
+          .build();
   private static final QueryStatistics.StatementType STATEMENT_TYPE =
       QueryStatistics.StatementType.SELECT;
   private static final Long TOTAL_BYTES_BILLED = 24L;
@@ -189,6 +197,7 @@ public class JobStatisticsTest {
           .setEstimatedBytesProcessed(ESTIMATE_BYTES_PROCESSED)
           .setNumDmlAffectedRows(NUM_DML_AFFECTED_ROWS)
           .setDmlStats(DML_STATS)
+          .setExportDataStats(EXPORT_DATA_STATS)
           .setReferenceTables(REFERENCED_TABLES)
           .setStatementType(STATEMENT_TYPE)
           .setTotalBytesBilled(TOTAL_BYTES_BILLED)
@@ -293,6 +302,7 @@ public void testBuilder() {
     assertEquals(ESTIMATE_BYTES_PROCESSED, QUERY_STATISTICS.getEstimatedBytesProcessed());
     assertEquals(NUM_DML_AFFECTED_ROWS, QUERY_STATISTICS.getNumDmlAffectedRows());
     assertEquals(DML_STATS, QUERY_STATISTICS.getDmlStats());
+    assertEquals(EXPORT_DATA_STATS, QUERY_STATISTICS.getExportDataStats());
     assertEquals(REFERENCED_TABLES, QUERY_STATISTICS.getReferencedTables());
     assertEquals(STATEMENT_TYPE, QUERY_STATISTICS.getStatementType());
     assertEquals(TOTAL_BYTES_BILLED, QUERY_STATISTICS.getTotalBytesBilled());
@@ -448,6 +458,8 @@ private void compareQueryStatistics(QueryStatistics expected, QueryStatistics va
     assertEquals(expected.getMetadataCacheStats(), value.getMetadataCacheStats());
     assertEquals(expected.getStatementType(), value.getStatementType());
     assertEquals(expected.getTimeline(), value.getTimeline());
+    assertEquals(expected.getDmlStats(), value.getDmlStats());
+    assertEquals(expected.getExportDataStats(), value.getExportDataStats());
   }
 
   private void compareStatistics(JobStatistics expected, JobStatistics value) {
diff --git a/google-cloud-bigquery/src/test/java/com/google/cloud/bigquery/it/ITBigQueryTest.java b/google-cloud-bigquery/src/test/java/com/google/cloud/bigquery/it/ITBigQueryTest.java
@@ -6734,4 +6734,28 @@ static GoogleCredentials loadCredentials(String credentialFile) {
     }
     return null;
   }
+
+  @Test
+  public void testQueryExportStatistics() throws InterruptedException {
+    String query =
+        String.format(
+            "EXPORT DATA OPTIONS(\n"
+                + "  uri='gs://%s/*.csv',\n"
+                + "  format='CSV',\n"
+                + "  overwrite=true,\n"
+                + "  header=true,\n"
+                + "  field_delimiter=';') AS\n"
+                + "SELECT num FROM UNNEST([1,2,3]) AS num",
+            BUCKET);
+    QueryJobConfiguration config =
+        QueryJobConfiguration.newBuilder(query).setDefaultDataset(DatasetId.of(DATASET)).build();
+    Job job = bigquery.create(JobInfo.of(JobId.of(), config));
+    job = job.waitFor();
+
+    QueryStatistics queryStatistics = job.getStatistics();
+    assertNotNull(queryStatistics);
+    assertNotNull(queryStatistics.getExportDataStats());
+    assertEquals(1L, queryStatistics.getExportDataStats().getFileCount().longValue());
+    assertEquals(3L, queryStatistics.getExportDataStats().getRowCount().longValue());
+  }
 }
