Go初心者がcobraを使ってコマンドラインツールを作ってみた話

Adways Advent Calendar 2018 18日目の記事です。

http://blog.engineer.adways.net/entry/advent_calendar_2018


 

こんにちは、最近Goの勉強を始めた高木です。
自分はがっつりチュートリアルをやるよりは実際に何か簡単なものを作ってみた方が覚えるだろう派の人間なので、業務で使えるコマンドラインツールをサンプルで作ってみました。
もちろんGoのチュートリアル、「A Tour of Go」は軽くやりました。

作ったもの

僕の所属している部署のレビューはslack標準のBotを使用してレビュアーをランダムで選定し、assigneeに指定してマージリクエストを出すというフローになっています。
Botを呼び出すトリガーは共通のものを使用しているため、自分が選ばれた場合はやり直す必要があるので正直面倒くさいです。
そこで、レビュアーを選定してマージリクエストのURLを発行するまでを自動でやってくれるコマンドラインツール wao(wonderful automatic operation) を作ることにしました。
(なお、現時点では自分を外した専用のトリガーを各自用意するという考えに誰も気づいていないと仮定します。)

使用したライブラリ

コマンドラインツールを作るライブラリはたくさんありますが、いろいろな方が使用していてドキュメントも充実していることからcobra選びました。
viperは設定ファイルを扱いやすくするライブラリです。同じ作者なので相性は抜群だと思います。
他にもslackやgitlabを操作するライブラリなどを使っていますが、コマンドラインツール関連はこの2つです。

cobraの使い方

試しにhogeというコマンドラインツールを作ってみます。ディレクトリ名がそのままツール名になります。

$ go get https://github.com/spf13/cobra
$ mkdir hoge && cd hoge
$ cobra init
Using config file: /home/vagrant/.cobra.yaml
Your Cobra application is ready at
/home/vagrant/.go/src/hoge

Give it a try by going there and running `go run main.go`.
Add commands to it by running `cobra add [cmdname]`.

cobra initで生成されるファイル群です。

$ tree .
.
├── LICENSE
├── cmd
│   └── root.go
└── main.go
// File: main.go
package main

import "hoge/cmd"

func main() {
    cmd.Execute()
}
// File: cmd/root.go
package cmd

import (
    "fmt"
    "os"

    homedir "github.com/mitchellh/go-homedir"
    "github.com/spf13/cobra"
    "github.com/spf13/viper"
)

var cfgFile string

// rootCmd represents the base command when called without any subcommands
var rootCmd = &cobra.Command{
    Use:   "hoge",
    Short: "A brief description of your application",
    Long: `A longer description that spans multiple lines and likely contains
examples and usage of using your application. For example:

Cobra is a CLI library for Go that empowers applications.
This application is a tool to generate the needed files
to quickly create a Cobra application.`,
    // Uncomment the following line if your bare application
    // has an action associated with it:
    //  Run: func(cmd *cobra.Command, args []string) { },
}

// Execute adds all child commands to the root command and sets flags appropriately.
// This is called by main.main(). It only needs to happen once to the rootCmd.
func Execute() {
    if err := rootCmd.Execute(); err != nil {
        fmt.Println(err)
        os.Exit(1)
    }
}

func init() {
    cobra.OnInitialize(initConfig)

    // Here you will define your flags and configuration settings.
    // Cobra supports persistent flags, which, if defined here,
    // will be global for your application.
    rootCmd.PersistentFlags().StringVar(&cfgFile, "config", "", "config file (default is $HOME/.hoge.yaml)")

    // Cobra also supports local flags, which will only run
    // when this action is called directly.
    rootCmd.Flags().BoolP("toggle", "t", false, "Help message for toggle")
}

// initConfig reads in config file and ENV variables if set.
func initConfig() {
    if cfgFile != "" {
        // Use config file from the flag.
        viper.SetConfigFile(cfgFile)
    } else {
        // Find home directory.
        home, err := homedir.Dir()
        if err != nil {
            fmt.Println(err)
            os.Exit(1)
        }

        // Search config in home directory with name ".hoge" (without extension).
        viper.AddConfigPath(home)
        viper.SetConfigName(".hoge")
    }

    viper.AutomaticEnv() // read in environment variables that match

    // If a config file is found, read it in.
    if err := viper.ReadInConfig(); err == nil {
        fmt.Println("Using config file:", viper.ConfigFileUsed())
    }
}

サブコマンドの追加はcobra add コマンド名で出来ます。
helloというサブコマンドを作ってみます。

$ cobra add hello
Using config file: /home/vagrant/.cobra.yaml
hello created at /home/vagrant/.go/src/hoge/cmd/hello.go

cmd/hello.goというファイルが生成されました。

// File: cmd/hello.go
package cmd

import (
    "fmt"

    "github.com/spf13/cobra"
)

// helloCmd represents the hello command
var helloCmd = &cobra.Command{
    Use:   "hello",
    Short: "A brief description of your command",
    Long: `A longer description that spans multiple lines and likely contains examples
and usage of using your command. For example:

Cobra is a CLI library for Go that empowers applications.
This application is a tool to generate the needed files
to quickly create a Cobra application.`,
    Run: func(cmd *cobra.Command, args []string) {
        fmt.Println("hello called")
    },
}

func init() {
    rootCmd.AddCommand(helloCmd)

    // Here you will define your flags and configuration settings.

    // Cobra supports Persistent Flags which will work for this command
    // and all subcommands, e.g.:
    // helloCmd.PersistentFlags().String("foo", "", "A help for foo")

    // Cobra supports local flags which will only run when this command
    // is called directly, e.g.:
    // helloCmd.Flags().BoolP("toggle", "t", false, "Help message for toggle")
}

go installでビルドすると使えるようになります。

$ go install
$ hoge hello
hello called

「hello called」と出力されました。簡単ですね。

作ってみる

まず作るコマンドラインツールの簡単な仕様

サブコマンドはpushのみ
wao push で現在のブランチをpush、レビュアーの選定し、マージリクエストの作成を行う
- 現在のブランチをpushする
  ※ git pushの代わりにwao pushを使う形になる
- slackのbotで自分以外のレビュアーが選ばれるまで繰り返す
  ※ botが返す名前はgitlabのユーザー名
- レビュアーをassigneeにセットしてマージリクエストを作成
  ※ マージリクエストのタイトルはデフォルトでマージ元のブランチ名とする
- 最後にマージリクエストのURLを出力して終了

設定ファイルはこんな感じになりました。 

# File: $HOME/.wao.toml
[Slack]
Token = "slack token"
ChannelID = "channel id"
TriggerWord = "reviewer"

[Gitlab]
Token = "gitlab token"
BaseURL = "https://gitlab.adways.net/api/v4"
UserName = "takagi_yoshiki"
DefaultMergeBranch = "develop"

弄るのは基本cmd/root.rbcmd/push.rbの2つのファイルだけになります。
本当はいい感じにpackage化したかったんですけど間に合わず。。。

// File: cmd/root.go
package cmd

import (
    "fmt"
    "os"

    homedir "github.com/mitchellh/go-homedir"
    "github.com/spf13/cobra"
    "github.com/spf13/viper"
)

// 読み込む設定ファイル名
var cfgFile string

// 読み込む設定の型
type Config struct {
    Slack  SlackConfig
    Gitlab GitlabConfig
}

type SlackConfig struct {
    Token       string
    ChannelID   string
    TriggerWord string
}

type GitlabConfig struct {
    Token              string
    BaseURL            string
    UserName           string
    DefaultMergeBranch string
}

// 読み込んだ設定ファイルの構造体
var config Config

var rootCmd = &cobra.Command{
    Use:   "wao",
    Short: "Wonderful Automatic Operation",
    Long: "Wonderful Automatic Operation",
}

func Execute() {
    if err := rootCmd.Execute(); err != nil {
        fmt.Println(err)
        os.Exit(1)
    }
}

func init() {
    cobra.OnInitialize(initConfig)

    // 設定ファイル名をフラグで受け取る
    rootCmd.PersistentFlags().StringVar(&cfgFile, "config", "", "config file (default is $HOME/.wao.toml)")
}

func initConfig() {
    if cfgFile != "" {
        viper.SetConfigFile(cfgFile)
    } else {
        home, err := homedir.Dir()
        if err != nil {
            fmt.Println(err)
            os.Exit(1)
        }

        viper.AddConfigPath(home)
        viper.SetConfigName(".wao")
    }

    // 設定ファイルを読み込む
    if err := viper.ReadInConfig(); err != nil {
        fmt.Println(err)
        os.Exit(1)
    }

    // 設定ファイルの内容を構造体にコピーする
    if err := viper.Unmarshal(&config); err != nil {
        fmt.Println(err)
        os.Exit(1)
    }
}
// File: cmd/push.go
package cmd

import (
    "fmt"
    "os/exec"
    "strings"

    "github.com/nlopes/slack"
    "github.com/spf13/cobra"
    "github.com/xanzy/go-gitlab"
)

var pushCmd = &cobra.Command{
    Use:   "push",
    Short: "Select reviewer and Create merge request",
    Long: "Select reviewer and Create merge request",
    RunE: push,
}

func push(cmd *cobra.Command, args []string) error {
    userName, err := selectReviewer()
    if err != nil {
        return err
    }

    userID, err := fetchUserID(userName)
    if err != nil {
        return err
    }

    // 現在のブランチ名を取得
    out, err := exec.Command("git", "symbolic-ref", "--short", "HEAD").Output()
    if err != nil {
        return err
    }

    // stringにキャストして、末尾の改行を取り除く
    sourceBranch := strings.TrimRight(string(out), "\n")

    err = exec.Command("git", "push", "origin", sourceBranch).Run()
    if err != nil {
        return err
    }

    projectID, err := fetchProjectID()
    if err != nil {
        return err
    }

    mergeRequestURL, err := fetchMergeRequestURL(sourceBranch, userID, projectID)
    if err != nil {
        return err
    }

    fmt.Println(mergeRequestURL)

    return nil
}

// レビュアーを選定
func selectReviewer() (string, error) {
    botUserID := "USLACKBOT"
    channelID := config.Slack.ChannelID
    triggerWord := config.Slack.TriggerWord
    userName := config.Gitlab.UserName

    api := slack.New(config.Slack.Token)
    rtm := api.NewRTM()
    go rtm.ManageConnection()

    for msg := range rtm.IncomingEvents {
        switch ev := msg.Data.(type) {
        case *slack.HelloEvent:
            // 始めにbotを呼び出すトリガーを発言する
            rtm.SendMessage(rtm.NewOutgoingMessage(triggerWord, channelID))

        case *slack.MessageEvent:
            // 指定のチャンネルではない、かつボットの発言じゃなかった場合に次に進む
            if ev.Channel != channelID || ev.User != botUserID {
                continue
            }

            // 自分以外のレビュアーが選ばれたら終了
            if ev.Text != userName {
                return ev.Text, nil
            }

            // 自分が選ばれてしまったら再度トリガーを発言する
            rtm.SendMessage(rtm.NewOutgoingMessage(triggerWord, channelID))

        case *slack.RTMError:
            return "", fmt.Errorf(ev.Error())

        case *slack.InvalidAuthEvent:
            return "", fmt.Errorf("Invalid credentials")
        }
    }

    return "", fmt.Errorf("Failed Select Reviewer")
}

// ユーザーIDを取得
func fetchUserID(userName string) (int, error) {
    git := gitlab.NewClient(nil, config.Gitlab.Token)
    git.SetBaseURL(config.Gitlab.BaseURL)

    opt := &gitlab.ListUsersOptions{
        Username: gitlab.String(userName),
    }

    users, _, err := git.Users.ListUsers(opt)

    if err != nil {
        return 0, err
    }

    if len(users) == 0 {
        return 0, fmt.Errorf("User Not Found")
    }

    user := users[0]

    return user.ID, nil
}

// プロジェクトIDを取得
func fetchProjectID() (int, error) {
    out, err := exec.Command("git", "config", "--get", "remote.origin.url").Output()

    if err != nil {
        return 0, err
    }

    remoteOriginURL := strings.TrimRight(string(out), "\n")

    git := gitlab.NewClient(nil, config.Gitlab.Token)
    git.SetBaseURL(config.Gitlab.BaseURL)

    opt := &gitlab.ListProjectsOptions{
        Membership: gitlab.Bool(true),
        ListOptions: gitlab.ListOptions{
            PerPage: 100,
        },
    }

    projects, resp, err := git.Projects.ListProjects(opt)

    for _, project := range projects {
        if remoteOriginURL == project.SSHURLToRepo {
            return project.ID, nil
        }
    }

    for page := 2; page < resp.TotalPages; page++ {
        opt.ListOptions.Page = page
        projects, _, err := git.Projects.ListProjects(opt)

        if err != nil {
            return 0, err
        }

        for _, project := range projects {
            if remoteOriginURL == project.SSHURLToRepo {
                return project.ID, nil
            }
        }
    }

    return 0, fmt.Errorf("Project Not Found")
}

// マージリクエストのURLを取得
func fetchMergeRequestURL(sourceBranch string, userID int, projectID int) (string, error) {
    git := gitlab.NewClient(nil, config.Gitlab.Token)
    git.SetBaseURL(config.Gitlab.BaseURL)

    opt := &gitlab.CreateMergeRequestOptions{
        Title:        gitlab.String(sourceBranch),
        SourceBranch: gitlab.String(sourceBranch),
        TargetBranch: gitlab.String(config.Gitlab.DefaultMergeBranch),
        AssigneeID:   gitlab.Int(userID),
    }

    mergeRequest, _, err := git.MergeRequests.CreateMergeRequest(projectID, opt)

    if err != nil {
        return "", err
    }

    return mergeRequest.WebURL, nil
}

func init() {
    rootCmd.AddCommand(pushCmd)
}

動かしてみる

f:id:AdwaysEngineerBlog:20181225151811g:plain

どや!

f:id:AdwaysEngineerBlog:20181225151322p:plain

ちゃんとassigneeにも入っていることが確認できました。

テスト

うっ

今後やりたいこと

  • 一つのファイルの中でいろいろ関数を書いているので適切なファイル分けをしたい
  • flagsを使ってオプションでマージ先ブランチなどを自由に指定できるようにしたい
  • 設定ファイルでslackのチャンネルをチャンネルIDではなくチャンネル名で指定できるようにしたい

などなど

おわりに

自分なりにお作法を調べながら書きましたが、おかしな書き方をしているところがあったら指摘していただけると嬉しいです。
調べていると公式ドキュメントやチュートリアルにたどり着くので、結局のところがっつり公式ドキュメントやチュートリアルを読み進めた方が早かったかもしれません。
cobraを使えばコマンドラインツールを簡単に作れることが分かったので、これからしょうもないコマンドラインツールを量産していくつもりです。
現在業務ではRubyがメインの言語なのですが、より一層理解を深めてGoを業務に取り入れられたらいいなと思います。

今日で2018年のアドベントカレンダーも終わりですね。
それではよいクリスマスにしましょう!メリークリスマス!

参考