github.com/rohankumardubey/aresdb@v0.0.2-0.20190517170215-e54e3ca06b9c/cmd/arescli/main.go (about)

     1  package main
     2  
     3  import (
     4  	"encoding/json"
     5  	"fmt"
     6  	"github.com/abiosoft/ishell"
     7  	"github.com/fatih/color"
     8  	"github.com/spf13/cobra"
     9  	"io/ioutil"
    10  	"net/http"
    11  	"strings"
    12  )
    13  
    14  const contentType = "application/json"
    15  
    16  type shellContext struct {
    17  	host        string
    18  	port        int
    19  	clusterMode bool
    20  	client      http.Client
    21  }
    22  
    23  // script global context
    24  var ctx shellContext
    25  
    26  func show(c *ishell.Context) {
    27  	if len(c.Args) != 1 {
    28  		c.Println(color.New(color.FgRed).Println("invalid argument for show command"))
    29  	} else {
    30  		arg := c.Args[0]
    31  		switch arg {
    32  		case "tables":
    33  			if ctx.clusterMode {
    34  				// TODO: @shz implement cluster mode with gateway.controller
    35  				c.Println("Not Implemented yet")
    36  			} else {
    37  				// local mode
    38  				resp, err := http.Get(fmt.Sprintf("http://%s:%d/schema/tables", ctx.host, ctx.port))
    39  				if err != nil {
    40  					c.Println(color.New(color.FgRed).Sprintf(err.Error()))
    41  					return
    42  				}
    43  				if resp.StatusCode != http.StatusOK {
    44  					c.Println(color.New(color.FgRed).Sprintf("Got code %d from aresdb server", resp.StatusCode))
    45  					return
    46  				}
    47  				var tables []string
    48  				var data []byte
    49  				data, err = ioutil.ReadAll(resp.Body)
    50  				if err != nil {
    51  					c.Println(color.New(color.FgRed).Sprintf("error reading response: %s", err))
    52  					return
    53  				}
    54  				err = json.Unmarshal(data, &tables)
    55  				if err != nil {
    56  					c.Println(color.New(color.FgRed).Sprintf("error decoding response: %s", err))
    57  					return
    58  				}
    59  				c.Println(strings.Join(tables, " "))
    60  			}
    61  		case "configs":
    62  			c.Printf("%+v\n", ctx)
    63  		}
    64  	}
    65  }
    66  
    67  func aql(c *ishell.Context) {
    68  	c.Println("aql query ending with semicolon ';':")
    69  	lines := c.ReadMultiLines(";")
    70  	lines = strings.Trim(lines, ";")
    71  	resp, err := http.Post(fmt.Sprintf("http://%s:%d/query/aql", ctx.host, ctx.port), contentType, strings.NewReader(lines))
    72  	if err != nil {
    73  		c.Println(color.New(color.FgRed).Sprintf(err.Error()))
    74  		return
    75  	}
    76  	if resp.StatusCode != http.StatusOK {
    77  		c.Println(color.New(color.FgRed).Sprintf("Got code %d from aresdb server", resp.StatusCode))
    78  		return
    79  	}
    80  	var data []byte
    81  	data, err = ioutil.ReadAll(resp.Body)
    82  	if err != nil {
    83  		c.Println(color.New(color.FgRed).Sprintf("error reading response: %s", err))
    84  		return
    85  	}
    86  
    87  	var result map[string]interface{}
    88  	err = json.Unmarshal(data, &result)
    89  	if err != nil {
    90  		c.Println(color.New(color.FgRed).Sprintf("error decoding response: %s", err))
    91  		return
    92  	}
    93  	data, err = json.MarshalIndent(result, "", "  ")
    94  	if err != nil {
    95  		c.Println(color.New(color.FgRed).Sprintf("error formatting response: %s", err))
    96  		return
    97  	}
    98  	c.ShowPaged(string(data))
    99  }
   100  
   101  func Execute() {
   102  
   103  	// ishell shell
   104  	shell := ishell.New()
   105  
   106  	shell.Println("Welcome to AresDB Cli!")
   107  	shell.AddCmd(&ishell.Cmd{
   108  		Name: "show",
   109  		Help: "`show tables` will show all tables in current cluster",
   110  		Func: show,
   111  		Completer: func(args []string) []string {
   112  			return []string{"tables", "configs"}
   113  		},
   114  	})
   115  
   116  	shell.AddCmd(&ishell.Cmd{
   117  		Name: "aql",
   118  		Help: "start a new aql query",
   119  		Func: aql,
   120  	})
   121  
   122  	//TODO: add sql cmd
   123  
   124  	// cobra command
   125  	cmd := &cobra.Command{
   126  		Use:     "arescli",
   127  		Short:   "AresDB cli",
   128  		Long:    "AresDB command line tool to interact with the backend",
   129  		Example: "arescli --host localhost --port 9374",
   130  		Run: func(cmd *cobra.Command, args []string) {
   131  			// read args
   132  			var err error
   133  			ctx.host, err = cmd.Flags().GetString("host")
   134  			if err != nil {
   135  				panic("failed to get aresdb host")
   136  			}
   137  			ctx.port, err = cmd.Flags().GetInt("port")
   138  			if err != nil {
   139  				panic("failed to get aresdb port")
   140  			}
   141  			ctx.clusterMode, err = cmd.Flags().GetBool("cluster")
   142  			if err != nil {
   143  				panic("failed to get aresdb cluster mode config")
   144  			}
   145  
   146  			// config http client
   147  			ctx.client = http.Client{}
   148  
   149  			if len(args) > 1 {
   150  				shell.Process(args[1:]...)
   151  			} else {
   152  				shell.Run()
   153  				shell.Close()
   154  			}
   155  		},
   156  	}
   157  
   158  	cmd.Flags().StringP("host", "", "localhost", "host of aresdb service, in cluster mode, host of controller")
   159  	cmd.Flags().IntP("port", "p", 9374, "port of aresdb service, in cluster mode, port of controller")
   160  	cmd.Flags().BoolP("cluster", "c", false, "whether to use cluster mode")
   161  	cmd.Execute()
   162  }
   163  
   164  func main() {
   165  	Execute()
   166  }