test cobra

This commit is contained in:
wangchuxiao 2023-03-08 11:45:54 +08:00
parent 60876e865b
commit e55167d0bf
10 changed files with 223 additions and 144 deletions

View File

@ -15,17 +15,16 @@ import (
) )
func main() { func main() {
rootCmd := cmd.NewRootCmd() apiCmd := cmd.NewApiCmd()
rootCmd.AddPortFlag() apiCmd.AddPortFlag()
rootCmd.AddRunE(run) apiCmd.AddApi(run)
if err := rootCmd.Execute(); err != nil { if err := apiCmd.Execute(); err != nil {
fmt.Println(err) fmt.Println(err)
os.Exit(1) os.Exit(1)
} }
} }
func run(rootCmd cmd.RootCmd) error { func run(port int) error {
port := rootCmd.GetPortFlag()
if port == 0 { if port == 0 {
port = config.Config.Api.GinPort[0] port = config.Config.Api.GinPort[0]
} }

View File

@ -1,108 +1,22 @@
package main package main
import ( import (
"OpenIM/internal/tools"
"OpenIM/pkg/common/cmd" "OpenIM/pkg/common/cmd"
"context"
"fmt" "fmt"
"github.com/spf13/cobra"
"os" "os"
) )
var seqCmd = &cobra.Command{
Use: "seq",
Short: "seq operation",
RunE: func(cmdLines *cobra.Command, args []string) error {
msgTool, err := tools.InitMsgTool()
if err != nil {
return err
}
userID, _ := cmdLines.Flags().GetString("userID")
superGroupID, _ := cmdLines.Flags().GetString("superGroupID")
fixAll, _ := cmdLines.Flags().GetBool("fixAll")
ctx := context.Background()
switch {
case cmdLines.Parent() == getCmd:
switch {
case userID != "":
msgTool.ShowUserSeqs(ctx, userID)
case superGroupID != "":
msgTool.ShowSuperGroupSeqs(ctx, superGroupID)
}
case cmdLines.Parent() == fixCmd:
switch {
case userID != "":
_, _, err = msgTool.GetAndFixUserSeqs(ctx, userID)
case superGroupID != "":
err = msgTool.FixGroupSeq(ctx, userID)
case fixAll:
err = msgTool.FixAllSeq(ctx)
}
}
return err
},
}
var msgCmd = &cobra.Command{
Use: "msg",
Short: "msg operation",
RunE: func(cmdLines *cobra.Command, args []string) error {
msgTool, err := tools.InitMsgTool()
if err != nil {
return err
}
userID, _ := cmdLines.Flags().GetString("userID")
superGroupID, _ := cmdLines.Flags().GetString("superGroupID")
clearAll, _ := cmdLines.Flags().GetBool("clearAll")
ctx := context.Background()
switch {
case cmdLines.Parent() == getCmd:
switch {
case userID != "":
msgTool.ShowUserSeqs(ctx, userID)
case superGroupID != "":
msgTool.ShowSuperGroupSeqs(ctx, superGroupID)
}
case cmdLines.Parent() == clearCmd:
switch {
case userID != "":
msgTool.ClearUsersMsg(ctx, []string{userID})
case superGroupID != "":
msgTool.ClearSuperGroupMsg(ctx, []string{superGroupID})
case clearAll:
msgTool.AllUserClearMsgAndFixSeq()
}
}
return nil
},
}
var getCmd = &cobra.Command{
Use: "get",
Short: "get operation",
}
var fixCmd = &cobra.Command{
Use: "fix",
Short: "fix seq operation",
}
var clearCmd = &cobra.Command{
Use: "clear",
Short: "clear operation",
}
func main() { func main() {
rootCmd := cmd.NewMsgUtilsCmd() msgUtilsCmd := cmd.NewMsgUtilsCmd()
rootCmd.Command.PersistentFlags().StringP("userID", "u", "", "openIM userID") msgUtilsCmd.AddSuperGroupIDFlag()
rootCmd.Command.PersistentFlags().StringP("groupID", "u", "", "openIM superGroupID") msgUtilsCmd.AddUserIDFlag()
seqCmd.Flags().BoolP("fixAll", "c", false, "openIM fix all seqs") seqCmd := cmd.NewSeqCmd()
msgCmd.Flags().BoolP("clearAll", "c", false, "openIM clear all timeout msgs") msgCmd := cmd.NewMsgCmd()
getCmd.AddCommand(seqCmd, msgCmd) cmd.GetCmd.AddCommand(seqCmd.Command, msgCmd.Command)
fixCmd.AddCommand(seqCmd) cmd.FixCmd.AddCommand(seqCmd.Command)
clearCmd.AddCommand(msgCmd) cmd.GetCmd.AddCommand(msgCmd.Command)
rootCmd.AddCommand(getCmd, fixCmd, clearCmd) msgUtilsCmd.AddCommand(cmd.GetCmd, cmd.FixCmd, cmd.ClearCmd)
if err := rootCmd.Execute(); err != nil { if err := msgUtilsCmd.Execute(); err != nil {
fmt.Println(err) fmt.Println(err)
os.Exit(1) os.Exit(1)
} }

View File

@ -2,17 +2,16 @@ package main
import ( import (
"OpenIM/internal/tools" "OpenIM/internal/tools"
"OpenIM/pkg/common/config" "OpenIM/pkg/common/cmd"
"fmt" "fmt"
"time" "os"
) )
func main() { func main() {
fmt.Println(time.Now(), "start cronTask") cronTaskCmd := cmd.NewCronTaskCmd()
if err := config.InitConfig(""); err != nil { cronTaskCmd.AddRunE(tools.StartCronTask)
panic(err.Error()) if err := cronTaskCmd.Execute(); err != nil {
} fmt.Println(err)
if err := tools.StartCronTask(); err != nil { os.Exit(1)
panic(err.Error())
} }
} }

17
pkg/common/cmd/api.go Normal file
View File

@ -0,0 +1,17 @@
package cmd
import "github.com/spf13/cobra"
type ApiCmd struct {
*RootCmd
}
func NewApiCmd() *ApiCmd {
return &ApiCmd{NewRootCmd()}
}
func (a *ApiCmd) AddApi(f func(port int) error) {
a.Command.RunE = func(cmd *cobra.Command, args []string) error {
return f(a.GetPortFlag())
}
}

View File

@ -0,0 +1,17 @@
package cmd
import "github.com/spf13/cobra"
type CronTaskCmd struct {
*RootCmd
}
func NewCronTaskCmd() *CronTaskCmd {
return &CronTaskCmd{NewRootCmd()}
}
func (c *CronTaskCmd) AddRunE(f func() error) {
c.Command.RunE = func(cmd *cobra.Command, args []string) error {
return f()
}
}

View File

@ -1,18 +1,20 @@
package cmd package cmd
import (
"OpenIM/internal/tools"
"context"
"github.com/spf13/cobra"
)
type MsgUtilsCmd struct { type MsgUtilsCmd struct {
*RootCmd *RootCmd
userID string userID string
userIDFlag bool
superGroupID string superGroupID string
superGroupIDFlag bool
clearAll bool clearAll bool
clearAllFlag bool
fixAll bool fixAll bool
fixAllFlag bool
} }
func NewMsgUtilsCmd() MsgUtilsCmd { func NewMsgUtilsCmd() MsgUtilsCmd {
@ -21,22 +23,135 @@ func NewMsgUtilsCmd() MsgUtilsCmd {
func (m *MsgUtilsCmd) AddUserIDFlag() { func (m *MsgUtilsCmd) AddUserIDFlag() {
m.Command.PersistentFlags().StringP("userID", "u", "", "openIM userID") m.Command.PersistentFlags().StringP("userID", "u", "", "openIM userID")
m.userIDFlag = true
} }
func (m *MsgUtilsCmd) getUserIDFlag() { func (m *MsgUtilsCmd) GetUserIDFlag() string {
m.Command.PersistentFlags().StringP("userID", "u", "", "openIM userID") return m.userID
}
func (m *MsgUtilsCmd) AddGroupIDFlag() {
m.Command.PersistentFlags().StringP("super-groupID", "u", "", "openIM superGroupID")
}
func (m *MsgUtilsCmd) AddClearAllFlag() {
m.Command.PersistentFlags().BoolP("clearAll", "c", false, "openIM clear all timeout msgs")
} }
func (m *MsgUtilsCmd) AddFixAllFlag() { func (m *MsgUtilsCmd) AddFixAllFlag() {
m.Command.PersistentFlags().BoolP("fixAll", "c", false, "openIM fix all seqs") m.Command.PersistentFlags().BoolP("fixAll", "c", false, "openIM fix all seqs")
} }
func (m *MsgUtilsCmd) GetFixAllFlag() bool {
return m.fixAll
}
func (m *MsgUtilsCmd) AddSuperGroupIDFlag() {
m.Command.PersistentFlags().StringP("super-groupID", "u", "", "openIM superGroupID")
}
func (m *MsgUtilsCmd) GetSuperGroupIDFlag() string {
return m.superGroupID
}
func (m *MsgUtilsCmd) AddClearAllFlag() bool {
return m.clearAll
}
func (m *MsgUtilsCmd) GetClearAllFlag() bool {
return m.clearAll
}
type SeqCmd struct {
Command *cobra.Command
}
func (SeqCmd) RunCommand(cmdLines *cobra.Command, args []string) error {
msgTool, err := tools.InitMsgTool()
if err != nil {
return err
}
userID, _ := cmdLines.Flags().GetString("userID")
superGroupID, _ := cmdLines.Flags().GetString("superGroupID")
fixAll, _ := cmdLines.Flags().GetBool("fixAll")
ctx := context.Background()
switch {
case cmdLines.Parent() == GetCmd:
switch {
case userID != "":
msgTool.ShowUserSeqs(ctx, userID)
case superGroupID != "":
msgTool.ShowSuperGroupSeqs(ctx, superGroupID)
}
case cmdLines.Parent() == FixCmd:
switch {
case userID != "":
_, _, err = msgTool.GetAndFixUserSeqs(ctx, userID)
case superGroupID != "":
err = msgTool.FixGroupSeq(ctx, userID)
case fixAll:
err = msgTool.FixAllSeq(ctx)
}
}
return err
}
func NewSeqCmd() SeqCmd {
seqCmd := SeqCmd{&cobra.Command{
Use: "seq",
Short: "seq operation",
}}
seqCmd.Command.Flags().BoolP("fixAll", "c", false, "openIM fix all seqs")
seqCmd.Command.RunE = seqCmd.RunCommand
return seqCmd
}
type MsgCmd struct {
Command *cobra.Command
}
func NewMsgCmd() MsgCmd {
msgCmd := MsgCmd{&cobra.Command{
Use: "msg",
Short: "msg operation",
}}
msgCmd.Command.RunE = msgCmd.RunCommand
msgCmd.Command.Flags().BoolP("clearAll", "c", false, "openIM clear all timeout msgs")
return msgCmd
}
func (*MsgCmd) RunCommand(cmdLines *cobra.Command, args []string) error {
msgTool, err := tools.InitMsgTool()
if err != nil {
return err
}
userID, _ := cmdLines.Flags().GetString("userID")
superGroupID, _ := cmdLines.Flags().GetString("superGroupID")
clearAll, _ := cmdLines.Flags().GetBool("clearAll")
ctx := context.Background()
switch {
case cmdLines.Parent() == GetCmd:
switch {
case userID != "":
msgTool.ShowUserSeqs(ctx, userID)
case superGroupID != "":
msgTool.ShowSuperGroupSeqs(ctx, superGroupID)
}
case cmdLines.Parent() == ClearCmd:
switch {
case userID != "":
msgTool.ClearUsersMsg(ctx, []string{userID})
case superGroupID != "":
msgTool.ClearSuperGroupMsg(ctx, []string{superGroupID})
case clearAll:
msgTool.AllUserClearMsgAndFixSeq()
}
}
return nil
}
var GetCmd = &cobra.Command{
Use: "get",
Short: "get operation",
}
var FixCmd = &cobra.Command{
Use: "fix",
Short: "fix seq operation",
}
var ClearCmd = &cobra.Command{
Use: "clear",
Short: "clear operation",
}

View File

@ -0,0 +1 @@
package cmd

View File

@ -7,12 +7,10 @@ import (
) )
type RootCmd struct { type RootCmd struct {
Command cobra.Command Command cobra.Command
port int port int
portFlag bool
prometheusPort int prometheusPort int
prometheusPortFlag bool
} }
func NewRootCmd() (rootCmd *RootCmd) { func NewRootCmd() (rootCmd *RootCmd) {
@ -22,12 +20,8 @@ func NewRootCmd() (rootCmd *RootCmd) {
Short: "Start the server", Short: "Start the server",
Long: `Start the server`, Long: `Start the server`,
PersistentPreRunE: func(cmd *cobra.Command, args []string) error { PersistentPreRunE: func(cmd *cobra.Command, args []string) error {
//if rootCmd.portFlag {
rootCmd.port = rootCmd.getPortFlag(cmd) rootCmd.port = rootCmd.getPortFlag(cmd)
//}
//if rootCmd.prometheusPortFlag {
rootCmd.prometheusPort = rootCmd.getPrometheusPortFlag(cmd) rootCmd.prometheusPort = rootCmd.getPrometheusPortFlag(cmd)
//}
return rootCmd.getConfFromCmdAndInit(cmd) return rootCmd.getConfFromCmdAndInit(cmd)
}, },
} }
@ -54,7 +48,6 @@ func (r *RootCmd) init() {
func (r *RootCmd) AddPortFlag() { func (r *RootCmd) AddPortFlag() {
r.Command.Flags().IntP(constant.FlagPort, "p", 0, "server listen port") r.Command.Flags().IntP(constant.FlagPort, "p", 0, "server listen port")
r.portFlag = true
} }
func (r *RootCmd) getPortFlag(cmd *cobra.Command) int { func (r *RootCmd) getPortFlag(cmd *cobra.Command) int {
@ -68,7 +61,6 @@ func (r *RootCmd) GetPortFlag() int {
func (r *RootCmd) AddPrometheusPortFlag() { func (r *RootCmd) AddPrometheusPortFlag() {
r.Command.Flags().StringP(constant.PrometheusPort, "pp", "", "server listen port") r.Command.Flags().StringP(constant.PrometheusPort, "pp", "", "server listen port")
r.prometheusPortFlag = true
} }
func (r *RootCmd) getPrometheusPortFlag(cmd *cobra.Command) int { func (r *RootCmd) getPrometheusPortFlag(cmd *cobra.Command) int {
@ -76,6 +68,10 @@ func (r *RootCmd) getPrometheusPortFlag(cmd *cobra.Command) int {
return port return port
} }
func (r *RootCmd) GetPrometheusPortFlag() int {
return r.prometheusPort
}
func (r *RootCmd) getConfFromCmdAndInit(cmdLines *cobra.Command) error { func (r *RootCmd) getConfFromCmdAndInit(cmdLines *cobra.Command) error {
configFolderPath, _ := cmdLines.Flags().GetString(constant.FlagConf) configFolderPath, _ := cmdLines.Flags().GetString(constant.FlagConf)
return config.InitConfig(configFolderPath) return config.InitConfig(configFolderPath)

21
pkg/common/cmd/rpc.go Normal file
View File

@ -0,0 +1,21 @@
package cmd
import (
"OpenIM/pkg/discoveryregistry"
"github.com/spf13/cobra"
"google.golang.org/grpc"
)
type RpcCmd struct {
*RootCmd
}
func NewRpcCmd() *RpcCmd {
return &RpcCmd{NewRootCmd()}
}
func (r *RpcCmd) AddRpc(f func(port, rpcRegisterName string, prometheusPort int, rpcFn func(client discoveryregistry.SvcDiscoveryRegistry, server *grpc.Server) error, options ...grpc.ServerOption) error) {
r.Command.RunE = func(cmd *cobra.Command, args []string) error {
return f(r.port, r.prometheusPort)
}
}

View File

@ -502,6 +502,9 @@ func (c *config) initConfig(config interface{}, configName, configFolderPath str
configFolderPath = DefaultFolderPath configFolderPath = DefaultFolderPath
} }
configPath := filepath.Join(configFolderPath, configName) configPath := filepath.Join(configFolderPath, configName)
defer func() {
fmt.Println("use config", configPath)
}()
_, err := os.Stat(configPath) _, err := os.Stat(configPath)
if err != nil { if err != nil {
if !os.IsNotExist(err) { if !os.IsNotExist(err) {
@ -527,9 +530,6 @@ func (c *config) GetConfFromRegistry(registry discoveryregistry.SvcDiscoveryRegi
} }
func InitConfig(configFolderPath string) error { func InitConfig(configFolderPath string) error {
defer func() {
fmt.Println("use config folder", configFolderPath)
}()
err := Config.initConfig(&Config, FileName, configFolderPath) err := Config.initConfig(&Config, FileName, configFolderPath)
if err != nil { if err != nil {
return err return err