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

Question

所以我创建了一个cobra命令

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

我建议采用以下方法：

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

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

Answer 2

您可以创建一个方法来初始化数据库，并在RUN参数中使用它。尝试以下操作：

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函数也适用于其他命令，则可以提取该函数本身。

var versionCmd = &cobra.Command{
.....
Run: runFunc,
}
func runFunc(cmd *cobra.Command, args []string) {
// do everythig here
}