pass pig parameters by tempfile

luigi/contrib/pig.py

@@ -22,7 +22,7 @@
     # pig home directory
     home: /usr/share/pig
 """
-
+from contextlib import contextmanager
 import logging
 import os
 import select
@@ -95,27 +95,32 @@ def pig_script_path(self):
         """
         raise NotImplementedError("subclass should define pig_script_path")
 
+    @contextmanager
     def _build_pig_cmd(self):
         opts = self.pig_options()
 
-        for k, v in six.iteritems(self.pig_parameters()):
-            opts.append("-p")
-            opts.append("%s=%s" % (k, v))
+        line = lambda k, v: ('%s=%s%s' % (k, v, os.linesep)).encode('utf-8')
+        with tempfile.NamedTemporaryFile() as param_file, tempfile.NamedTemporaryFile() as prop_file:
+            if self.pig_parameters():
+                items = six.iteritems(self.pig_parameters())
+                param_file.writelines(line(k, v) for (k, v) in items)
+                opts.append('-param_file')
+                opts.append(param_file.name)
 
-        if self.pig_properties():
-            with open('pig_property_file', 'w') as prop_file:
-                prop_file.writelines(["%s=%s%s" % (k, v, os.linesep) for (k, v) in six.iteritems(self.pig_properties())])
-            opts.append('-propertyFile')
-            opts.append('pig_property_file')
+            if self.pig_properties():
+                items = six.iteritems(self.pig_properties())
+                prop_file.writelines(line(k, v) for (k, v) in items)
+                opts.append('-propertyFile')
+                opts.append(prop_file.name)
 
-        cmd = [self.pig_command_path()] + opts + ["-f", self.pig_script_path()]
+            cmd = [self.pig_command_path()] + opts + ["-f", self.pig_script_path()]
 
-        logger.info(subprocess.list2cmdline(cmd))
-        return cmd
+            logger.info(subprocess.list2cmdline(cmd))
+            yield cmd
 
     def run(self):
-        cmd = self._build_pig_cmd()
-        self.track_and_progress(cmd)
+        with self._build_pig_cmd() as cmd:
+            self.track_and_progress(cmd)
 
     def track_and_progress(self, cmd):
         temp_stdout = tempfile.TemporaryFile()

test/contrib/pig_test.py

@@ -101,47 +101,71 @@ def test_run__success(self, mock):
         arglist_result = []
         p = subprocess.Popen
         subprocess.Popen = _get_fake_Popen(arglist_result, 0)
-        try:
-            job = ComplexTestJob()
-            job.run()
-            self.assertEqual([['/usr/share/pig/bin/pig', '-x', 'local', '-p',
-                               'YOUR_PARAM_NAME=Your param value',
-                               '-propertyFile', 'pig_property_file', '-f',
-                               'my_complex_pig_script.pig']], arglist_result)
-
-            # Check property file
-            with open('pig_property_file') as pprops_file:
-                pprops = pprops_file.readlines()
-                self.assertEqual(1, len(pprops))
-                self.assertEqual('pig.additional.jars=/path/to/your/jar\n', pprops[0])
-        finally:
-            subprocess.Popen = p
+
+        with tempfile.NamedTemporaryFile(delete=False) as param_file_mock, \
+                tempfile.NamedTemporaryFile(delete=False) as prop_file_mock, \
+                patch('luigi.contrib.pig.tempfile.NamedTemporaryFile',
+                      side_effect=[param_file_mock, prop_file_mock]):
+            try:
+                job = ComplexTestJob()
+                job.run()
+                self.assertEqual([['/usr/share/pig/bin/pig', '-x', 'local',
+                                   '-param_file', param_file_mock.name,
+                                   '-propertyFile', prop_file_mock.name,
+                                   '-f', 'my_complex_pig_script.pig']],
+                                 arglist_result)
+
+                # Check param file
+                with open(param_file_mock.name) as pparams_file:
+                    pparams = pparams_file.readlines()
+                    self.assertEqual(1, len(pparams))
+                    self.assertEqual('YOUR_PARAM_NAME=Your param value\n', pparams[0])
+
+                # Check property file
+                with open(prop_file_mock.name) as pprops_file:
+                    pprops = pprops_file.readlines()
+                    self.assertEqual(1, len(pprops))
+                    self.assertEqual('pig.additional.jars=/path/to/your/jar\n', pprops[0])
+            finally:
+                subprocess.Popen = p
 
     @patch('subprocess.Popen')
     def test_run__fail(self, mock):
         arglist_result = []
         p = subprocess.Popen
         subprocess.Popen = _get_fake_Popen(arglist_result, 1)
-        try:
-            job = ComplexTestJob()
-            job.run()
-        except PigJobError as e:
-            p = e
-            self.assertEqual('stderr', p.err)
-            self.assertEqual([['/usr/share/pig/bin/pig', '-x', 'local', '-p',
-                               'YOUR_PARAM_NAME=Your param value',
-                               '-propertyFile', 'pig_property_file', '-f',
-                               'my_complex_pig_script.pig']], arglist_result)
-
-            # Check property file
-            with open('pig_property_file') as pprops_file:
-                pprops = pprops_file.readlines()
-                self.assertEqual(1, len(pprops))
-                self.assertEqual('pig.additional.jars=/path/to/your/jar\n', pprops[0])
-        else:
-            self.fail("Should have thrown PigJobError")
-        finally:
-            subprocess.Popen = p
+
+        with tempfile.NamedTemporaryFile(delete=False) as param_file_mock, \
+                tempfile.NamedTemporaryFile(delete=False) as prop_file_mock, \
+                patch('luigi.contrib.pig.tempfile.NamedTemporaryFile',
+                      side_effect=[param_file_mock, prop_file_mock]):
+            try:
+                job = ComplexTestJob()
+                job.run()
+            except PigJobError as e:
+                p = e
+                self.assertEqual('stderr', p.err)
+                self.assertEqual([['/usr/share/pig/bin/pig', '-x', 'local',
+                                   '-param_file', param_file_mock.name,
+                                   '-propertyFile', prop_file_mock.name, '-f',
+                                   'my_complex_pig_script.pig']],
+                                 arglist_result)
+
+                # Check param file
+                with open(param_file_mock.name) as pparams_file:
+                    pparams = pparams_file.readlines()
+                    self.assertEqual(1, len(pparams))
+                    self.assertEqual('YOUR_PARAM_NAME=Your param value\n', pparams[0])
+
+                # Check property file
+                with open(prop_file_mock.name) as pprops_file:
+                    pprops = pprops_file.readlines()
+                    self.assertEqual(1, len(pprops))
+                    self.assertEqual('pig.additional.jars=/path/to/your/jar\n', pprops[0])
+            else:
+                self.fail("Should have thrown PigJobError")
+            finally:
+                subprocess.Popen = p
 
 
 def _get_fake_Popen(arglist_result, return_code, *args, **kwargs):