From 0ea06378b14b2ae31206499aa2ed7154f8ca59c2 Mon Sep 17 00:00:00 2001
From: Kevin Risden <risdenk@avalonconsult.com>
Date: Wed, 13 Apr 2016 16:45:50 -0500
Subject: [PATCH] [scripts] Fix bin/ycsb for Python 2.6 and bin/ycsb
 check_output

---
 bin/ycsb | 45 ++++++++++++++++++++++++++++++++++++++-------
 1 file changed, 38 insertions(+), 7 deletions(-)

diff --git a/bin/ycsb b/bin/ycsb
index ef2c18598f..1595801b20 100755
--- a/bin/ycsb
+++ b/bin/ycsb
@@ -23,11 +23,13 @@ import os
 import shlex
 import sys
 import subprocess
+
 try:
+    mod = __import__('argparse')
     import argparse
 except ImportError:
     print >> sys.stderr, '[ERROR] argparse not found. Try installing it via "pip".'
-    raise
+    exit(1)
 
 BASE_URL = "https://github.com/brianfrankcooper/YCSB/tree/master/"
 COMMANDS = {
@@ -120,12 +122,41 @@ def usage():
 
     return output.getvalue()
 
-def check_output(cmd):
-    p = subprocess.Popen(cmd, stdout=subprocess.PIPE)
-    stdout, _ = p.communicate()
-    if p.returncode:
-        raise subprocess.CalledProcessError(p.returncode, cmd)
-    return stdout
+# Python 2.6 doesn't have check_output. Add the method as it is in Python 2.7
+# Based on https://github.com/python/cpython/blob/2.7/Lib/subprocess.py#L545
+def check_output(*popenargs, **kwargs):
+    r"""Run command with arguments and return its output as a byte string.
+
+    If the exit code was non-zero it raises a CalledProcessError.  The
+    CalledProcessError object will have the return code in the returncode
+    attribute and output in the output attribute.
+
+    The arguments are the same as for the Popen constructor.  Example:
+
+    >>> check_output(["ls", "-l", "/dev/null"])
+    'crw-rw-rw- 1 root root 1, 3 Oct 18  2007 /dev/null\n'
+
+    The stdout argument is not allowed as it is used internally.
+    To capture standard error in the result, use stderr=STDOUT.
+
+    >>> check_output(["/bin/sh", "-c",
+    ...               "ls -l non_existent_file ; exit 0"],
+    ...              stderr=STDOUT)
+    'ls: non_existent_file: No such file or directory\n'
+    """
+    if 'stdout' in kwargs:
+        raise ValueError('stdout argument not allowed, it will be overridden.')
+    process = subprocess.Popen(stdout=subprocess.PIPE, *popenargs, **kwargs)
+    output, unused_err = process.communicate()
+    retcode = process.poll()
+    if retcode:
+        cmd = kwargs.get("args")
+        if cmd is None:
+            cmd = popenargs[0]
+        error = subprocess.CalledProcessError(retcode, cmd)
+        error.output = output
+        raise error
+    return output
 
 def debug(message):
     print >> sys.stderr, "[DEBUG] ", message