/*
 * Copyright (C) 2018 Google Inc.
 *
 * Licensed 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 com.google.cloud.teleport.spanner;

import com.google.cloud.teleport.spanner.ddl.Ddl;
import com.google.cloud.teleport.spanner.ddl.Table;
import java.util.stream.Collectors;
import org.apache.beam.sdk.io.gcp.spanner.ReadOperation;
import org.apache.beam.sdk.transforms.DoFn;
import org.apache.beam.sdk.transforms.PTransform;
import org.apache.beam.sdk.transforms.ParDo;
import org.apache.beam.sdk.values.PCollection;

/** Given a Cloud Spanner {@link Ddl} generates a "read all" operation per table. */
class BuildReadFromTableOperations
    extends PTransform<PCollection<Ddl>, PCollection<ReadOperation>> {

  @Override
  public PCollection<ReadOperation> expand(PCollection<Ddl> ddl) {
    return ddl.apply(
        "Read from table operations",
        ParDo.of(
            new DoFn<Ddl, ReadOperation>() {

              @ProcessElement
              public void processElement(ProcessContext c) {
                Ddl ddl = c.element();
                for (Table table : ddl.allTables()) {
                  String columnsListAsString =
                      table.columns().stream()
                          .map(x -> "t.`" + x.name() + "`")
                          .collect(Collectors.joining(","));
                  // Also have to export table name to be able to identify which row belongs to
                  // which table.
                  ReadOperation read =
                      ReadOperation.create()
                          .withQuery(
                              String.format(
                                  "SELECT \"%s\" AS _spanner_table, %s FROM `%s` AS t",
                                  table.name(), columnsListAsString, table.name()));
                  c.output(read);
                }
              }
            }));
  }
}

@@ -633,16 +633,66 @@
 anner;%0A%0A
+import com.google.cloud.spanner.PartitionOptions;%0A
 import c
@@ -733,16 +733,16 @@
 dl.Ddl;%0A
-
 import c
@@ -1257,16 +1257,263 @@
 on%3E%3E %7B%0A%0A
+  // The number of read partitions have to be capped so that in case the Partition token is large%0A  // (which can happen with a table with a lot of columns), the PartitionResponse size is bounded.%0A  private static final int maxPartitions = 1000;%0A%0A
   @Overr
@@ -2102,24 +2102,173 @@
 ning(%22,%22));%0A
+%0A                  PartitionOptions partitionOptions =%0A                      PartitionOptions.newBuilder().setMaxPartitions(maxPartitions).build();%0A%0A
             
@@ -2640,16 +2640,16 @@
  AS t%22,%0A
-
         
@@ -2723,16 +2723,82 @@
 name()))
+%0A                          .withPartitionOptions(partitionOptions)
 ;%0A      

Use a PartitionOption of maxPartitions=1000 to bound the size of the PartitionResponse which was causing issues with export of tables with large schema.