PIG-5315: pig.script is not set for scripts run via PigServer (satishsaley via rohini)

git-svn-id: https://svn.apache.org/repos/asf/pig/trunk@1817125 13f79535-47bb-0310-9956-ffa450edef68
diff --git a/CHANGES.txt b/CHANGES.txt
index 711bc36..fd01d2d 100644
--- a/CHANGES.txt
+++ b/CHANGES.txt
@@ -60,6 +60,8 @@
  
 BUG FIXES
 
+PIG-5315: pig.script is not set for scripts run via PigServer (satishsaley via rohini)
+
 PIG-5310: MergeJoin throwing NullPointer Exception (satishsaley via rohini)
 
 PIG-5314: Abort method is not implemented in PigProcessor (satishsaley via rohini)
diff --git a/src/org/apache/pig/PigServer.java b/src/org/apache/pig/PigServer.java
index ac8f818..c35fd31 100644
--- a/src/org/apache/pig/PigServer.java
+++ b/src/org/apache/pig/PigServer.java
@@ -43,6 +43,7 @@
 import java.util.Set;
 import java.util.concurrent.atomic.AtomicInteger;
 
+import org.apache.commons.io.IOUtils;
 import org.apache.commons.lang.StringUtils;
 import org.apache.commons.logging.Log;
 import org.apache.commons.logging.LogFactory;
@@ -775,7 +776,10 @@
      */
     public void registerScript(InputStream in, Map<String,String> params,List<String> paramsFiles) throws IOException {
         try {
-            String substituted = pigContext.doParamSubstitution(in, paramMapToList(params), paramsFiles);
+            String script = IOUtils.toString(in);
+            ScriptState.get().setScript(script);
+            String substituted = pigContext.doParamSubstitution(new BufferedReader(new StringReader(script)),
+                    paramMapToList(params), paramsFiles);
             GruntParser grunt = new GruntParser(new StringReader(substituted), this);
             grunt.setInteractive(false);
             grunt.parseStopOnError(true);
@@ -855,6 +859,7 @@
         FileInputStream fis = null;
         try{
             fis = new FileInputStream(fileName);
+            ScriptState.get().setFileName(fileName);
             registerScript(fis, params, paramsFiles);
         }catch (FileNotFoundException e){
             log.error(e.getLocalizedMessage());
diff --git a/src/org/apache/pig/tools/grunt/GruntParser.java b/src/org/apache/pig/tools/grunt/GruntParser.java
index eca6d2d..a5d4495 100644
--- a/src/org/apache/pig/tools/grunt/GruntParser.java
+++ b/src/org/apache/pig/tools/grunt/GruntParser.java
@@ -73,6 +73,7 @@
 import org.apache.pig.tools.pigstats.JobStats;
 import org.apache.pig.tools.pigstats.PigStats;
 import org.apache.pig.tools.pigstats.PigStats.JobGraph;
+import org.apache.pig.tools.pigstats.ScriptState;
 import org.apache.pig.validator.BlackAndWhitelistFilter;
 import org.apache.pig.validator.PigCommandFilter;
 import org.fusesource.jansi.Ansi;
@@ -494,9 +495,12 @@
                 setBatchOn();
                 mPigServer.setJobName(script);
                 try {
-                    loadScript(script, true, false, mLoadOnly, params, files);
+                    FetchFileRet scriptFile = FileLocalizer.fetchFile(mConf, script);
+                    ScriptState.get().beginNestedScript(scriptFile.file);
+                    loadScript(scriptFile, script, true, false, mLoadOnly, params, files);
                     executeBatch();
                 } finally {
+                    ScriptState.get().endNestedScript();
                     discardBatch();
                 }
             } else {
@@ -508,6 +512,11 @@
     }
 
     private void loadScript(String script, boolean batch, boolean loadOnly, boolean illustrate,
+            List<String> params, List<String> files) throws IOException, ParseException {
+        loadScript(FileLocalizer.fetchFile(mConf, script), script, batch, loadOnly, illustrate, params, files);
+    }
+
+    private void loadScript(FetchFileRet fetchFile, String script, boolean batch, boolean loadOnly, boolean illustrate,
                             List<String> params, List<String> files)
         throws IOException, ParseException {
 
@@ -524,7 +533,6 @@
         pc.setParamFiles(files);
 
         try {
-            FetchFileRet fetchFile = FileLocalizer.fetchFile(mConf, script);
             String cmds = runPreprocessor(fetchFile.file.getAbsolutePath(), params, files);
 
             if (mInteractive && !batch) { // Write prompt and echo commands
diff --git a/src/org/apache/pig/tools/pigstats/ScriptState.java b/src/org/apache/pig/tools/pigstats/ScriptState.java
index e755f1c..9b6f8c2 100644
--- a/src/org/apache/pig/tools/pigstats/ScriptState.java
+++ b/src/org/apache/pig/tools/pigstats/ScriptState.java
@@ -26,6 +26,7 @@
 import java.util.HashSet;
 import java.util.List;
 import java.util.Map;
+import java.util.Stack;
 import java.util.UUID;
 import java.util.jar.Attributes;
 import java.util.jar.JarFile;
@@ -182,6 +183,8 @@
 
     protected List<PigProgressNotificationListener> listeners = Lists.newArrayList();
 
+    private Stack<ScriptInfo> scripts = new Stack<>();
+
     protected ScriptState(String id) {
         this.id = id;
         this.serializedScript = "";
@@ -439,6 +442,31 @@
         return featureLongToString(scriptFeatures);
     }
 
+    /**
+     * Stores information about the current script and pushes it onto a stack.
+     *
+     * @param scriptFile
+     * @throws IOException
+     */
+    public void beginNestedScript(File scriptFile) throws IOException {
+        ScriptInfo scriptInfo = new ScriptInfo();
+        scriptInfo.fileName = this.fileName;
+        scriptInfo.serializedScript = this.serializedScript;
+        scriptInfo.truncatedScript = this.truncatedScript;
+        scripts.push(scriptInfo);
+        this.setScript(scriptFile);
+        this.setFileName(scriptFile.getName());
+    }
+
+    public void endNestedScript() {
+        if (!scripts.isEmpty()) {
+            ScriptInfo scriptInfo = scripts.pop();
+            // Change the current script information
+            this.fileName = scriptInfo.fileName;
+            this.serializedScript = scriptInfo.serializedScript;
+            this.truncatedScript = scriptInfo.truncatedScript;
+        }
+    }
     static class LogicalPlanFeatureVisitor extends LogicalRelationalNodesVisitor {
 
         private BitSet feature;
@@ -712,4 +740,10 @@
             }
         }
     }
+
+    private static class ScriptInfo {
+        String serializedScript;
+        String truncatedScript;
+        String fileName;
+    }
 }
diff --git a/test/org/apache/pig/pigunit/pig/PigServer.java b/test/org/apache/pig/pigunit/pig/PigServer.java
index d348711..920ebcc 100644
--- a/test/org/apache/pig/pigunit/pig/PigServer.java
+++ b/test/org/apache/pig/pigunit/pig/PigServer.java
@@ -20,9 +20,11 @@
 import java.util.Map;
 import java.util.Properties;
 
+import org.apache.commons.io.IOUtils;
 import org.apache.pig.ExecType;
 import org.apache.pig.backend.executionengine.ExecException;
 import org.apache.pig.impl.util.Utils;
+import org.apache.pig.tools.pigstats.ScriptState;
 
 /**
  * Slightly modified PigServer that accepts a list of Pig aliases to override.
@@ -49,6 +51,8 @@
   public void registerScript(String fileName, Map<String, String> aliasOverride)
       throws IOException {
     try {
+      ScriptState.get().setFileName(fileName);
+      ScriptState.get().setScript(IOUtils.toString(new FileInputStream(fileName)));
       InputStream compositeStream = Utils.getCompositeStream(new FileInputStream(fileName), pigContext.getProperties());
       GruntParser grunt = new GruntParser(new InputStreamReader(compositeStream), this, aliasOverride);
       grunt.setInteractive(false);
diff --git a/test/org/apache/pig/test/TestPigServerLocal.java b/test/org/apache/pig/test/TestPigServerLocal.java
index 3817b9d..f50eac3 100644
--- a/test/org/apache/pig/test/TestPigServerLocal.java
+++ b/test/org/apache/pig/test/TestPigServerLocal.java
@@ -20,6 +20,7 @@
 import static org.apache.pig.builtin.mock.Storage.resetData;
 import static org.apache.pig.builtin.mock.Storage.tuple;
 import static org.junit.Assert.assertEquals;
+import static org.junit.Assert.assertNotNull;
 import static org.junit.Assert.assertTrue;
 import static org.junit.Assert.fail;
 
@@ -54,6 +55,7 @@
 import org.apache.pig.tools.grunt.Grunt;
 import org.apache.pig.tools.grunt.GruntParser;
 import org.apache.pig.tools.pigstats.PigStats;
+import org.apache.pig.tools.pigstats.ScriptState;
 import org.junit.Before;
 import org.junit.Test;
 
@@ -167,6 +169,7 @@
         params.put("input", "test/org/apache/pig/test/data/passwd");
         String script="a = load '$input' using PigStorage(':');";
         pig.registerScript(new ByteArrayInputStream(script.getBytes("UTF-8")),params);
+        assertEquals("ScriptState contains different script", script, ScriptState.get().getScript());
         Iterator<Tuple> iter=pig.openIterator("a");
         int index=0;
         List<Tuple> expectedTuples=Util.readFile2TupleList("test/org/apache/pig/test/data/passwd", ":");