如何将数据库连接传递给所有 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，则可以提取该函数本身。Runvar 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) errorfunc 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)    }}有时，出于测试目的，传递数据库对象以外的其他参数非常方便。下面是一个实现及其测试。