简体   繁体   中英

How do I pass the database connection to all cobra commands?

So I have a cobra command created

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")
        }
    
       ....
    },
}

you can see the database connection code inside. Now I have to add the same Database connection code in all of the commands. Is there a way I can avoid this?

You can create a method to initialize db and use that inside the RUN arg. Try following:

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
}

Alternatively, if the Run function is gerenic for other commands as well, you can extract that function itself.

var versionCmd = &cobra.Command{
    .....

    Run: runFunc,
}

func runFunc(cmd *cobra.Command, args []string) {
  // do everythig here
}

I'd suggest the following approach:

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
}

It's also really easy to test:

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

Sometimes it's handy to pass other parameters apart from the database object for testing purposes. Here is an implementation and itstest .

The technical post webpages of this site follow the CC BY-SA 4.0 protocol. If you need to reprint, please indicate the site URL or the original address.Any question please contact:yoyou2525@163.com.

 
粤ICP备18138465号  © 2020-2024 STACKOOM.COM