/*
 * Copyright © 2019 Cask Data, 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 io.cdap.plugin.salesforce.plugin.source.batch;

import com.google.common.collect.ImmutableMap;
import com.google.gson.Gson;
import io.cdap.cdap.api.data.batch.InputFormatProvider;
import io.cdap.plugin.salesforce.SalesforceConstants;
import io.cdap.plugin.salesforce.plugin.OAuthInfo;
import io.cdap.plugin.salesforce.plugin.source.batch.util.SalesforceSourceConstants;

import java.util.List;
import java.util.Map;
import java.util.Objects;
import javax.annotation.Nullable;

/**
 * InputFormatProvider used by cdap to provide configurations to mapreduce job
 */
public class SalesforceInputFormatProvider implements InputFormatProvider {

  private static final Gson GSON = new Gson();

  private final Map<String, String> conf;

  public SalesforceInputFormatProvider(SalesforceBaseSourceConfig config,
                                       List<String> queries,
                                       Map<String, String> schemas,
                                       @Nullable String sObjectNameField) {
    ImmutableMap.Builder<String, String> configBuilder = new ImmutableMap.Builder<String, String>()
      .put(SalesforceSourceConstants.CONFIG_QUERIES, GSON.toJson(queries))
      .put(SalesforceSourceConstants.CONFIG_SCHEMAS, GSON.toJson(schemas));

    OAuthInfo oAuthInfo = config.getOAuthInfo();
    if (oAuthInfo != null) {
      configBuilder
        .put(SalesforceConstants.CONFIG_OAUTH_TOKEN, oAuthInfo.getAccessToken())
        .put(SalesforceConstants.CONFIG_OAUTH_INSTANCE_URL, oAuthInfo.getInstanceURL());
    } else {
      configBuilder
        .put(SalesforceConstants.CONFIG_USERNAME, Objects.requireNonNull(config.getUsername()))
        .put(SalesforceConstants.CONFIG_PASSWORD, Objects.requireNonNull(config.getPassword()))
        .put(SalesforceConstants.CONFIG_CONSUMER_KEY, Objects.requireNonNull(config.getConsumerKey()))
        .put(SalesforceConstants.CONFIG_CONSUMER_SECRET, Objects.requireNonNull(config.getConsumerSecret()))
        .put(SalesforceConstants.CONFIG_LOGIN_URL, Objects.requireNonNull(config.getLoginUrl()));
    }

    if (sObjectNameField != null) {
      configBuilder.put(SalesforceSourceConstants.CONFIG_SOBJECT_NAME_FIELD, sObjectNameField);
    }

    this.conf = configBuilder.build();
  }

  public SalesforceInputFormatProvider(SalesforceSourceConfig config,
                                       List<String> queries,
                                       Map<String, String> schemas,
                                       @Nullable String sObjectNameField) {
    ImmutableMap.Builder<String, String> builder = new ImmutableMap.Builder<String, String>()
      .put(SalesforceConstants.CONFIG_USERNAME, config.getUsername())
      .put(SalesforceConstants.CONFIG_PASSWORD, config.getPassword())
      .put(SalesforceConstants.CONFIG_CONSUMER_KEY, config.getConsumerKey())
      .put(SalesforceConstants.CONFIG_CONSUMER_SECRET, config.getConsumerSecret())
      .put(SalesforceConstants.CONFIG_LOGIN_URL, config.getLoginUrl())
      .put(SalesforceSourceConstants.CONFIG_QUERIES, GSON.toJson(queries))
      .put(SalesforceSourceConstants.CONFIG_SCHEMAS, GSON.toJson(schemas))
      .put(SalesforceSourceConstants.CONFIG_PK_CHUNK_ENABLE, String.valueOf(config.getEnablePKChunk()))
      .put(SalesforceSourceConstants.CONFIG_CHUNK_SIZE, String.valueOf(config.getChunkSize()))
      .put(SalesforceSourceConstants.CONFIG_CHUNK_PARENT, config.getParent());

    if (sObjectNameField != null) {
      builder.put(SalesforceSourceConstants.CONFIG_SOBJECT_NAME_FIELD, sObjectNameField);
    }

    this.conf = builder.build();
  }

  @Override
  public Map<String, String> getInputFormatConfiguration() {
    return conf;
  }

  @Override
  public String getInputFormatClassName() {
    return SalesforceInputFormat.class.getName();
  }
}

@@ -2819,938 +2819,155 @@
     
-this.conf = configBuilder.build();%0A  %7D%0A%0A  public SalesforceInputFormatProvider(SalesforceSourceConfig config,%0A                                       List%3CString%3E queries,%0A                                       Map%3CString, String%3E schemas,%0A                                       @Nullable String sObjectNameField) %7B%0A    ImmutableMap.Builder%3CString, String%3E builder = new ImmutableMap.Builder%3CString, String%3E()%0A      .put(SalesforceConstants.CONFIG_USERNAME, config.getUsername())%0A      .put(SalesforceConstants.CONFIG_PASSWORD, config.getPassword())%0A      .put(SalesforceConstants.CONFIG_CONSUMER_KEY, config.getConsumerKey())%0A      .put(SalesforceConstants.CONFIG_CONSUMER_SECRET, config.getConsumerSecret())%0A      .put(SalesforceConstants.CONFIG_LOGIN_URL, config.getLoginUrl())%0A      .put(SalesforceSourceConstants.CONFIG_QUERIES, GSON.toJson(queries))%0A      .put(SalesforceSourceConstants.CONFIG_SCHEMAS, GSON.toJson(schemas))%0A
+if (config instanceof SalesforceSourceConfig) %7B%0A      SalesforceSourceConfig sourceConfig = (SalesforceSourceConfig) config;%0A      configBuilder%0A  
     
@@ -3030,33 +3030,39 @@
  String.valueOf(
-c
+sourceC
 onfig.getEnableP
@@ -3068,24 +3068,26 @@
 PKChunk()))%0A
+  
       .put(S
@@ -3145,17 +3145,23 @@
 valueOf(
-c
+sourceC
 onfig.ge
@@ -3171,24 +3171,26 @@
 unkSize()))%0A
+  
       .put(S
@@ -3235,17 +3235,23 @@
 PARENT, 
-c
+sourceC
 onfig.ge
@@ -3265,135 +3265,8 @@
 ());
-%0A%0A    if (sObjectNameField != null) %7B%0A      builder.put(SalesforceSourceConstants.CONFIG_SOBJECT_NAME_FIELD, sObjectNameField);
 %0A   
@@ -3285,17 +3285,23 @@
 .conf = 
-b
+configB
 uilder.b

Fix the OAuth support that was broken when chunking was added