github.com/cayleygraph/cayley@v0.7.7/cmd/cayley/command/repl.go (about)

     1  package command
     2  
     3  import (
     4  	"context"
     5  	"encoding/json"
     6  	"fmt"
     7  	"io/ioutil"
     8  	"os"
     9  	"os/signal"
    10  	"strings"
    11  	"time"
    12  
    13  	"github.com/spf13/cobra"
    14  	"github.com/spf13/viper"
    15  
    16  	"github.com/cayleygraph/cayley/clog"
    17  	"github.com/cayleygraph/cayley/internal/repl"
    18  	"github.com/cayleygraph/cayley/query"
    19  )
    20  
    21  const (
    22  	keyQueryTimeout = "query.timeout"
    23  )
    24  
    25  func getContext() (context.Context, func()) {
    26  	ctx, cancel := context.WithCancel(context.Background())
    27  	ch := make(chan os.Signal, 1)
    28  	signal.Notify(ch, os.Interrupt)
    29  	go func() {
    30  		select {
    31  		case <-ch:
    32  		case <-ctx.Done():
    33  		}
    34  		signal.Stop(ch)
    35  		cancel()
    36  	}()
    37  	return ctx, cancel
    38  }
    39  
    40  func registerQueryFlags(cmd *cobra.Command) {
    41  	langs := query.Languages()
    42  	cmd.Flags().Bool("init", false, "initialize the database before using it")
    43  	cmd.Flags().String("lang", "gizmo", `query language to use ("`+strings.Join(langs, `", "`)+`")`)
    44  	cmd.Flags().DurationP("timeout", "t", 30*time.Second, "elapsed time until an individual query times out")
    45  	viper.BindPFlag(keyQueryTimeout, cmd.Flags().Lookup("timeout"))
    46  	registerLoadFlags(cmd)
    47  }
    48  
    49  func NewReplCmd() *cobra.Command {
    50  	cmd := &cobra.Command{
    51  		Use:   "repl",
    52  		Short: "Drop into a REPL of the given query language.",
    53  		RunE: func(cmd *cobra.Command, args []string) error {
    54  			printBackendInfo()
    55  			p := mustSetupProfile(cmd)
    56  			defer mustFinishProfile(p)
    57  
    58  			h, err := openForQueries(cmd)
    59  			if err != nil {
    60  				return err
    61  			}
    62  			defer h.Close()
    63  
    64  			ctx, cancel := getContext()
    65  			defer cancel()
    66  
    67  			timeout := viper.GetDuration("timeout")
    68  			lang, _ := cmd.Flags().GetString("lang")
    69  			return repl.Repl(ctx, h, lang, timeout)
    70  		},
    71  	}
    72  	registerQueryFlags(cmd)
    73  	return cmd
    74  }
    75  
    76  func NewQueryCmd() *cobra.Command {
    77  	cmd := &cobra.Command{
    78  		Use:     "query",
    79  		Aliases: []string{"qu"},
    80  		Short:   "Run a query in a specified database and print results.",
    81  		RunE: func(cmd *cobra.Command, args []string) error {
    82  			var querystr string
    83  			if len(args) == 0 {
    84  				bytes, err := ioutil.ReadAll(os.Stdin)
    85  				if err != nil {
    86  					return fmt.Errorf("Error occured while reading from stdin : %s.", err)
    87  				}
    88  				querystr = string(bytes)
    89  			} else if len(args) == 1 {
    90  				querystr = args[0]
    91  			} else {
    92  				return fmt.Errorf("Query accepts only one argument, the query string or nothing for reading from stdin.")
    93  			}
    94  			clog.Infof("Query:\n%s", querystr)
    95  			printBackendInfo()
    96  			p := mustSetupProfile(cmd)
    97  			defer mustFinishProfile(p)
    98  
    99  			h, err := openForQueries(cmd)
   100  			if err != nil {
   101  				return err
   102  			}
   103  			defer h.Close()
   104  
   105  			ctx, cancel := getContext()
   106  			defer cancel()
   107  
   108  			timeout := viper.GetDuration("timeout")
   109  			if timeout > 0 {
   110  				ctx, cancel = context.WithTimeout(ctx, timeout)
   111  				defer cancel()
   112  			}
   113  			lang, _ := cmd.Flags().GetString("lang")
   114  			limit, err := cmd.Flags().GetInt("limit")
   115  			if err != nil {
   116  				return err
   117  			}
   118  			enc := json.NewEncoder(os.Stdout)
   119  			it, err := query.Execute(ctx, h, lang, querystr, query.Options{
   120  				Collation: query.JSON,
   121  				Limit:     limit,
   122  			})
   123  			if err != nil {
   124  				return err
   125  			}
   126  			defer it.Close()
   127  			for i := 0; it.Next(ctx) && (limit <= 0 || i < limit); i++ {
   128  				if err = enc.Encode(it.Result()); err != nil {
   129  					return err
   130  				}
   131  			}
   132  			return it.Err()
   133  		},
   134  	}
   135  	registerQueryFlags(cmd)
   136  	cmd.Flags().IntP("limit", "n", 100, "limit a number of results")
   137  	return cmd
   138  }