Adds a "read from query results" example

flink-connector-bigquery-examples/src/main/java/com/google/cloud/flink/bigquery/examples/BigQueryExample.java

@@ -34,20 +34,34 @@
 /**
  * A simple Flink application using DataStream API and BigQuery connector.
  *
- * <p>The Flink pipeline will try to read the specified BigQuery table, limiting the element count
- * to the specified row restriction and limit, returning {@link GenericRecord} representing the
- * rows, and finally print out some aggregated values given the provided payload's field. The
- * sequence of operations in this pipeline is: <i>source > flatMap > keyBy > max > print</i>.
+ * <p>The Flink pipeline will try to read the specified BigQuery table according to given the
+ * command line arguments, returning {@link GenericRecord} representing the rows, and finally print
+ * out some aggregated values given the provided payload's field. The sequence of operations in this
+ * pipeline is: <i>source > flatMap > keyBy > sum > print</i>.
  *
- * <p>Flink command line format to execute this application: <br>
- * flink run {additional runtime params} {path to this jar}/BigQueryExample.jar <br>
- * --gcp-project {required; project ID which contains the BigQuery table} <br>
- * --bq-dataset {required; name of BigQuery dataset containing the desired table} <br>
- * --bq-table {required; name of BigQuery table to read} <br>
- * --agg-prop {required; record property to aggregate in Flink job} <br>
- * --restriction {optional; SQL-like filter applied at the BigQuery table before reading} <br>
- * --limit {optional; maximum records to read from BigQuery table} <br>
- * --checkpoint-interval {optional; time interval between state checkpoints in milliseconds}
+ * <p>This example module should be used in one of the following two ways.
+ *
+ * <ol>
+ *   <li>Specify the BQ dataset and table with an optional row restriction. Flink command line
+ *       format to execute this mode is: <br>
+ *       flink run {additional runtime params} {path to this jar}/BigQueryExample.jar <br>
+ *       --gcp-project {required; project ID which contains the BigQuery table} <br>
+ *       --bq-dataset {required; name of BigQuery dataset containing the desired table} <br>
+ *       --bq-table {required; name of BigQuery table to read} <br>
+ *       --agg-prop {required; record property to aggregate in Flink job} <br>
+ *       --restriction {optional; SQL filter applied at the BigQuery table before reading} <br>
+ *       --limit {optional; maximum records to read from BigQuery table} <br>
+ *       --checkpoint-interval {optional; time interval between state checkpoints in milliseconds}
+ *   <li>Specify SQL query to fetch data from BQ dataset. For example, "SELECT * FROM
+ *       some_dataset.INFORMATION_SCHEMA.PARTITIONS". Flink command line format to execute this mode
+ *       is: <br>
+ *       flink run {additional runtime params} {path to this jar}/BigQueryExample.jar <br>
+ *       --gcp-project {required; project ID which contains the BigQuery table} <br>
+ *       --query {required; SQL query to fetch data from BigQuery table} <br>
+ *       --agg-prop {required; record property to aggregate in Flink job} <br>
+ *       --limit {optional; maximum records to read from BigQuery table} <br>
+ *       --checkpoint-interval {optional; time interval between state checkpoints in milliseconds}
+ * </ol>
  *
  * <p>Note on row restriction: In case a restriction relies on temporal reference, something like
  * {@code "TIMESTAMP_TRUNC(ingestion_timestamp, HOUR) = '2023-06-20 19:00:00'"}, and if launching
@@ -63,36 +77,73 @@ public static void main(String[] args) throws Exception {
         // parse input arguments
         final ParameterTool parameterTool = ParameterTool.fromArgs(args);
 
-        if (parameterTool.getNumberOfParameters() < 4) {
+        if (parameterTool.getNumberOfParameters() < 1) {
             LOG.error(
                     "Missing parameters!\n"
-                            + "Usage: flink run <additional runtime params> BigQuery.jar"
-                            + " --gcp-project <gcp-project>"
+                            + "Usage: flink run <additional runtime params> <jar>"
+                            + " --gcp-project <gcp project id>"
                             + " --bq-dataset <dataset name>"
                             + " --bq-table <table name>"
-                            + " --agg-prop <record property>"
+                            + " --agg-prop <record property to aggregate>"
                             + " --restriction <row filter predicate>"
-                            + " --limit <limit records returned>"
-                            + " --checkpoint-interval <milliseconds between state checkpoints>");
+                            + " --limit <limit on records returned>"
+                            + " --checkpoint-interval <milliseconds between state checkpoints>"
+                            + " --query <SQL query to get data from BQ table>");
             return;
         }
 
         String projectName = parameterTool.getRequired("gcp-project");
-        String datasetName = parameterTool.getRequired("bq-dataset");
-        String tableName = parameterTool.getRequired("bq-table");
-        String rowRestriction = parameterTool.get("restriction", "").replace("\\u0027", "'");
-        Integer recordLimit = parameterTool.getInt("limit", -1);
+        String query = parameterTool.get("query", "");
         String recordPropertyToAggregate = parameterTool.getRequired("agg-prop");
+        Integer recordLimit = parameterTool.getInt("limit", -1);
         Long checkpointInterval = parameterTool.getLong("checkpoint-interval", 60000L);
 
-        runFlinkJob(
-                projectName,
-                datasetName,
-                tableName,
-                recordPropertyToAggregate,
-                rowRestriction,
-                recordLimit,
-                checkpointInterval);
+        if (!query.isEmpty()) {
+            runFlinkQueryJob(
+                    projectName, query, recordPropertyToAggregate, recordLimit, checkpointInterval);
+        } else {
+            String datasetName = parameterTool.getRequired("bq-dataset");
+            String tableName = parameterTool.getRequired("bq-table");
+            String rowRestriction = parameterTool.get("restriction", "").replace("\\u0027", "'");
+
+            runFlinkJob(
+                    projectName,
+                    datasetName,
+                    tableName,
+                    recordPropertyToAggregate,
+                    rowRestriction,
+                    recordLimit,
+                    checkpointInterval);
+        }
+    }
+
+    private static void runFlinkQueryJob(
+            String projectName,
+            String query,
+            String recordPropertyToAggregate,
+            Integer limit,
+            Long checkpointInterval)
+            throws Exception {
+
+        final StreamExecutionEnvironment env = StreamExecutionEnvironment.getExecutionEnvironment();
+        env.enableCheckpointing(checkpointInterval);
+
+        /**
+         * we will be reading avro generic records from BigQuery, and in this case we are assuming
+         * the GOOGLE_APPLICATION_CREDENTIALS env variable will be present in the execution runtime.
+         * In case of needing authenticate differently, the credentials builder (part of the
+         * BigQueryConnectOptions) should enable capturing the credentials from various sources.
+         */
+        BigQuerySource<GenericRecord> bqSource =
+                BigQuerySource.readAvrosFromQuery(query, projectName, limit);
+
+        env.fromSource(bqSource, WatermarkStrategy.noWatermarks(), "BigQueryQuerySource")
+                .flatMap(new FlatMapper(recordPropertyToAggregate))
+                .keyBy(mappedTuple -> mappedTuple.f0)
+                .sum("f1")
+                .print();
+
+        env.execute("Flink BigQuery query example");
     }
 
     private static void runFlinkJob(
@@ -130,10 +181,10 @@ private static void runFlinkJob(
         env.fromSource(bqSource, WatermarkStrategy.noWatermarks(), "BigQuerySource")
                 .flatMap(new FlatMapper(recordPropertyToAggregate))
                 .keyBy(mappedTuple -> mappedTuple.f0)
-                .max("f1")
+                .sum("f1")
                 .print();
 
-        env.execute("Flink BigQuery Example");
+        env.execute("Flink BigQuery example");
     }
 
     static class FlatMapper implements FlatMapFunction<GenericRecord, Tuple2<String, Integer>> {
@@ -147,9 +198,7 @@ public FlatMapper(String recordPropertyToAggregate) {
         @Override
         public void flatMap(GenericRecord record, Collector<Tuple2<String, Integer>> out)
                 throws Exception {
-            out.collect(
-                    Tuple2.<String, Integer>of(
-                            (String) record.get(recordPropertyToAggregate).toString(), 1));
+            out.collect(Tuple2.of((String) record.get(recordPropertyToAggregate).toString(), 1));
         }
     }
 }