IMPALA-7673: Support values from other variables in Impala shell --var

Prior to this patch, Impala shell --var could not accept values from
other variables unlike the one in Impala interactive shell with the SET
command.  This patch refactors the logic of variable substitution to
use the same logic in both interactive and command line shells.

Example:
$ impala-shell.sh \
    --var="msg1=1" \
    --var="msg2=\${var:msg1}2" \
    --var="msg3=\${var:msg1}\${var:msg2}"

[localhost:21000] default> select ${var:msg3};
Query: select 112
+-----+
| 112 |
+-----+
| 112 |
+-----+

Testing:
- Added a new shell test
- Ran all shell tests

Change-Id: Ib5b9fda329c45f2e5682f3cbc76d29ceca2e226a
Reviewed-on: http://gerrit.cloudera.org:8080/11623
Reviewed-by: Impala Public Jenkins <impala-public-jenkins@cloudera.com>
Tested-by: Impala Public Jenkins <impala-public-jenkins@cloudera.com>

shell/impala_shell.py

@@ -543,32 +543,6 @@ class ImpalaShell(object, cmd.Cmd):
         print_to_stderr("Failed to reconnect and close (try %i/%i): %s" % (
             cancel_try + 1, ImpalaShell.CANCELLATION_TRIES, err_msg))
 
-  def _replace_variables(self, query):
-    """Replaces variable within the query text with their corresponding values"""
-    errors = False
-    matches = set(map(lambda v: v.upper(), re.findall(r'(?<!\\)\${([^}]+)}', query)))
-    for name in matches:
-      value = None
-      # Check if syntax is correct
-      var_name = self._get_var_name(name)
-      if var_name is None:
-        print_to_stderr('Error: Unknown substitution syntax (%s). ' % (name,) + \
-                        'Use ${VAR:var_name}.')
-        errors = True
-      else:
-        # Replaces variable value
-        if self.set_variables and var_name in self.set_variables:
-          value = self.set_variables[var_name]
-          regexp = re.compile(r'(?<!\\)\${%s}' % (name,), re.IGNORECASE)
-          query = regexp.sub(value, query)
-        else:
-          print_to_stderr('Error: Unknown variable %s' % (var_name))
-          errors = True
-    if errors:
-      return None
-    else:
-      return query
-
   def set_prompt(self, db):
     self.prompt = ImpalaShell.PROMPT_FORMAT.format(
         host=self.impalad[0], port=self.impalad[1], db=db)
@@ -598,7 +572,7 @@ class ImpalaShell(object, cmd.Cmd):
        the interactive case, when cmdloop is called.
     """
     # Replace variables in the statement before it's executed
-    line = self._replace_variables(line)
+    line = replace_variables(self.set_variables, line)
     # Cmd is an old-style class, hence we need to call the method directly
     # instead of using super()
     # TODO: This may have to be changed to a super() call once we move to Python 3
@@ -673,18 +647,6 @@ class ImpalaShell(object, cmd.Cmd):
     except KeyError:
       return False
 
-  def _get_var_name(self, name):
-    """Look for a namespace:var_name pattern in an option name.
-       Return the variable name if it's a match or None otherwise.
-    """
-    ns_match = re.match(r'^([^:]*):(.*)', name)
-    if ns_match is not None:
-      ns = ns_match.group(1)
-      var_name = ns_match.group(2)
-      if ns in ImpalaShell.VAR_PREFIXES:
-        return var_name
-    return None
-
   def _print_with_set(self, print_level):
     self._print_options(print_level)
     print "\nVariables:"
@@ -721,7 +683,7 @@ class ImpalaShell(object, cmd.Cmd):
         return CmdStatus.ERROR
     option_upper = tokens[0].upper()
     # Check if it's a variable
-    var_name = self._get_var_name(option_upper)
+    var_name = get_var_name(option_upper)
     if var_name is not None:
       # Set the variable
       self.set_variables[var_name] = tokens[1]
@@ -745,7 +707,7 @@ class ImpalaShell(object, cmd.Cmd):
       return CmdStatus.ERROR
     option = args.upper()
     # Check if it's a variable
-    var_name = self._get_var_name(option)
+    var_name = get_var_name(option)
     if var_name is not None:
       if self.set_variables.get(var_name):
         print 'Unsetting variable %s' % var_name
@@ -1539,9 +1501,53 @@ def parse_variables(keyvals):
         parser.print_help()
         sys.exit(1)
       else:
-        vars[match.groups()[0].upper()] = match.groups()[1]
+        vars[match.groups()[0].upper()] = replace_variables(vars, match.groups()[1])
   return vars
 
+
+def replace_variables(set_variables, string):
+  """Replaces variable within the string with their corresponding values using the
+  given set_variables."""
+  errors = False
+  matches = set(map(lambda v: v.upper(), re.findall(r'(?<!\\)\${([^}]+)}', string)))
+  for name in matches:
+    value = None
+    # Check if syntax is correct
+    var_name = get_var_name(name)
+    if var_name is None:
+      print_to_stderr('Error: Unknown substitution syntax (%s). ' % (name,) +
+                      'Use ${VAR:var_name}.')
+      errors = True
+    else:
+      # Replaces variable value
+      if set_variables and var_name in set_variables:
+        value = set_variables[var_name]
+        if value is None:
+          errors = True
+        else:
+          regexp = re.compile(r'(?<!\\)\${%s}' % (name,), re.IGNORECASE)
+          string = regexp.sub(value, string)
+      else:
+        print_to_stderr('Error: Unknown variable %s' % (var_name))
+        errors = True
+  if errors:
+    return None
+  else:
+    return string
+
+
+def get_var_name(name):
+  """Look for a namespace:var_name pattern in an option name.
+     Return the variable name if it's a match or None otherwise.
+  """
+  ns_match = re.match(r'^([^:]*):(.*)', name)
+  if ns_match is not None:
+    ns = ns_match.group(1)
+    var_name = ns_match.group(2)
+    if ns in ImpalaShell.VAR_PREFIXES:
+      return var_name
+  return None
+
 def execute_queries_non_interactive_mode(options, query_options):
   """Run queries in non-interactive mode."""
   if options.query_file:

tests/shell/test_shell_commandline.py

@@ -535,6 +535,31 @@ class TestImpalaShell(ImpalaTestSuite):
     assert ("Error: Could not parse key-value \"foo\". It must follow the pattern "
              "\"KEY=VALUE\".") in result.stderr
 
+    # IMPALA-7673: Test that variable substitution in command line can accept values
+    # from other variables just like the one in interactive shell.
+    result = run_impala_shell_cmd('--var="msg1=1" --var="msg2=${var:msg1}2" '
+                                  '--var="msg3=${var:msg1}${var:msg2}" '
+                                  '--query="select ${var:msg3}"')
+    self._validate_shell_messages(result.stderr, ['112', 'Fetched 1 row(s)'],
+                                  should_exist=True)
+
+    # Test with an escaped variable.
+    result = run_impala_shell_cmd('--var="msg1=1" --var="msg2=${var:msg1}2" '
+                                  '--var="msg3=\${var:msg1}${var:msg2}" '
+                                  '--query="select \'${var:msg3}\'"')
+    self._validate_shell_messages(result.stderr, ['${var:msg1}12', 'Fetched 1 row(s)'],
+                                  should_exist=True)
+
+    # Referencing a non-existent variable will result in an error.
+    result = run_impala_shell_cmd('--var="msg1=1" --var="msg2=${var:doesnotexist}2" '
+                                  '--var="msg3=\${var:msg1}${var:msg2}" '
+                                  '--query="select \'${var:msg3}\'"',
+                                  expect_success=False)
+    self._validate_shell_messages(result.stderr,
+                                  ['Error: Unknown variable DOESNOTEXIST',
+                                   'Could not execute command: select \'${var:msg3}\''],
+                                  should_exist=True)
+
   # Checks if 'messages' exists/does not exist in 'result_stderr' based on the value of
   # 'should_exist'
   def _validate_shell_messages(self, result_stderr, messages, should_exist=True):