Prepare for Hadoop authentication extensions

4 files changed, 34 insertions, 21 deletions

diff --git a/vespa-hadoop/src/main/java/com/yahoo/vespa/hadoop/mapreduce/VespaOutputFormat.java b/vespa-hadoop/src/main/java/com/yahoo/vespa/hadoop/mapreduce/VespaOutputFormat.java
index b4056fa4665..bef51e9ae08 100644
--- a/vespa-hadoop/src/main/java/com/yahoo/vespa/hadoop/mapreduce/VespaOutputFormat.java
+++ b/vespa-hadoop/src/main/java/com/yahoo/vespa/hadoop/mapreduce/VespaOutputFormat.java
@@ -18,7 +18,7 @@ import java.util.Properties;
 @SuppressWarnings("rawtypes")
 public class VespaOutputFormat extends OutputFormat {
 
-    private final Properties configOverride;
+    final Properties configOverride;
 
     public VespaOutputFormat() {
         super();
diff --git a/vespa-hadoop/src/main/java/com/yahoo/vespa/hadoop/mapreduce/VespaRecordWriter.java b/vespa-hadoop/src/main/java/com/yahoo/vespa/hadoop/mapreduce/VespaRecordWriter.java
index 5e03a11a356..6a99a7bc08b 100644
--- a/vespa-hadoop/src/main/java/com/yahoo/vespa/hadoop/mapreduce/VespaRecordWriter.java
+++ b/vespa-hadoop/src/main/java/com/yahoo/vespa/hadoop/mapreduce/VespaRecordWriter.java
@@ -50,11 +50,11 @@ public class VespaRecordWriter extends RecordWriter {
 
     private boolean initialized = false;
     private FeedClient feedClient;
-
     private final VespaCounters counters;
-    private final VespaConfiguration configuration;
     private final int progressInterval;
 
+    final VespaConfiguration configuration;
+
     VespaRecordWriter(VespaConfiguration configuration, VespaCounters counters) {
         this.counters = counters;
         this.configuration = configuration;
@@ -99,16 +99,7 @@ public class VespaRecordWriter extends RecordWriter {
         }
     }
 
-
-    private void initialize() {
-        if (!configuration.dryrun() && configuration.randomSartupSleepMs() > 0) {
-            int delay = new Random().nextInt(configuration.randomSartupSleepMs());
-            log.info("VespaStorage: Delaying startup by " + delay + " ms");
-            try {
-                Thread.sleep(delay);
-            } catch (Exception e) {}
-        }
-
+    protected ConnectionParams.Builder configureConnectionParams() {
         ConnectionParams.Builder connParamsBuilder = new ConnectionParams.Builder();
         connParamsBuilder.setDryRun(configuration.dryrun());
         connParamsBuilder.setUseCompression(configuration.useCompression());
@@ -121,20 +112,41 @@ public class VespaRecordWriter extends RecordWriter {
         if (configuration.proxyPort() >= 0) {
             connParamsBuilder.setProxyPort(configuration.proxyPort());
         }
+        return connParamsBuilder;
+    }
 
+    protected FeedParams.Builder configureFeedParams() {
         FeedParams.Builder feedParamsBuilder = new FeedParams.Builder();
         feedParamsBuilder.setDataFormat(configuration.dataFormat());
         feedParamsBuilder.setRoute(configuration.route());
         feedParamsBuilder.setMaxSleepTimeMs(configuration.maxSleepTimeMs());
         feedParamsBuilder.setMaxInFlightRequests(configuration.maxInFlightRequests());
         feedParamsBuilder.setLocalQueueTimeOut(3600*1000); //1 hour queue timeout
+        return feedParamsBuilder;
+    }
+
+    protected SessionParams.Builder configureSessionParams() {
+        SessionParams.Builder sessionParamsBuilder = new SessionParams.Builder();
+        sessionParamsBuilder.setThrottlerMinSize(configuration.throttlerMinSize());
+        sessionParamsBuilder.setClientQueueSize(configuration.maxInFlightRequests()*2);
+        return sessionParamsBuilder;
+    }
+    
+    private void initialize() {
+        if (!configuration.dryrun() && configuration.randomSartupSleepMs() > 0) {
+            int delay = new Random().nextInt(configuration.randomSartupSleepMs());
+            log.info("VespaStorage: Delaying startup by " + delay + " ms");
+            try {
+                Thread.sleep(delay);
+            } catch (Exception e) {}
+        }
+
+        ConnectionParams.Builder connParamsBuilder = configureConnectionParams();
+        FeedParams.Builder feedParamsBuilder = configureFeedParams();
+        SessionParams.Builder sessionParams = configureSessionParams();
 
-        SessionParams.Builder sessionParams = new SessionParams.Builder();
-        sessionParams.setThrottlerMinSize(configuration.throttlerMinSize());
         sessionParams.setConnectionParams(connParamsBuilder.build());
         sessionParams.setFeedParams(feedParamsBuilder.build());
-        sessionParams.setClientQueueSize(configuration.maxInFlightRequests()*2);
-
 
         String endpoints = configuration.endpoint();
         StringTokenizer tokenizer = new StringTokenizer(endpoints, ",");
diff --git a/vespa-hadoop/src/main/java/com/yahoo/vespa/hadoop/mapreduce/util/VespaConfiguration.java b/vespa-hadoop/src/main/java/com/yahoo/vespa/hadoop/mapreduce/util/VespaConfiguration.java
index 340f1cac55a..59da262428f 100644
--- a/vespa-hadoop/src/main/java/com/yahoo/vespa/hadoop/mapreduce/util/VespaConfiguration.java
+++ b/vespa-hadoop/src/main/java/com/yahoo/vespa/hadoop/mapreduce/util/VespaConfiguration.java
@@ -125,7 +125,7 @@ public class VespaConfiguration {
     }
 
 
-    private String getString(String name) {
+    public String getString(String name) {
         if (override != null && override.containsKey(name)) {
             return override.getProperty(name);
         }
@@ -133,7 +133,7 @@ public class VespaConfiguration {
     }
 
 
-    private int getInt(String name, int defaultValue) {
+    public int getInt(String name, int defaultValue) {
         if (override != null && override.containsKey(name)) {
             return Integer.parseInt(override.getProperty(name));
         }
@@ -141,7 +141,7 @@ public class VespaConfiguration {
     }
 
 
-    private boolean getBoolean(String name, boolean defaultValue) {
+    public boolean getBoolean(String name, boolean defaultValue) {
         if (override != null && override.containsKey(name)) {
             return Boolean.parseBoolean(override.getProperty(name));
         }
diff --git a/vespa-hadoop/src/main/java/com/yahoo/vespa/hadoop/pig/VespaStorage.java b/vespa-hadoop/src/main/java/com/yahoo/vespa/hadoop/pig/VespaStorage.java
index a77805b4153..f0283636725 100644
--- a/vespa-hadoop/src/main/java/com/yahoo/vespa/hadoop/pig/VespaStorage.java
+++ b/vespa-hadoop/src/main/java/com/yahoo/vespa/hadoop/pig/VespaStorage.java
@@ -34,13 +34,14 @@ public class VespaStorage extends StoreFunc {
     private String signature = null;
     private RecordWriter recordWriter = null;
     private ResourceSchema resourceSchema = null;
-    private Properties properties = new Properties();
 
     private static final String PROPERTY_CREATE_DOC_OP = "create-document-operation";
     private static final String PROPERTY_ID_TEMPLATE = "docid";
     private static final String PROPERTY_OPERATION = "operation";
     private static final String PROPERTY_RESOURCE_SCHEMA = "resource_schema";
 
+    Properties properties = new Properties();
+
     public VespaStorage() {
         createDocOp = false;
         template = null;