diff --git a/CHANGES.md b/CHANGES.md index 44e7568cb8ee..2624185256e6 100644 --- a/CHANGES.md +++ b/CHANGES.md @@ -80,6 +80,7 @@ ## Bugfixes +* (Python) An unknown pipeline option is no longer given the value of a later argument that does not directly follow it, which made `Pipeline.run()` and `get_all_options(retain_unknown_options=True)` fail with `expected one argument` for flags like `--unknown_flag --job_name=foo bar` ([#22434](https://github.com/apache/beam/issues/22434)). * Fixed X (Java/Python) ([#X](https://github.com/apache/beam/issues/X)). ## Security Fixes diff --git a/sdks/python/apache_beam/options/pipeline_options.py b/sdks/python/apache_beam/options/pipeline_options.py index 113dcb00e190..7445551029dd 100644 --- a/sdks/python/apache_beam/options/pipeline_options.py +++ b/sdks/python/apache_beam/options/pipeline_options.py @@ -24,6 +24,7 @@ import json import logging import os +import sys from typing import Any from typing import Callable from typing import Optional @@ -91,6 +92,26 @@ def _f(value): return _f +def _find_positions(args, flags): + """Returns the index in flags of each of args. + + args is expected to be a subsequence of flags, such as the unknown args that + argparse returns from parse_known_args. The position of an arg that cannot be + found is None. + """ + positions = [] + start = 0 + for arg in args: + try: + position = flags.index(arg, start) + except ValueError: + positions.append(None) + continue + positions.append(position) + start = position + 1 + return positions + + class _BeamArgumentParser(argparse.ArgumentParser): """An ArgumentParser that supports ValueProvider options. @@ -535,6 +556,13 @@ def add_new_arg(arg, **kwargs): parser.add_argument(arg, **kwargs) seen.add(arg) + # Unknown args are returned in their original order, but they are not + # necessarily adjacent in the original flags, e.g. a known option may + # sit between an unknown flag and a stray value. Find where each one + # came from so that a value is only given to the flag right before it. + all_flags = sys.argv[1:] if self._flags is None else list(self._flags) + positions = _find_positions(unknown_args, all_flags) + i = 0 while i < len(unknown_args): # End of argument parsing. @@ -545,7 +573,12 @@ def add_new_arg(arg, **kwargs): if not unknown_args[i].startswith('-'): i += 1 continue - if i + 1 >= len(unknown_args) or unknown_args[i + 1].startswith('-'): + has_value = ( + '=' not in unknown_args[i] and i + 1 < len(unknown_args) and + not unknown_args[i + 1].startswith('-') and + positions[i] is not None and positions[i] + 1 < len(all_flags) and + not all_flags[positions[i] + 1].startswith('-')) + if not has_value: split = unknown_args[i].split('=', 1) if len(split) == 1: add_new_arg(unknown_args[i], action='store_true') diff --git a/sdks/python/apache_beam/options/pipeline_options_test.py b/sdks/python/apache_beam/options/pipeline_options_test.py index 79ea0aa650f8..387e81d4426e 100644 --- a/sdks/python/apache_beam/options/pipeline_options_test.py +++ b/sdks/python/apache_beam/options/pipeline_options_test.py @@ -407,6 +407,94 @@ def test_retain_unknown_options_unary_single_dash_store_true(self): result = options.get_all_options(retain_unknown_options=True) self.assertEqual(result['i'], True) + # Unknown args are not necessarily adjacent in the original flags, so a + # value must only be attributed to an unknown flag it directly follows. + # See https://github.com/apache/beam/issues/22434. + @parameterized.expand([ + (['--unknown_option', '--job_name=foo', 'bar'], ), + (['--unknown_option', '--job_name', 'foo', 'bar'], ), + ]) + def test_retain_unknown_options_unary_followed_by_known_option(self, flags): + options = PipelineOptions(flags) + result = options.get_all_options(retain_unknown_options=True) + self.assertEqual(result['unknown_option'], True) + self.assertEqual(result['job_name'], 'foo') + + def test_retain_unknown_options_mixed_with_known_options(self): + options = PipelineOptions([ + '--unknown_unary', + '--job_name=foo', + 'stray_value', + '--unknown_binary', + 'some_value', + '--project', + 'bar', + ]) + result = options.get_all_options(retain_unknown_options=True) + self.assertEqual(result['unknown_unary'], True) + self.assertEqual(result['unknown_binary'], 'some_value') + self.assertEqual(result['job_name'], 'foo') + self.assertEqual(result['project'], 'bar') + + def test_retain_unknown_options_binary_equals_followed_by_value(self): + options = PipelineOptions(['--unknown_option=some_value', 'stray_value']) + result = options.get_all_options(retain_unknown_options=True) + self.assertEqual(result['unknown_option'], 'some_value') + self.assertNotIn('unknown_option=some_value', result) + + def test_retain_unknown_options_repeated_unary_store_true(self): + options = PipelineOptions( + ['--unknown_option', '--job_name=foo', 'bar', '--unknown_option']) + result = options.get_all_options(retain_unknown_options=True) + self.assertEqual(result['unknown_option'], True) + + def test_retain_unknown_options_repeated_binary_keeps_last_value(self): + options = PipelineOptions([ + '--unknown_option', + 'first', + '--job_name=foo', + '--unknown_option', + 'second', + ]) + result = options.get_all_options(retain_unknown_options=True) + self.assertEqual(result['unknown_option'], 'second') + + def test_retain_unknown_options_unary_before_separator(self): + options = PipelineOptions( + ['--unknown_option', '--', 'bar', '--after_separator']) + result = options.get_all_options(retain_unknown_options=True) + self.assertEqual(result['unknown_option'], True) + self.assertNotIn('after_separator', result) + + def test_retain_unknown_options_unary_single_dash_followed_by_known_option( + self): + options = PipelineOptions(['-i', '--job_name=foo', 'bar']) + result = options.get_all_options(retain_unknown_options=True) + self.assertEqual(result['i'], True) + self.assertEqual(result['job_name'], 'foo') + + def test_retain_unknown_options_does_not_take_value_of_known_positional(self): + options = PipelineOptions(['--unknown_option', 'some_value']) + + def add_extra_options(parser): + parser.add_argument('extra_positional', nargs='?') + + result = options.get_all_options( + add_extra_args_fn=add_extra_options, retain_unknown_options=True) + self.assertEqual(result['unknown_option'], True) + self.assertEqual(result['extra_positional'], 'some_value') + + def test_retain_unknown_options_value_then_known_positional(self): + options = PipelineOptions(['--unknown_option', 'some_value', 'other']) + + def add_extra_options(parser): + parser.add_argument('extra_positional', nargs='?') + + result = options.get_all_options( + add_extra_args_fn=add_extra_options, retain_unknown_options=True) + self.assertEqual(result['unknown_option'], 'some_value') + self.assertEqual(result['extra_positional'], 'other') + def test_override_options(self): base_flags = ['--num_workers', '5'] options = PipelineOptions(base_flags)