如何将数据库连接传递给所有 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
}

或者，如果Run function 也适用于其他命令，您可以提取 function 本身。

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)
    }
}

有时，为了测试目的，传递数据库 object 之外的其他参数会很方便。 这是一个实现及其测试。