serve.go

Documentation: oss.terrastruct.com/d2/d2plugin

     1  package d2plugin
     2  
     3  import (
     4  	"context"
     5  	"encoding/json"
     6  	"errors"
     7  	"fmt"
     8  	"io"
     9  
    10  	"github.com/spf13/pflag"
    11  
    12  	"oss.terrastruct.com/d2/d2graph"
    13  	"oss.terrastruct.com/util-go/xmain"
    14  )
    15  
    16  // Serve returns a xmain.RunFunc that will invoke the plugin p as necessary to service the
    17  // calling d2 CLI.
    18  //
    19  // See implementation of d2plugin-dagre in the ./cmd directory.
    20  //
    21  // Also see execPlugin in exec.go for the d2 binary plugin protocol.
    22  func Serve(p Plugin) xmain.RunFunc {
    23  	return func(ctx context.Context, ms *xmain.State) (err error) {
    24  		if !ms.Opts.Flags.Parsed() {
    25  			fs, err := p.Flags(ctx)
    26  			if err != nil {
    27  				return err
    28  			}
    29  			for _, f := range fs {
    30  				f.AddToOpts(ms.Opts)
    31  			}
    32  			err = ms.Opts.Flags.Parse(ms.Opts.Args)
    33  			if !errors.Is(err, pflag.ErrHelp) && err != nil {
    34  				return xmain.UsageErrorf("failed to parse flags: %v", err)
    35  			}
    36  			if errors.Is(err, pflag.ErrHelp) {
    37  				// At some point we want to write a friendly help.
    38  				return info(ctx, p, ms)
    39  			}
    40  		}
    41  
    42  		if len(ms.Opts.Flags.Args()) < 1 {
    43  			return xmain.UsageErrorf("expected first argument to be subcmd name")
    44  		}
    45  
    46  		err = HydratePluginOpts(ctx, ms, p)
    47  		if err != nil {
    48  			return err
    49  		}
    50  
    51  		subcmd := ms.Opts.Flags.Arg(0)
    52  		switch subcmd {
    53  		case "info":
    54  			return info(ctx, p, ms)
    55  		case "flags":
    56  			return flags(ctx, p, ms)
    57  		case "layout":
    58  			return layout(ctx, p, ms)
    59  		case "postprocess":
    60  			return postProcess(ctx, p, ms)
    61  		default:
    62  			return xmain.UsageErrorf("unrecognized command: %s", subcmd)
    63  		}
    64  	}
    65  }
    66  
    67  func info(ctx context.Context, p Plugin, ms *xmain.State) error {
    68  	info, err := p.Info(ctx)
    69  	if err != nil {
    70  		return err
    71  	}
    72  	b, err := json.Marshal(info)
    73  	if err != nil {
    74  		return err
    75  	}
    76  	_, err = ms.Stdout.Write(b)
    77  	if err != nil {
    78  		return err
    79  	}
    80  	return nil
    81  }
    82  
    83  func flags(ctx context.Context, p Plugin, ms *xmain.State) error {
    84  	flags, err := p.Flags(ctx)
    85  	if err != nil {
    86  		return err
    87  	}
    88  	b, err := json.Marshal(flags)
    89  	if err != nil {
    90  		return err
    91  	}
    92  	_, err = ms.Stdout.Write(b)
    93  	if err != nil {
    94  		return err
    95  	}
    96  	return nil
    97  }
    98  
    99  func layout(ctx context.Context, p Plugin, ms *xmain.State) error {
   100  	in, err := io.ReadAll(ms.Stdin)
   101  	if err != nil {
   102  		return err
   103  	}
   104  	var g d2graph.Graph
   105  	if err := d2graph.DeserializeGraph(in, &g); err != nil {
   106  		return fmt.Errorf("failed to unmarshal input to graph: %s", in)
   107  	}
   108  	err = p.Layout(ctx, &g)
   109  	if err != nil {
   110  		return err
   111  	}
   112  	b, err := d2graph.SerializeGraph(&g)
   113  	if err != nil {
   114  		return err
   115  	}
   116  	_, err = ms.Stdout.Write(b)
   117  	if err != nil {
   118  		return err
   119  	}
   120  	return nil
   121  }
   122  
   123  func postProcess(ctx context.Context, p Plugin, ms *xmain.State) error {
   124  	in, err := io.ReadAll(ms.Stdin)
   125  	if err != nil {
   126  		return err
   127  	}
   128  
   129  	out, err := p.PostProcess(ctx, in)
   130  	if err != nil {
   131  		return err
   132  	}
   133  
   134  	_, err = ms.Stdout.Write(out)
   135  	if err != nil {
   136  		return err
   137  	}
   138  	return nil
   139  }
   140
View as plain text