如何將數據庫連接傳遞給所有 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 之外的其他參數會很方便。 這是一個實現及其測試。