[Python] Add beam pipe to the experimental beam command (#40891)

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)).

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'