// Path: src/java/com/google/cloud/bigquery/dwhassessment/extractiontool/db/AvroHelper.java
// public static GenericRecord parseRowToAvro(ResultSet row, Schema schema) throws SQLException {
//   GenericRecordBuilder recordBuilder = new GenericRecordBuilder(schema);
//   ResultSetMetaData metaData = row.getMetaData();
//   for (int columnIndex = 1; columnIndex <= row.getMetaData().getColumnCount(); columnIndex++) {
//     recordBuilder.set(
//         metaData.getColumnName(columnIndex), getRowObject(metaData, row, columnIndex));
//   }
//   return recordBuilder.build();
// }
// 
// Path: src/java/com/google/cloud/bigquery/dwhassessment/extractiontool/dumper/DataEntityManager.java
// public interface DataEntityManager extends Closeable {
// 
//   /**
//    * Get the output stream for an entity.
//    *
//    * @param name The name of the output entity.
//    */
//   OutputStream getEntityOutputStream(String name) throws IOException;
// 
//   /**
//    * Indicate whether the data entity allows resumable processing.
//    *
//    * @return true if some progress can be retained after the writing is interrupted.
//    */
//   boolean isResumable();
// 
//   /**
//    * Get the absolute path of a file given its name.
//    *
//    * @param name The name of the file.
//    */
//   Path getAbsolutePath(String name);
// }


import static com.google.cloud.bigquery.dwhassessment.extractiontool.db.AvroHelper.parseRowToAvro;
import static java.nio.file.StandardCopyOption.ATOMIC_MOVE;
import com.google.cloud.bigquery.dwhassessment.extractiontool.dumper.DataEntityManager;
import com.google.common.annotations.VisibleForTesting;
import com.google.common.base.Preconditions;
import com.google.common.collect.ImmutableList;
import com.google.common.collect.ImmutableMap;
import com.google.common.collect.ImmutableSet;
import java.io.IOException;
import java.nio.file.Files;
import java.sql.Connection;
import java.sql.ResultSet;
import java.sql.SQLException;
import java.sql.Timestamp;
import java.time.Instant;
import java.time.temporal.ChronoUnit;
import java.util.Locale;
import java.util.function.Supplier;
import java.util.logging.Logger;
import org.apache.avro.Schema;
import org.apache.avro.generic.GenericRecord;
/*
 * Copyright 2021 Google LLC
 *
 * 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
 *
 *     https://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.bigquery.dwhassessment.extractiontool.db;



/**
 * Implementation of script manager. Manages mapping from script name to SQL script. Executes script
 * and writes results to an output stream.
 */
public class ScriptManagerImpl implements ScriptManager {

  private static final Logger LOGGER = Logger.getLogger(ScriptManagerImpl.class.getName());

  private final ImmutableMap<String, Supplier<String>> scriptsMap;
  private final ImmutableMap<String, ImmutableList<String>> sortingColumnsMap;
  private final ScriptRunner scriptRunner;

  public ScriptManagerImpl(
      ScriptRunner scriptRunner,
      ImmutableMap<String, Supplier<String>> scriptsMap,
      ImmutableMap<String, ImmutableList<String>> sortingColumnsMap) {
    this.scriptRunner = scriptRunner;
    this.scriptsMap = scriptsMap;
    this.sortingColumnsMap = sortingColumnsMap;
  }

  @Override
  public void executeScript(
      Connection connection,
      boolean dryRun,
      SqlTemplateRenderer sqlTemplateRenderer,
      String scriptName,
// Path: src/java/com/google/cloud/bigquery/dwhassessment/extractiontool/db/AvroHelper.java
// public static GenericRecord parseRowToAvro(ResultSet row, Schema schema) throws SQLException {
//   GenericRecordBuilder recordBuilder = new GenericRecordBuilder(schema);
//   ResultSetMetaData metaData = row.getMetaData();
//   for (int columnIndex = 1; columnIndex <= row.getMetaData().getColumnCount(); columnIndex++) {
//     recordBuilder.set(
//         metaData.getColumnName(columnIndex), getRowObject(metaData, row, columnIndex));
//   }
//   return recordBuilder.build();
// }
// 
// Path: src/java/com/google/cloud/bigquery/dwhassessment/extractiontool/dumper/DataEntityManager.java
// public interface DataEntityManager extends Closeable {
// 
//   /**
//    * Get the output stream for an entity.
//    *
//    * @param name The name of the output entity.
//    */
//   OutputStream getEntityOutputStream(String name) throws IOException;
// 
//   /**
//    * Indicate whether the data entity allows resumable processing.
//    *
//    * @return true if some progress can be retained after the writing is interrupted.
//    */
//   boolean isResumable();
// 
//   /**
//    * Get the absolute path of a file given its name.
//    *
//    * @param name The name of the file.
//    */
//   Path getAbsolutePath(String name);
// }


// Path: src/java/com/google/cloud/bigquery/dwhassessment/extractiontool/db/ScriptManagerImpl.java
import static com.google.cloud.bigquery.dwhassessment.extractiontool.db.AvroHelper.parseRowToAvro;
import static java.nio.file.StandardCopyOption.ATOMIC_MOVE;
import com.google.cloud.bigquery.dwhassessment.extractiontool.dumper.DataEntityManager;
import com.google.common.annotations.VisibleForTesting;
import com.google.common.base.Preconditions;
import com.google.common.collect.ImmutableList;
import com.google.common.collect.ImmutableMap;
import com.google.common.collect.ImmutableSet;
import java.io.IOException;
import java.nio.file.Files;
import java.sql.Connection;
import java.sql.ResultSet;
import java.sql.SQLException;
import java.sql.Timestamp;
import java.time.Instant;
import java.time.temporal.ChronoUnit;
import java.util.Locale;
import java.util.function.Supplier;
import java.util.logging.Logger;
import org.apache.avro.Schema;
import org.apache.avro.generic.GenericRecord;

/*
 * Copyright 2021 Google LLC
 *
 * 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
 *
 *     https://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.bigquery.dwhassessment.extractiontool.db;



/**
 * Implementation of script manager. Manages mapping from script name to SQL script. Executes script
 * and writes results to an output stream.
 */
public class ScriptManagerImpl implements ScriptManager {

  private static final Logger LOGGER = Logger.getLogger(ScriptManagerImpl.class.getName());

  private final ImmutableMap<String, Supplier<String>> scriptsMap;
  private final ImmutableMap<String, ImmutableList<String>> sortingColumnsMap;
  private final ScriptRunner scriptRunner;

  public ScriptManagerImpl(
      ScriptRunner scriptRunner,
      ImmutableMap<String, Supplier<String>> scriptsMap,
      ImmutableMap<String, ImmutableList<String>> sortingColumnsMap) {
    this.scriptRunner = scriptRunner;
    this.scriptsMap = scriptsMap;
    this.sortingColumnsMap = sortingColumnsMap;
  }

  @Override
  public void executeScript(
      Connection connection,
      boolean dryRun,
      SqlTemplateRenderer sqlTemplateRenderer,
      String scriptName,
      DataEntityManager dataEntityManager,