diff --git a/CHANGES.md b/CHANGES.md index c16264b8ca4a..32662c309522 100644 --- a/CHANGES.md +++ b/CHANGES.md @@ -94,6 +94,7 @@ ### New Feature * X feature added (Java/Python) ([#X](https://github.com/apache/beam/issues/X)). +* (Python) Added an experimental `beam` command. `beam run FILE.yaml` runs a Beam YAML pipeline like `python -m apache_beam.yaml.main --yaml_pipeline_file=FILE.yaml`, with documented exit codes, and Ctrl-C cancels the pipeline ([#40891](https://github.com/apache/beam/issues/40891)). * (Python) Added support for Protobuf 7.x ([#40580](https://github.com/apache/beam/issues/40580)). * (Go) Added the standard `millis_instant` and `decimal` logical types for cross-language rows ([#39684](https://github.com/apache/beam/issues/39684)). * [Beam SQL] Added `PRIMARY KEY (...) NOT ENFORCED` to `CREATE EXTERNAL TABLE`, and `ALTER TABLE ... ADD/DROP PRIMARY KEY` ([#40826](https://github.com/apache/beam/issues/40826)). diff --git a/sdks/python/apache_beam/cli/__init__.py b/sdks/python/apache_beam/cli/__init__.py new file mode 100644 index 000000000000..82de2c27aff5 --- /dev/null +++ b/sdks/python/apache_beam/cli/__init__.py @@ -0,0 +1,21 @@ +# +# Licensed to the Apache Software Foundation (ASF) under one or more +# contributor license agreements. See the NOTICE file distributed with +# this work for additional information regarding copyright ownership. +# The ASF licenses this file to You under the Apache License, Version 2.0 +# (the "License"); you may not use this file except in compliance with +# the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# + +"""The `beam` command-line tool. Experimental: commands and flags may change. + +The `beam` console script runs `apache_beam.cli.main.main`. +""" diff --git a/sdks/python/apache_beam/cli/main.py b/sdks/python/apache_beam/cli/main.py new file mode 100644 index 000000000000..e7b00f859abf --- /dev/null +++ b/sdks/python/apache_beam/cli/main.py @@ -0,0 +1,277 @@ +# +# Licensed to the Apache Software Foundation (ASF) under one or more +# contributor license agreements. See the NOTICE file distributed with +# this work for additional information regarding copyright ownership. +# The ASF licenses this file to You under the Apache License, Version 2.0 +# (the "License"); you may not use this file except in compliance with +# the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# + +"""Entry point of the experimental `beam` command. + +`beam run FILE` runs a Beam YAML pipeline file the same way as +`python -m apache_beam.yaml.main --yaml_pipeline_file=FILE`. It builds the +pipeline with that module's own functions, so Beam YAML features (Jinja +variables, providers, ...) and runners behave identically, and every flag the +command doesn't define is passed through unchanged as a Beam pipeline option. + +On top of that, it gives command-line users and scripts what they rely on: +errors without tracebacks unless `--verbose` is given, documented exit codes, +Ctrl-C cancelling the pipeline, and a last line naming the runner that ran it. +""" + +import argparse +import logging +import os +import sys +import traceback + +import apache_beam as beam +from apache_beam.options.pipeline_options import StandardOptions +from apache_beam.runners.runner import PipelineState + +# Exit codes, also listed in `beam run --help`. Usage errors use 2, the code +# argparse itself exits with. +EXIT_OK = 0 +EXIT_PIPELINE_FAILED = 1 +EXIT_USAGE = 2 +EXIT_MISSING_PACKAGE = 4 +EXIT_INTERRUPTED = 130 # 128 + SIGINT, the shell convention for Ctrl-C. + +_RUN_EPILOG = """\ +examples: + beam run pipeline.yaml + beam run pipeline.yaml --runner=dataflow --project=my-project \\ + --region=us-central1 --temp_location=gs://my-bucket/tmp + +Every other --flag is passed to the pipeline as a Beam pipeline option. +Write them after FILE, as --flag=value. + +exit codes: + 0 the pipeline finished (or was submitted with --no_wait_until_finish) + 1 the pipeline failed + 2 usage error, or the pipeline could not be built + 4 a required package is missing + 130 interrupted (Ctrl-C) +""" + + +def main(argv=None): + """Runs the `beam` command and returns its exit code.""" + args, pipeline_args = _build_parser().parse_known_args(argv) + # The same logging setup as `python -m apache_beam.yaml.main`. + logging.getLogger().setLevel(logging.INFO) + try: + return _run(args.file, pipeline_args, args.verbose) + except _CliError as e: + print(f'error: {e}', file=sys.stderr) + return e.exit_code + except KeyboardInterrupt: + print('error: interrupted', file=sys.stderr) + return EXIT_INTERRUPTED + + +def _build_parser(): + parser = argparse.ArgumentParser( + prog='beam', + description='Apache Beam command-line tool (experimental).', + allow_abbrev=False) + commands = parser.add_subparsers( + dest='command', required=True, metavar='COMMAND') + # allow_abbrev=False on each parser, so that a pipeline option is never + # taken for an abbreviation of one of this command's flags. + run = commands.add_parser( + 'run', + help='run a Beam YAML pipeline file', + description=( + 'Run a Beam YAML pipeline file, like ' + 'python -m apache_beam.yaml.main --yaml_pipeline_file=FILE.'), + epilog=_RUN_EPILOG, + formatter_class=argparse.RawDescriptionHelpFormatter, + allow_abbrev=False) + run.add_argument( + 'file', + metavar='FILE', + help='the pipeline file: a local path, or a path on any Beam ' + 'filesystem such as gs://') + run.add_argument( + '--verbose', action='store_true', help='show tracebacks for errors') + return parser + + +class _CliError(Exception): + """An error reported to the user as one message, without a traceback.""" + def __init__(self, message, exit_code): + super().__init__(message) + self.exit_code = exit_code + + +def _run(file, pipeline_args, verbose): + # Paths with a scheme (gs://, s3://, ...) are opened by Beam's filesystems + # when the pipeline is built, which reports its own errors for them. + if '://' not in file and not os.path.isfile(file): + hint = '' + if pipeline_args: + # `beam run --runner dataflow p.yaml` makes argparse take "dataflow" + # for FILE. + hint = '\n pipeline options go after FILE, written as --flag=value' + raise _CliError(f'no such file: {file!r}{hint}', EXIT_USAGE) + + # Beam YAML is imported here, not at the top of the module, because it needs + # the [yaml] extra; a missing extra is then reported with exit code 4 and a + # hint instead of a traceback. + try: + from apache_beam.yaml import main as yaml_main + except ImportError as e: + _maybe_print_traceback(verbose) + raise _CliError( + f'a required package is missing: {e}\n' + " Beam YAML needs the [yaml] extra: pip install 'apache-beam[yaml]'", + EXIT_MISSING_PACKAGE) from None + + # Build and run in two steps (yaml_main.run does both in a `with` block) so + # that build errors and pipeline failures get different exit codes. Enter + # `Pipeline.__enter__` across both steps so Java expansion service + # subprocesses stay cached between `constructor(pipeline)` and + # `pipeline.run()` (when artifacts are staged), and `--beam_services` + # overrides apply to both. + pipeline = None + try: + try: + options, constructor, display_data = ( + yaml_main.build_pipeline_components_from_argv( + ['--yaml_pipeline_file=' + file] + pipeline_args)) + pipeline = beam.Pipeline(options=options, display_data=display_data) + pipeline.__enter__() + constructor(pipeline) + except ImportError as e: + # A transform's extra is missing, for example [gcp] for BigQuery. + _maybe_print_traceback(verbose) + raise _CliError( + f'a required package is missing: {e}\n' + ' install the extra the pipeline needs, for example: ' + "pip install 'apache-beam[yaml,gcp]'", + EXIT_MISSING_PACKAGE) from None + except Exception as e: # pylint: disable=broad-except + _maybe_print_traceback(verbose) + raise _CliError( + f'could not build the pipeline: {e}', EXIT_USAGE) from None + + no_wait = options.view_as(StandardOptions).no_wait_until_finish + state = None + try: + result = pipeline.run() + if not no_wait: + try: + state = result.wait_until_finish() + except KeyboardInterrupt: + # main() reports the interruption after the pipeline is cancelled. + _cancel(result) + raise + except Exception as e: # pylint: disable=broad-except + _maybe_print_traceback(verbose) + more = '' if verbose else '\n rerun with --verbose for the traceback' + raise _CliError( + f'pipeline failed: {_summarize(e)}{more}', + EXIT_PIPELINE_FAILED) from None + finally: + if pipeline is not None and hasattr(pipeline, '_extra_context'): + pipeline._extra_context.close() + + if no_wait: + print( + f'Submitted {_describe_runner(options, result)} ' + '(not waiting for completion).') + return EXIT_OK + + if state != PipelineState.DONE: + raise _CliError(f'pipeline finished in state {state}', EXIT_PIPELINE_FAILED) + + print(f'Done. Ran {_describe_runner(options, result)}.') + return EXIT_OK + + +def _maybe_print_traceback(verbose): + if verbose: + traceback.print_exc() + + +_PYTHON_TRACEBACK = 'Traceback (most recent call last):' + + +def _summarize(error, max_lines=10): + """Returns the part of a pipeline error that says what went wrong. + + Runner errors often embed whole stack traces. Prism and the default runner + put a line of context ("Pipeline job-001 failed in state FAILED: bundle ... + failed:") before the worker's Python traceback, whose last line names the + exception and the failing step; Java runners add `\\tat ...` frames. Stack + frames are indented, so only unindented lines are kept, and for a Python + traceback only the context and the last exception line. `--verbose` still + prints everything. + """ + message = str(error) + if type(error) is not RuntimeError and not message.startswith( + type(error).__name__): + # Errors raised directly (for example a ZeroDivisionError from the + # in-process runner) don't name their type in the message. + message = f'{type(error).__name__}: {message}' + lines = [ + line.rstrip() for line in message.splitlines() + if line.strip() and not line[0].isspace() + ] + if any(_PYTHON_TRACEBACK in line for line in lines): + context = lines[0].split(_PYTHON_TRACEBACK)[0].rstrip() + causes = [ + line for line in lines if _PYTHON_TRACEBACK not in line and + not line.startswith('During handling of the above exception') + ] + lines = [line for line in (context, causes[-1] if causes else '') if line] + if not lines: + return repr(error) + shown = '\n '.join(lines[:max_lines]) + return shown + ('\n ...' if len(lines) > max_lines else '') + + +def _cancel(result): + """Cancels the pipeline after Ctrl-C, saying whether that worked. + + A local runner usually receives the same Ctrl-C and stops by itself, in + which case cancelling fails harmlessly. + """ + try: + result.cancel() + print('Cancelled the pipeline.', file=sys.stderr) + except Exception as e: # pylint: disable=broad-except + print( + f'Could not cancel the pipeline ({e}). If it is still running, ' + "stop it with the runner's tools.", + file=sys.stderr) + + +def _describe_runner(options, result): + """Says which runner ran the pipeline, for the last line of output. + + Without --runner (or with --runner=DirectRunner), the Python SDK runs the + pipeline on Prism and falls back to the Python direct runner when Prism + can't run it. The type of the result tells the two apart. + """ + runner = options.view_as(StandardOptions).runner + name = str(runner or 'DirectRunner').rsplit('.', 1)[-1].lower() + if name not in ('direct', 'directrunner'): + return f'on {runner}' + if 'portable_runner' in type(result).__module__: + return 'locally on Prism' + return 'locally on the Python direct runner' + + +if __name__ == '__main__': + sys.exit(main()) diff --git a/sdks/python/apache_beam/cli/main_test.py b/sdks/python/apache_beam/cli/main_test.py new file mode 100644 index 000000000000..96af4ae4dea4 --- /dev/null +++ b/sdks/python/apache_beam/cli/main_test.py @@ -0,0 +1,390 @@ +# +# Licensed to the Apache Software Foundation (ASF) under one or more +# contributor license agreements. See the NOTICE file distributed with +# this work for additional information regarding copyright ownership. +# The ASF licenses this file to You under the Apache License, Version 2.0 +# (the "License"); you may not use this file except in compliance with +# the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# + +import contextlib +import glob +import io +import json +import os +import shutil +import tempfile +import unittest +from unittest import mock + +import apache_beam as beam +from apache_beam.cli import main as cli +from apache_beam.options.pipeline_options import PipelineOptions +from apache_beam.utils import subprocess_server + +try: + import jsonschema + from apache_beam.yaml import main as yaml_main +except ImportError: + jsonschema = None + yaml_main = None + +# Paths are formatted with `json.dumps` so Windows backslashes (such as `\U` +# and `\t` in `C:\Users\...\temp`) are escaped as valid YAML strings. +PIPELINE = ''' +pipeline: + type: chain + transforms: + - type: ReadFromCsv + config: {path: %(input)s} + - type: Filter + config: {language: python, keep: "age > 20"} + - type: WriteToJson + config: {path: %(output)s} +''' + +FAILING_PIPELINE = ''' +pipeline: + type: chain + transforms: + - type: ReadFromCsv + config: {path: %(input)s} + - type: MapToFields + config: {language: python, fields: {x: "age // 0"}} + - type: WriteToJson + config: {path: %(output)s} +''' + +INVALID_PIPELINE = ''' +pipeline: + type: chain + transforms: + - type: NoSuchTransform +''' + + +def _run_cli(argv): + """Runs `beam` with argv; returns (exit code, stdout, stderr).""" + out, err = io.StringIO(), io.StringIO() + with contextlib.redirect_stdout(out), contextlib.redirect_stderr(err): + code = cli.main(argv) + return code, out.getvalue(), err.getvalue() + + +def _read_lines(prefix): + lines = [] + for path in glob.glob(prefix + '*'): + with open(path) as f: + lines.extend(f.read().splitlines()) + return sorted(lines) + + +class ParserTest(unittest.TestCase): + def test_usage_errors_exit_with_2(self): + for argv in ([], ['run'], ['nosuchcommand', 'x']): + with self.subTest(argv=argv): + with contextlib.redirect_stderr(io.StringIO()): + with self.assertRaises(SystemExit) as e: + cli.main(argv) + self.assertEqual(e.exception.code, cli.EXIT_USAGE) + + def test_missing_file(self): + code, _, err = _run_cli(['run', 'no/such/pipeline.yaml']) + self.assertEqual(code, cli.EXIT_USAGE) + self.assertIn("no such file: 'no/such/pipeline.yaml'", err) + self.assertNotIn('pipeline options go after FILE', err) + + def test_missing_file_hints_at_option_order(self): + # argparse takes "dataflow" for FILE when an option uses a space. + code, _, err = _run_cli(['run', '--runner', 'dataflow', 'p.yaml']) + self.assertEqual(code, cli.EXIT_USAGE) + self.assertIn("no such file: 'dataflow'", err) + self.assertIn('pipeline options go after FILE', err) + + +@unittest.skipIf( + yaml_main is None or jsonschema is None, 'requires the [yaml] extra') +class RunTest(unittest.TestCase): + def setUp(self): + self.tmp = tempfile.mkdtemp() + self.addCleanup(shutil.rmtree, self.tmp, ignore_errors=True) + self.input = os.path.join(self.tmp, 'users.csv') + with open(self.input, 'w') as f: + f.write('name,age\nAda,36\nLinus,17\nGrace,45\n') + + def _pipeline_file(self, template, output='users.json', input_path=None): + path = os.path.join(self.tmp, 'pipeline.yaml') + with open(path, 'w') as f: + f.write( + template % { + 'input': json.dumps(input_path or self.input), + 'output': json.dumps(os.path.join(self.tmp, output)), + }) + return path + + def test_runs_like_yaml_main(self): + # The CLI and `python -m apache_beam.yaml.main` write the same output for + # the same pipeline file and options. + path = self._pipeline_file(PIPELINE, 'cli_users.json') + code, out, _ = _run_cli(['run', path, '--runner=DirectRunner']) + self.assertEqual(code, cli.EXIT_OK) + self.assertRegex(out, r'Done\. Ran locally on (Prism|the Python direct)') + + reference = self._pipeline_file(PIPELINE, 'reference_users.json') + yaml_main.run( + ['--yaml_pipeline_file=' + reference, '--runner=DirectRunner']) + cli_rows = _read_lines(os.path.join(self.tmp, 'cli_users.json')) + self.assertEqual(len(cli_rows), 2) + self.assertEqual( + cli_rows, _read_lines(os.path.join(self.tmp, 'reference_users.json'))) + + def test_backslash_sequences_in_paths(self): + # Regression test for Windows temp paths (`C:\Users\...\tmp...`), where + # unescaped `\U` and `\t` inside double-quoted YAML strings fail parsing. + tricky_input = os.path.join(self.tmp, r'C_\Users\tmp_users.csv') + shutil.copyfile(self.input, tricky_input) + path = self._pipeline_file( + PIPELINE, 'backslash_users.json', input_path=tricky_input) + code, _, _ = _run_cli(['run', path, '--runner=DirectRunner']) + self.assertEqual(code, cli.EXIT_OK) + self.assertEqual( + len(_read_lines(os.path.join(self.tmp, 'backslash_users.json'))), 2) + + def test_pipeline_failure_exits_with_1(self): + path = self._pipeline_file(FAILING_PIPELINE) + code, _, err = _run_cli(['run', path, '--runner=DirectRunner']) + self.assertEqual(code, cli.EXIT_PIPELINE_FAILED) + # Only the command's own message is checked: with an in-process worker, + # Beam's logging also prints the worker's traceback, as it does for + # `python -m apache_beam.yaml.main`. + start = err.index('error: pipeline failed:') + end = err.index('rerun with --verbose', start) + message = err[start:end] + self.assertIn('ZeroDivisionError', message) + self.assertNotIn('Traceback', message) + + def test_invalid_pipeline_exits_with_2(self): + path = self._pipeline_file(INVALID_PIPELINE) + code, _, err = _run_cli(['run', path]) + self.assertEqual(code, cli.EXIT_USAGE) + self.assertIn('could not build the pipeline:', err) + self.assertIn('NoSuchTransform', err) + self.assertNotIn('Traceback', err) + + def test_verbose_shows_traceback(self): + path = self._pipeline_file(INVALID_PIPELINE) + code, _, err = _run_cli(['run', path, '--verbose']) + self.assertEqual(code, cli.EXIT_USAGE) + self.assertIn('Traceback', err) + + def test_other_flags_are_passed_through_in_order(self): + path = self._pipeline_file(PIPELINE) + with mock.patch.object(yaml_main, + 'build_pipeline_components_from_argv', + side_effect=ValueError('stop')) as build: + _run_cli([ + 'run', + path, + '--runner=DirectRunner', + '--jinja_variables={"a": 1}', + '--verbose', + '--project', + 'p' + ]) + self.assertEqual( + build.call_args.args[0], + [ + '--yaml_pipeline_file=' + path, + '--runner=DirectRunner', + '--jinja_variables={"a": 1}', + '--project', + 'p' + ]) + + def test_missing_package_exits_with_4(self): + path = self._pipeline_file(PIPELINE) + with mock.patch.object(yaml_main, + 'build_pipeline_components_from_argv', + side_effect=ImportError('No module named x')): + code, _, err = _run_cli(['run', path]) + self.assertEqual(code, cli.EXIT_MISSING_PACKAGE) + self.assertIn('No module named x', err) + self.assertIn('pip install', err) + + def test_pipeline_context_active_across_build_and_run(self): + # Java expansion services (Sql, Kafka, JDBC, Iceberg, ...) are started + # during `constructor(pipeline)` and queried again for artifacts during + # `pipeline.run()`, so `Pipeline.__enter__`'s subprocess cache and + # `--beam_services` overrides must stay active across both steps and be + # cleaned up afterwards. + path = self._pipeline_file(PIPELINE) + real_build = yaml_main.build_pipeline_components_from_argv + snapshots = {} + + def capture(stage): + snapshots[stage] = ( + dict(subprocess_server.JavaJarServer._BEAM_SERVICES.replacements), + set(subprocess_server.SubprocessServer._cache._live_owners), + ) + + def wrapped_build(argv): + options, constructor, display_data = real_build(argv) + + def wrapped_constructor(pipeline): + capture('build') + constructor(pipeline) + + return options, wrapped_constructor, display_data + + result = mock.Mock() + result.wait_until_finish.return_value = cli.PipelineState.DONE + + def fake_run(_self): + capture('run') + return result + + owners_before = set(subprocess_server.SubprocessServer._cache._live_owners) + services_before = dict( + subprocess_server.JavaJarServer._BEAM_SERVICES.replacements) + + with ( + mock.patch.object(yaml_main, + 'build_pipeline_components_from_argv', + side_effect=wrapped_build), + mock.patch.object(beam.Pipeline, 'run', fake_run), + ): + code, _, _ = _run_cli( + ['run', path, '--beam_services={"target": "override"}']) + + self.assertEqual(code, cli.EXIT_OK) + build_services, build_owners = snapshots['build'] + run_services, run_owners = snapshots['run'] + self.assertTrue(build_owners - owners_before) + self.assertEqual(build_owners, run_owners) + self.assertEqual(build_services.get('target'), 'override') + self.assertEqual(run_services.get('target'), 'override') + self.assertEqual( + set(subprocess_server.SubprocessServer._cache._live_owners), + owners_before) + self.assertEqual( + dict(subprocess_server.JavaJarServer._BEAM_SERVICES.replacements), + services_before) + + # Also verify cleanup when `constructor(pipeline)` raises during build. + def failing_build(argv): + options, _, display_data = real_build(argv) + + def failing_constructor(_pipeline): + self.assertGreater( + len(subprocess_server.SubprocessServer._cache._live_owners), + len(owners_before)) + raise RuntimeError('expansion failed') + + return options, failing_constructor, display_data + + with mock.patch.object(yaml_main, + 'build_pipeline_components_from_argv', + side_effect=failing_build): + code, _, _ = _run_cli( + ['run', path, '--beam_services={"target": "override"}']) + self.assertEqual(code, cli.EXIT_USAGE) + self.assertEqual( + set(subprocess_server.SubprocessServer._cache._live_owners), + owners_before) + self.assertEqual( + dict(subprocess_server.JavaJarServer._BEAM_SERVICES.replacements), + services_before) + + def test_no_wait_until_finish(self): + path = self._pipeline_file(PIPELINE) + result = mock.Mock() + result.state = cli.PipelineState.RUNNING + with mock.patch.object(beam.Pipeline, 'run', + return_value=result) as run_mock: + code, out, _ = _run_cli(['run', path, '--no_wait_until_finish']) + self.assertEqual(code, cli.EXIT_OK) + run_mock.assert_called_once() + result.wait_until_finish.assert_not_called() + self.assertIn('not waiting for completion', out) + + def test_ctrl_c_cancels_and_exits_with_130(self): + path = self._pipeline_file(PIPELINE) + result = mock.Mock() + result.wait_until_finish.side_effect = KeyboardInterrupt + with mock.patch.object(beam.Pipeline, 'run', return_value=result): + code, _, err = _run_cli(['run', path]) + self.assertEqual(code, cli.EXIT_INTERRUPTED) + result.cancel.assert_called_once() + self.assertIn('Cancelled the pipeline.', err) + + +class HelpersTest(unittest.TestCase): + def test_cancel_failure_is_reported(self): + result = mock.Mock() + result.cancel.side_effect = RuntimeError('job not found') + err = io.StringIO() + with contextlib.redirect_stderr(err): + cli._cancel(result) + self.assertIn( + 'Could not cancel the pipeline (job not found)', err.getvalue()) + + def test_describe_runner(self): + prism_result = type('PipelineResult', (), {})() + type(prism_result).__module__ = ( + 'apache_beam.runners.portability.portable_runner') + direct_result = object() + cases = [ + ([], prism_result, 'locally on Prism'), + ([], direct_result, 'locally on the Python direct runner'), + (['--runner=DirectRunner'], prism_result, 'locally on Prism'), + (['--runner=dataflow'], direct_result, 'on dataflow'), + ] + for argv, result, expected in cases: + with self.subTest(argv=argv, expected=expected): + self.assertEqual( + cli._describe_runner(PipelineOptions(argv), result), expected) + + def test_summarize_keeps_the_first_lines(self): + error = RuntimeError( + 'Job failed.\n\n' + '\n'.join(f'line {i}' for i in range(20))) + summary = cli._summarize(error) + self.assertTrue(summary.startswith('Job failed.\n line 0')) + self.assertIn('line 8', summary) + self.assertNotIn('line 9', summary) + self.assertTrue(summary.endswith('...')) + + def test_summarize_embedded_python_traceback(self): + # The shape of a Prism failure: context, then the worker's traceback. + error = RuntimeError( + 'Pipeline job-001 failed in state FAILED: bundle b failed:' + 'Traceback (most recent call last):\n' + ' File "common.py", line 1, in process\n' + ' return fn(x)\n' + 'ZeroDivisionError: division by zero\n\n' + 'During handling of the above exception, another exception ' + 'occurred:\n\n' + 'Traceback (most recent call last):\n' + ' File "common.py", line 2, in process\n' + "ZeroDivisionError: division by zero [while running 'Map']\n") + self.assertEqual( + cli._summarize(error), + 'Pipeline job-001 failed in state FAILED: bundle b failed:\n' + " ZeroDivisionError: division by zero [while running 'Map']") + + def test_summarize_names_the_type_of_direct_errors(self): + self.assertEqual( + cli._summarize(ZeroDivisionError('division by zero')), + 'ZeroDivisionError: division by zero') + self.assertEqual(cli._summarize(RuntimeError('boom')), 'boom') + + +if __name__ == '__main__': + unittest.main() diff --git a/sdks/python/setup.py b/sdks/python/setup.py index 9f9618e30fbc..804514ebcc94 100644 --- a/sdks/python/setup.py +++ b/sdks/python/setup.py @@ -433,6 +433,10 @@ def get_portability_package_data(): ] }, ext_modules=extensions, + entry_points={ + # Experimental; see apache_beam/cli/main.py. + 'console_scripts': ['beam = apache_beam.cli.main:main'], + }, install_requires=[ 'cryptography>=50.0.0', 'fastavro>=0.23.6,<2',