如何将数据库连接传递给所有 Cobra 命令？

Question

所以我创建了一个眼镜蛇命令func init() {    rootCmd.AddCommand(versionCmd)}var versionCmd = &cobra.Command{    Use:   "version",    Short: "Show App version",    Long:  "Show App version.",    Run: func(cmd *cobra.Command, args []string) {        dir, _ := os.Getwd()        db, err := gorm.Open(sqlite.Open(fmt.Sprintf("%s/db/db.sqlite", dir)), &gorm.Config{})        if err != nil {            panic("failed to connect database")        }           ....    },}你可以看到里面的数据库连接代码。现在，我必须在所有命令中添加相同的数据库连接代码。有没有办法避免这种情况？

Answer 1

您可以创建一个方法来初始化 db，并在 RUN arg 中使用该方法。请尝试以下操作：

var versionCmd = &cobra.Command{

    .....

    Run: func(cmd *cobra.Command, args []string) {

       db := initDB()     

       ....

    },

}

func initDB() *gorm.DB {

  dir, _ := os.Getwd()      

  db, err := gorm.Open(sqlite.Open(fmt.Sprintf("%s/db/db.sqlite", dir)), &gorm.Config{})

  if err != nil {

    panic("failed to connect database")       

  }

  

  return db

}

或者，如果该函数对于其他命令也是 gerenic，则可以提取该函数本身。Run

var versionCmd = &cobra.Command{

    .....

    Run: runFunc,

}

func runFunc(cmd *cobra.Command, args []string) {

  // do everythig here

}

Answer 2

我建议采用以下方法：

type RunEFunc func(cmd *cobra.Command, args []string) error

func NewCmd(db *gorm.DB) *cobra.Command {

    cmd := &cobra.Command{

            ...

        RunE: runCmd(db),

    }

    return cmd

}

func runCmd(db *gorm.DB) RunEFunc {

    return func(cmd *cobra.Command, args []string) error {

        // Use db here

}

测试也非常容易：

func TestCmd(t *testing.T) {

    db := // Open db

    cmd := NewCmd(db)

    cmd.SetArgs([]string{"any here"})

    // You can also modify flags

    if err := cmd.Execute(); err != nil {

        t.Fatal(err)

    }

}

有时，出于测试目的，传递数据库对象以外的其他参数非常方便。下面是一个实现及其测试。