From 593b04500844d178906389a9dd99a06ea090f54d Mon Sep 17 00:00:00 2001 From: Jignesh Shah Date: Fri, 6 Nov 2026 15:07:02 +0000 Subject: [PATCH] [Python] Add beam pipe to the experimental beam command (#40891) beam pipe --from SOURCE --to SINK reads CSV, JSON, Parquet or Avro files and writes them as CSV, JSON or Avro by building a two-step Beam YAML pipeline. The format comes from the file extension (also behind a glob or shard suffix such as users.parquet*) unless a flag sets it. Existing output is never overwritten, the files written are listed, and --dry-run prints the pipeline as YAML to save and run with beam run. --- CHANGES.md | 2 +- sdks/python/apache_beam/cli/main.py | 227 ++++++++++++++++-- sdks/python/apache_beam/cli/main_test.py | 179 +++++++++++++- sdks/python/apache_beam/cli/pipe_spec.py | 133 ++++++++++ sdks/python/apache_beam/cli/pipe_spec_test.py | 104 ++++++++ 5 files changed, 623 insertions(+), 22 deletions(-) create mode 100644 sdks/python/apache_beam/cli/pipe_spec.py create mode 100644 sdks/python/apache_beam/cli/pipe_spec_test.py diff --git a/CHANGES.md b/CHANGES.md index 1eb1229261ad..0b7b38a3e4ba 100644 --- a/CHANGES.md +++ b/CHANGES.md @@ -95,7 +95,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 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. `beam pipe --from SOURCE --to SINK` reads CSV, JSON, Parquet or Avro files and writes CSV, JSON or Avro, and `--dry-run` prints the equivalent pipeline to save and run with `beam run` ([#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/main.py b/sdks/python/apache_beam/cli/main.py index b83060c0e2d3..7e5c93ec5ff4 100644 --- a/sdks/python/apache_beam/cli/main.py +++ b/sdks/python/apache_beam/cli/main.py @@ -23,7 +23,12 @@ 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: +`beam pipe --from SOURCE --to SINK` reads files in one format and writes them +in another, without writing a pipeline first. It builds a two-step Beam YAML +pipeline (see `pipe_spec`) and runs it like `beam run`. `--dry-run` prints +that pipeline instead, ready to save, edit and `beam run`. + +On top of that, both give 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. """ @@ -34,18 +39,31 @@ import sys import traceback +import yaml + import apache_beam as beam +from apache_beam.cli import pipe_spec +from apache_beam.io.filesystems import FileSystems 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 codes, also listed in `--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. +_EXIT_CODES = """ +exit codes: + 0 the pipeline finished (or --dry-run checked it, or --no_wait_until_finish submitted it) + 1 the pipeline failed + 2 usage error, or the pipeline could not be built + 4 a required package is missing + 130 interrupted (Ctrl-C) +""" + _RUN_EPILOG = """\ examples: beam run pipeline.yaml @@ -54,23 +72,32 @@ Every other --flag is passed to the pipeline as a Beam pipeline option. Write them after FILE, as --flag=value. +""" + _EXIT_CODES -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) -""" +_PIPE_EPILOG = """\ +examples: + beam pipe --from users.csv --to out/users.json + beam pipe --from 'data/*.parquet' --to out/data --to-format avro + beam pipe --from users.csv --to out/users.json --dry-run > pipeline.yaml + +Formats: csv, json and avro can be read and written; parquet can be read. +Paths can be local or on any Beam filesystem (gs://, s3://, ...). The +output path is a prefix: Beam writes sharded files that start with it, and +existing output is never overwritten. +Every other --flag is passed to the pipeline as a Beam pipeline option. +""" + _EXIT_CODES def main(argv=None): """Runs the `beam` command and returns its exit code.""" + argv = sys.argv[1:] if argv is None else argv 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) + if args.command == 'run': + return _run(args.file, pipeline_args, args.verbose) + return _pipe(args, pipeline_args, argv[argv.index('pipe') + 1:]) except _CliError as e: print(f'error: {e}', file=sys.stderr) return e.exit_code @@ -104,6 +131,50 @@ def _build_parser(): 'filesystem such as gs://') run.add_argument( '--verbose', action='store_true', help='show tracebacks for errors') + + pipe = commands.add_parser( + 'pipe', + help='read files in one format and write them in another', + description=( + 'Read files in one format and write them in another, by building ' + 'and running a two-step Beam YAML pipeline.'), + epilog=_PIPE_EPILOG, + formatter_class=argparse.RawDescriptionHelpFormatter, + allow_abbrev=False) + pipe.add_argument( + '--from', + dest='source', + required=True, + metavar='SOURCE', + help='files to read: a path or glob') + pipe.add_argument( + '--to', + dest='sink', + required=True, + metavar='SINK', + help='where to write: a path prefix') + pipe.add_argument( + '--format', + metavar='FMT', + help='format of both sides; overridden by --from-format / --to-format') + pipe.add_argument( + '--from-format', + metavar='FMT', + help='format of the source (default: from the file extension)') + pipe.add_argument( + '--to-format', + metavar='FMT', + help='format of the sink (default: from the file extension)') + # `--dry_run` is accepted and consumed too: Beam has a test-only pipeline + # option of that name, which some runners honor and others ignore. Passed + # through, it could run the pipeline the user asked not to run. + pipe.add_argument( + '--dry-run', + '--dry_run', + action='store_true', + help='check the pipeline and print it as YAML instead of running it') + pipe.add_argument( + '--verbose', action='store_true', help='show tracebacks for errors') return parser @@ -125,6 +196,39 @@ def _run(file, pipeline_args, verbose): hint = '\n pipeline options go after FILE, written as --flag=value' raise _CliError(f'no such file: {file!r}{hint}', EXIT_USAGE) + return _run_yaml(['--yaml_pipeline_file=' + file] + pipeline_args, verbose) + + +def _pipe(args, pipeline_args, pipe_argv): + """Runs `beam pipe`; `pipe_argv` is its command line, for --dry-run.""" + try: + pipe = pipe_spec.build_pipe_spec( + args.source, + args.sink, + args.from_format or args.format, + args.to_format or args.format) + except pipe_spec.PipeSpecError as e: + raise _CliError(str(e), EXIT_USAGE) from None + + # A dry run makes every check a real run makes, including building the + # pipeline, so that it predicts whether `beam run` on its output will work. + _refuse_existing_output(pipe.output_prefix) + return _run_yaml( + ['--yaml_pipeline=' + yaml.safe_dump(pipe.spec, sort_keys=False)] + + pipeline_args, + args.verbose, + local_paths=[args.source, args.sink], + dry_run_yaml=pipe_spec.to_yaml(pipe.spec, pipe_argv) + if args.dry_run else None, + output_prefix=pipe.output_prefix) + + +def _run_yaml( + yaml_main_argv, + verbose, + local_paths=(), + dry_run_yaml=None, + output_prefix=None): # 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. @@ -138,17 +242,16 @@ def _run(file, pipeline_args, verbose): 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. + # that build errors and pipeline failures get different exit codes, and so + # that `pipe --dry-run` can stop after building. 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)) + yaml_main.build_pipeline_components_from_argv(yaml_main_argv)) pipeline = beam.Pipeline(options=options, display_data=display_data) pipeline.__enter__() constructor(pipeline) @@ -165,6 +268,19 @@ def _run(file, pipeline_args, verbose): raise _CliError( f'could not build the pipeline: {e}', EXIT_USAGE) from None + if local_paths: + _warn_if_local_paths(local_paths, options) + if dry_run_yaml is not None: + # The YAML goes to stdout and every message to stderr, so that + # `> pipeline.yaml` captures exactly the pipeline. + print(dry_run_yaml, end='') + print( + 'Dry run: the pipeline was built and checked but not run. To run ' + 'it, save the YAML (for example with > pipeline.yaml) and use: ' + 'beam run pipeline.yaml', + file=sys.stderr) + return EXIT_OK + no_wait = options.view_as(StandardOptions).no_wait_until_finish state = None try: @@ -199,9 +315,81 @@ def _run(file, pipeline_args, verbose): raise _CliError(f'pipeline finished in state {state}', EXIT_PIPELINE_FAILED) print(f'Done. Ran {_describe_runner(options, result)}.') + if output_prefix is not None: + written = _matching_files(output_prefix) + print(f'Wrote {len(written)} file{"" if len(written) == 1 else "s"}:') + for path in written: + print(' ' + path) return EXIT_OK +def _refuse_existing_output(prefix): + """Fails if any file starts with the output prefix: pipe never overwrites.""" + existing = _matching_files(prefix) + if existing: + shown = '\n'.join(' ' + path for path in existing[:5]) + more = f'\n ... and {len(existing) - 5} more' if len(existing) > 5 else '' + raise _CliError( + f'output already exists for {prefix!r}; beam pipe never overwrites ' + f'data:\n{shown}{more}\n' + ' remove these files or choose another --to path', + EXIT_USAGE) + + +def _matching_files(prefix): + # LocalFileSystem.match calls os.path.dirname(prefix), which is '' for a + # bare filename in the current directory and causes _list('') to return no + # files. Prepending './' (or '.\' on Windows) gives dirname '.' so files in + # the current directory are found. + dot = ( + '.' + + os.sep if '://' not in prefix and not os.path.dirname(prefix) else '') + try: + match = FileSystems.match([dot + prefix + '*'])[0] + except ValueError as e: + # No filesystem handles the scheme, for example gs:// without the [gcp] + # extra. Beam's message names the extra to install. + raise _CliError(str(e), EXIT_MISSING_PACKAGE) from None + return sorted( + m.path[len(dot):] if dot and m.path.startswith(dot) else m.path + for m in match.metadata_list) + + +def _runner_name(runner): + """'DirectRunner', 'direct', None -> 'direct'; 'DataflowRunner' -> ... + + None means the SDK's default runner, DirectRunner. + """ + name = str(runner or 'DirectRunner').rsplit('.', 1)[-1].lower() + return name[:-len('runner')] if name.endswith('runner') else name + + +# Runners that execute on this machine, by _runner_name. +_LOCAL_RUNNERS = frozenset( + ['direct', 'switchingdirect', 'bundlebaseddirect', 'prism', 'fnapi']) + + +def _warn_if_local_paths(paths, options): + """Warns when a runner that isn't local would open local files. + + On a cluster the runner's workers open the files, so a local or relative + path refers to each worker's own disk or container. The job can then + succeed while reading nothing, or write where the user never looks. + """ + runner = options.view_as(StandardOptions).runner + local = [ + path for path in paths + if '://' not in path or path.lower().startswith('file://') + ] + if local and _runner_name(runner) not in _LOCAL_RUNNERS: + print( + f'warning: {", ".join(local)}: local paths are opened by the ' + f"{runner} runner's workers, not by this machine; use storage every " + 'worker can reach (for example gs://, s3:// or hdfs://, or a shared ' + 'mount at the same absolute path)', + file=sys.stderr) + + def _maybe_print_traceback(verbose): if verbose: traceback.print_exc() @@ -268,8 +456,7 @@ def _describe_runner(options, result): 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'): + if _runner_name(runner) != 'direct': return f'on {runner}' if 'portable_runner' in type(result).__module__: return 'locally on Prism' diff --git a/sdks/python/apache_beam/cli/main_test.py b/sdks/python/apache_beam/cli/main_test.py index 15c6488d439c..f7a99e874a9a 100644 --- a/sdks/python/apache_beam/cli/main_test.py +++ b/sdks/python/apache_beam/cli/main_test.py @@ -21,12 +21,14 @@ import json import os import shutil +import sys import tempfile import unittest from unittest import mock import apache_beam as beam from apache_beam.cli import main as cli +from apache_beam.io.filesystems import FileSystems from apache_beam.options.pipeline_options import PipelineOptions from apache_beam.utils import subprocess_server @@ -89,7 +91,14 @@ def _read_lines(prefix): class ParserTest(unittest.TestCase): def test_usage_errors_exit_with_2(self): - for argv in ([], ['run'], ['nosuchcommand', 'x']): + for argv in ( + [], + ['run'], + ['pipe'], + ['pipe', '--from', 'a.csv'], + ['pipe', '--to', 'b.csv'], + ['nosuchcommand', 'x'], + ): with self.subTest(argv=argv): with contextlib.redirect_stderr(io.StringIO()): with self.assertRaises(SystemExit) as e: @@ -172,6 +181,14 @@ def test_pipeline_failure_exits_with_1(self): self.assertIn('ZeroDivisionError', message) self.assertNotIn('Traceback', message) + # Also check when `wait_until_finish` returns a non-DONE state directly. + result = mock.Mock() + result.wait_until_finish.return_value = cli.PipelineState.FAILED + with mock.patch.object(beam.Pipeline, 'run', return_value=result): + code, _, err = _run_cli(['run', path]) + self.assertEqual(code, cli.EXIT_PIPELINE_FAILED) + self.assertIn('pipeline finished in state FAILED', err) + def test_invalid_pipeline_exits_with_2(self): path = self._pipeline_file(INVALID_PIPELINE) code, _, err = _run_cli(['run', path]) @@ -212,6 +229,11 @@ def test_other_flags_are_passed_through_in_order(self): def test_missing_package_exits_with_4(self): path = self._pipeline_file(PIPELINE) + with mock.patch.dict(sys.modules, {'apache_beam.yaml': None}): + code, _, err = _run_cli(['run', path]) + self.assertEqual(code, cli.EXIT_MISSING_PACKAGE) + self.assertIn('apache-beam[yaml]', err) + with mock.patch.object(yaml_main, 'build_pipeline_components_from_argv', side_effect=ImportError('No module named x')): @@ -327,7 +349,161 @@ def test_ctrl_c_cancels_and_exits_with_130(self): self.assertIn('Cancelled the pipeline.', err) +@unittest.skipIf( + yaml_main is None or jsonschema is None, 'requires the [yaml] extra') +class PipeTest(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\n') + self.output = os.path.join(self.tmp, 'out', 'users.json') + + def test_pipes_csv_to_json_and_lists_the_files(self): + code, out, _ = _run_cli([ + 'pipe', + '--from', + self.input, + '--to', + self.output, + '--runner=DirectRunner' + ]) + self.assertEqual(code, cli.EXIT_OK) + self.assertIn('Wrote 1 file:', out) + self.assertEqual([json.loads(line) for line in _read_lines(self.output)], + [{ + 'name': 'Ada', 'age': 36 + }, { + 'name': 'Linus', 'age': 17 + }]) + + # Also check bare filenames in the current working directory. + cwd = os.getcwd() + self.addCleanup(os.chdir, cwd) + os.chdir(self.tmp) + code, out, _ = _run_cli([ + 'pipe', + '--from', + 'users.csv', + '--to', + 'relative_users.json', + '--runner=DirectRunner' + ]) + self.assertEqual(code, cli.EXIT_OK) + self.assertIn('Wrote 1 file:\n relative_users.json', out) + code, _, err = _run_cli( + ['pipe', '--from', 'users.csv', '--to', 'relative_users.json']) + self.assertEqual(code, cli.EXIT_USAGE) + self.assertIn('beam pipe never overwrites data', err) + + def test_backslash_sequences_in_paths(self): + tricky_input = os.path.join(self.tmp, r'C_\Users\tmp_users.csv') + tricky_output = os.path.join(self.tmp, r'C_\Users\out\users.json') + os.makedirs(os.path.dirname(tricky_input), exist_ok=True) + os.makedirs(os.path.dirname(tricky_output), exist_ok=True) + shutil.copyfile(self.input, tricky_input) + code, out, _ = _run_cli([ + 'pipe', + '--from', + tricky_input, + '--to', + tricky_output, + '--runner=DirectRunner' + ]) + self.assertEqual(code, cli.EXIT_OK) + self.assertIn('Wrote 1 file:', out) + self.assertEqual(len(_read_lines(tricky_output)), 2) + + def test_refuses_existing_output(self): + os.makedirs(os.path.dirname(self.output), exist_ok=True) + for i in range(7): + with open(f'{self.output}-0000{i}-of-00007', 'w') as f: + f.write('keep me\n') + code, _, err = _run_cli(['pipe', '--from', self.input, '--to', self.output]) + self.assertEqual(code, cli.EXIT_USAGE) + self.assertIn('beam pipe never overwrites data', err) + self.assertIn('... and 2 more', err) + with open(f'{self.output}-00000-of-00007') as f: + self.assertEqual(f.read(), 'keep me\n') + + def test_missing_filesystem_extra_exits_with_4(self): + with mock.patch.object( + FileSystems, + 'match', + side_effect=ValueError('Unable to get filesystem from specified path')): + code, _, err = _run_cli( + ['pipe', '--from', self.input, '--to', 's3://bucket/users.json']) + self.assertEqual(code, cli.EXIT_MISSING_PACKAGE) + self.assertIn('Unable to get filesystem', err) + + def test_no_wait_and_failure_and_ctrl_c_exit_codes(self): + result = mock.Mock() + result.state = cli.PipelineState.RUNNING + with mock.patch.object(beam.Pipeline, 'run', return_value=result): + code, out, _ = _run_cli([ + 'pipe', + '--from', + self.input, + '--to', + self.output, + '--no_wait_until_finish' + ]) + self.assertEqual(code, cli.EXIT_OK) + self.assertIn('not waiting for completion', out) + self.assertNotIn('Wrote', out) + + result.wait_until_finish.side_effect = RuntimeError('worker crashed') + with mock.patch.object(beam.Pipeline, 'run', return_value=result): + code, _, err = _run_cli( + ['pipe', '--from', self.input, '--to', self.output]) + self.assertEqual(code, cli.EXIT_PIPELINE_FAILED) + self.assertIn('worker crashed', err) + + result.wait_until_finish.side_effect = KeyboardInterrupt + with mock.patch.object(beam.Pipeline, 'run', return_value=result): + code, _, err = _run_cli( + ['pipe', '--from', self.input, '--to', self.output]) + self.assertEqual(code, cli.EXIT_INTERRUPTED) + self.assertIn('Cancelled the pipeline.', err) + + def test_dry_run_prints_a_pipeline_that_beam_run_runs(self): + for flag in ('--dry-run', '--dry_run'): + with self.subTest(flag=flag): + code, out, err = _run_cli( + ['pipe', '--from', self.input, '--to', self.output, flag]) + self.assertEqual(code, cli.EXIT_OK) + self.assertTrue(out.startswith('# Generated by: beam pipe --from')) + self.assertIn('Dry run:', err) + self.assertEqual(glob.glob(self.output + '*'), []) + + saved = os.path.join(self.tmp, 'pipeline.yaml') + with open(saved, 'w') as f: + f.write(out) + code, _, _ = _run_cli(['run', saved, '--runner=DirectRunner']) + self.assertEqual(code, cli.EXIT_OK) + self.assertEqual(len(_read_lines(self.output)), 2) + + def test_flag_errors_exit_with_2(self): + code, _, err = _run_cli(['pipe', '--from', self.input, '--to', 'out/x']) + self.assertEqual(code, cli.EXIT_USAGE) + self.assertIn('add --to-format', err) + + class HelpersTest(unittest.TestCase): + def test_local_paths_warning_only_for_remote_runners(self): + for argv, warned in [([], False), (['--runner=PrismRunner'], False), + (['--runner=dataflow'], True)]: + with self.subTest(argv=argv): + err = io.StringIO() + with contextlib.redirect_stderr(err): + cli._warn_if_local_paths( + ['users.csv', 'file:///tmp/a.csv', 'gs://b/out'], + PipelineOptions(argv)) + self.assertEqual('users.csv' in err.getvalue(), warned) + self.assertEqual('file:///tmp/a.csv' in err.getvalue(), warned) + self.assertNotIn('gs://b/out', err.getvalue()) + def test_cancel_failure_is_reported(self): result = mock.Mock() result.cancel.side_effect = RuntimeError('job not found') @@ -385,6 +561,7 @@ def test_summarize_names_the_type_of_direct_errors(self): cli._summarize(ZeroDivisionError('division by zero')), 'ZeroDivisionError: division by zero') self.assertEqual(cli._summarize(RuntimeError('boom')), 'boom') + self.assertEqual(cli._summarize(RuntimeError('')), "RuntimeError('')") if __name__ == '__main__': diff --git a/sdks/python/apache_beam/cli/pipe_spec.py b/sdks/python/apache_beam/cli/pipe_spec.py new file mode 100644 index 000000000000..a424f3be0557 --- /dev/null +++ b/sdks/python/apache_beam/cli/pipe_spec.py @@ -0,0 +1,133 @@ +# +# 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. +# + +"""Builds the Beam YAML pipeline that `beam pipe` runs. + +`beam pipe --from SOURCE --to SINK` reads records from one set of files and +writes them to another, in the same or a different format. This module turns +those flags into a two-step Beam YAML pipeline: a read transform and a write +transform from Beam YAML's standard IO (`ReadFromCsv`, `WriteToJson`, ...). +The read and the write are chosen independently; Beam YAML's schemas connect +them. + +It never builds or runs a pipeline, so it is tested on plain data, and +`beam pipe --dry-run` prints exactly what it returns. +""" + +import os +import re +import shlex +from typing import NamedTuple + +import yaml + + +class PipeSpecError(ValueError): + """A problem with `beam pipe`'s flags, reported to the user as-is.""" + + +class PipeSpec(NamedTuple): + """The pipeline for one `beam pipe` command.""" + spec: dict + # The sink path. Beam file sinks treat it as a prefix and write sharded + # files that start with it, so the command checks for and lists files + # matching it. + output_prefix: str + + +# Format name -> suffix of the Beam YAML transforms (ReadFrom, WriteTo). +_READ_FORMATS = { + 'csv': 'Csv', 'json': 'Json', 'parquet': 'Parquet', 'avro': 'Avro' +} +_WRITE_FORMATS = {'csv': 'Csv', 'json': 'Json', 'avro': 'Avro'} +_EXTENSIONS = { + '.csv': 'csv', + '.json': 'json', + '.jsonl': 'json', + '.parquet': 'parquet', + '.avro': 'avro', +} +# Schemes handled by Beam filesystems. Anything else with '://' is rejected +# here with a clear message, rather than failing inside the pipeline. +_FILE_SCHEMES = ('gs', 's3', 'azfs', 'hdfs', 'file') + + +def build_pipe_spec(source, sink, from_format=None, to_format=None): + """Returns the PipeSpec for `beam pipe --from source --to sink`. + + Args: + source: The `--from` path or glob, on any Beam filesystem. + sink: The `--to` path prefix, on any Beam filesystem. + from_format: The source format; defaults to the file extension. + to_format: The sink format; defaults to the file extension. + + Raises: + PipeSpecError: if a location or format is invalid. The message names the + flag that fixes it. + """ + read = _transform( + 'ReadFrom', _READ_FORMATS, source, from_format, '--from-format') + write = _transform('WriteTo', _WRITE_FORMATS, sink, to_format, '--to-format') + spec = {'pipeline': {'type': 'chain', 'transforms': [read, write]}} + return PipeSpec(spec, sink) + + +# A trailing run of glob characters, dashes and Beam shard suffixes +# ("-00000-of-00001"), as in `out/users.json*` or `users.json-00000-of-00001`. +_SHARD_OR_GLOB_SUFFIX = re.compile(r'(-\d+-of-\d+|[-*?])+$') + + +def _extension(path): + """The file extension of a path or glob, for example '.json'. + + Beam file sinks write sharded files that start with the output path, so + reading one command's output in another uses a glob such as + `out/users.json*`. The extension is the one before that suffix. + """ + name = _SHARD_OR_GLOB_SUFFIX.sub( + '', os.path.basename(path.replace('\\', '/'))) + return os.path.splitext(name)[1].lower() + + +def _transform(verb, formats, path, fmt, flag): + if '://' in path and path.split('://', 1)[0].lower() not in _FILE_SCHEMES: + schemes = ', '.join(s + '://' for s in _FILE_SCHEMES) + raise PipeSpecError( + f'unsupported location {path!r}: use a local path or one of ' + f'{schemes}') + # Never guess: without a flag the extension decides, and an unknown + # extension is an error that names the flag to add. + fmt = (fmt or _EXTENSIONS.get(_extension(path), '')).lower() + if not fmt: + raise PipeSpecError( + f'cannot tell the format of {path!r} (no known file extension)\n' + f' add {flag}, for example: {flag} json') + if fmt not in formats: + raise PipeSpecError( + f'unsupported format {fmt!r} for {path!r}; use one of: ' + f'{", ".join(formats)}') + return {'type': verb + formats[fmt], 'config': {'path': path}} + + +def to_yaml(spec, argv): + """Renders a pipeline as a Beam YAML file to save and run with `beam run`. + + The file starts with a comment showing the command that generated it, so a + saved pipeline records where it came from. + """ + header = '# Generated by: beam pipe ' + shlex.join(argv) + return header + '\n' + yaml.safe_dump(spec, sort_keys=False) diff --git a/sdks/python/apache_beam/cli/pipe_spec_test.py b/sdks/python/apache_beam/cli/pipe_spec_test.py new file mode 100644 index 000000000000..b65015198420 --- /dev/null +++ b/sdks/python/apache_beam/cli/pipe_spec_test.py @@ -0,0 +1,104 @@ +# +# 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 unittest + +import yaml + +from apache_beam.cli import pipe_spec +from apache_beam.cli.pipe_spec import PipeSpecError +from apache_beam.cli.pipe_spec import build_pipe_spec + +CSV_TO_AVRO = ''' +pipeline: + type: chain + transforms: + - type: ReadFromCsv + config: {path: users.csv} + - type: WriteToAvro + config: {path: out/users.avro} +''' + + +class BuildPipeSpecTest(unittest.TestCase): + def test_csv_to_avro(self): + pipe = build_pipe_spec('users.csv', 'out/users.avro') + self.assertEqual(pipe.spec, yaml.safe_load(CSV_TO_AVRO)) + self.assertEqual(pipe.output_prefix, 'out/users.avro') + + def test_formats_come_from_extensions(self): + for path, transform in [ + ('a.csv', 'ReadFromCsv'), + ('logs/*.json', 'ReadFromJson'), + ('a.jsonl', 'ReadFromJson'), + ('gs://bucket/a.parquet', 'ReadFromParquet'), + ('GS://bucket/a.parquet', 'ReadFromParquet'), + ('a.AVRO', 'ReadFromAvro'), + # Output of an earlier command, read back. + ('out/users.parquet*', 'ReadFromParquet'), + ('out/users.avro-*', 'ReadFromAvro'), + ('out/users.json-00000-of-00001', 'ReadFromJson'), + # Windows backslash paths. + (r'C:\Users\Ada\users.csv', 'ReadFromCsv'), + (r'C:\Users\Ada\out\users.parquet*', 'ReadFromParquet'), + ]: + with self.subTest(path=path): + pipe = build_pipe_spec(path, 'out.csv') + self.assertEqual( + pipe.spec['pipeline']['transforms'][0]['type'], transform) + + def test_explicit_formats_override_extensions(self): + pipe = build_pipe_spec('data.txt', 'out', 'csv', 'json') + self.assertEqual([t['type'] for t in pipe.spec['pipeline']['transforms']], + ['ReadFromCsv', 'WriteToJson']) + + def test_errors_name_the_fix(self): + cases = [ + (('users', 'out.csv'), {}, 'add --from-format'), + ((r'C:\Users\ada.smith\users', 'out.csv'), {}, 'add --from-format'), + (('a.csv', 'out/users'), {}, 'add --to-format'), + (('a.csv', 'b.csv'), { + 'from_format': 'xml' + }, "unsupported format 'xml'"), + (('a.csv', 'out/users.parquet'), {}, + "unsupported format 'parquet' .* use one of: csv, json, avro"), + (('kafka://broker:9092/topic', 'b.csv'), {}, 'unsupported location'), + ] + for args, kwargs, message in cases: + with self.subTest(args=args): + with self.assertRaisesRegex(PipeSpecError, message): + build_pipe_spec(*args, **kwargs) + + +class ToYamlTest(unittest.TestCase): + def test_starts_with_the_command_and_round_trips(self): + for src, dst in [ + ('users.csv', 'out/users.avro'), + ( + r'C:\Users\runneradmin\AppData\Local\Temp\tmp123\users.csv', + r'C:\Users\runneradmin\AppData\Local\Temp\tmp123\out\users.json', + ), + ]: + with self.subTest(src=src): + pipe = build_pipe_spec(src, dst) + text = pipe_spec.to_yaml(pipe.spec, ['--from', src, '--to', dst]) + self.assertTrue(text.startswith('# Generated by: beam pipe --from ')) + self.assertEqual(yaml.safe_load(text), pipe.spec) + + +if __name__ == '__main__': + unittest.main()