use context

Author: mashiike

cmd/mysqlbatch/main.go

@@ -1,10 +1,13 @@
 package main
 
 import (
+	"context"
 	"flag"
 	"fmt"
 	"log"
 	"os"
+	"os/signal"
+	"syscall"
 
 	"github.com/mashiike/mysqlbatch"
 )
@@ -37,9 +40,11 @@ func main() {
 	if flag.NArg() == 1 {
 		conf.Database = flag.Arg(0)
 	}
+	ctx, stop := signal.NotifyContext(context.Background(), syscall.SIGTERM|syscall.SIGHUP|syscall.SIGINT)
+	defer stop()
 
 	executer := mysqlbatch.New(conf)
-	if err := executer.Execute(os.Stdin); err != nil {
+	if err := executer.ExecuteContext(ctx, os.Stdin); err != nil {
 		log.Println(err)
 		os.Exit(1)
 	}

executer.go

@@ -2,6 +2,7 @@ package mysqlbatch
 
 import (
 	"bufio"
+	"context"
 	"database/sql"
 	"fmt"
 	"io"
@@ -53,6 +54,10 @@ func New(config *Config) *Executer {
 }
 
 func (e *Executer) Execute(queryReader io.Reader) error {
+	return e.ExecuteContext(context.Background(), queryReader)
+}
+
+func (e *Executer) ExecuteContext(ctx context.Context, queryReader io.Reader) error {
 
 	db, err := sql.Open("mysql", e.dsn)
 	if err != nil {
@@ -67,7 +72,7 @@ func (e *Executer) Execute(queryReader io.Reader) error {
 		if query == "" {
 			continue
 		}
-		if _, err := db.Exec(query); err != nil {
+		if _, err := db.ExecContext(ctx, query); err != nil {
 			return errors.Wrap(err, "execute query failed")
 		}
 	}