Repository: ariga/atlas
Pull request #632 (2022-03-08 07:54:47)

cmd/action: use cobra context to cancel operations on interrupt / kill


Patch:
diff --git a/cmd/action/apply.go b/cmd/action/apply.go
index 89af1fa3..9e9b8fb4 100644
--- a/cmd/action/apply.go
+++ b/cmd/action/apply.go
@@ -71,20 +71,19 @@ func init() {
 }
 
 // CmdApplyRun is the command used when running CLI.
-func CmdApplyRun(*cobra.Command, []string) {
+func CmdApplyRun(cmd *cobra.Command, _ []string) {
 	if ApplyFlags.Web {
 		schemaCmd.PrintErrln("The Atlas UI is not available in this release.")
 		return
 	}
 	d, err := DefaultMux.OpenAtlas(ApplyFlags.URL)
 	cobra.CheckErr(err)
-	applyRun(d, ApplyFlags.URL, ApplyFlags.File, ApplyFlags.DryRun, ApplyFlags.AutoApprove)
+	applyRun(cmd.Context(), d, ApplyFlags.URL, ApplyFlags.File, ApplyFlags.DryRun, ApplyFlags.AutoApprove)
 }
 
-func applyRun(d *Driver, url string, file string, dryRun bool, autoApprove bool) {
-	ctx := context.Background()
+func applyRun(ctx context.Context, d *Driver, url string, file string, dryRun bool, autoApprove bool) {
 	schemas := ApplyFlags.Schema
-	if n, err := SchemaNameFromURL(url); n != "" {
+	if n, err := SchemaNameFromURL(ctx, url); n != "" {
 		cobra.CheckErr(err)
 		schemas = append(schemas, n)
 	}
diff --git a/cmd/action/diff.go b/cmd/action/diff.go
index 1eeae816..e25bdf29 100644
--- a/cmd/action/diff.go
+++ b/cmd/action/diff.go
@@ -5,7 +5,6 @@
 package action
 
 import (
-	"context"
 	"strings"
 
 	"github.com/spf13/cobra"
@@ -44,15 +43,15 @@ func init() {
 // cmdDiffRun connects to the given databases, and prints an SQL plan to get from
 // the "from" schema to the "to" schema.
 func cmdDiffRun(cmd *cobra.Command, flags *diffCmdOpts) {
+	ctx := cmd.Context()
 	fromDriver, err := DefaultMux.OpenAtlas(flags.fromURL)
 	cobra.CheckErr(err)
 	toDriver, err := DefaultMux.OpenAtlas(flags.toURL)
 	cobra.CheckErr(err)
-	fromName, err := SchemaNameFromURL(flags.fromURL)
+	fromName, err := SchemaNameFromURL(ctx, flags.fromURL)
 	cobra.CheckErr(err)
-	toName, err := SchemaNameFromURL(flags.toURL)
+	toName, err := SchemaNameFromURL(ctx, flags.toURL)
 	cobra.CheckErr(err)
-	ctx := context.Background()
 	fromSchema, err := fromDriver.InspectSchema(ctx, fromName, nil)
 	cobra.CheckErr(err)
 	toSchema, err := toDriver.InspectSchema(ctx, toName, nil)
diff --git a/cmd/action/inspect.go b/cmd/action/inspect.go
index 1e1917b6..2717f533 100644
--- a/cmd/action/inspect.go
+++ b/cmd/action/inspect.go
@@ -56,20 +56,19 @@ func init() {
 }
 
 // CmdInspectRun is the command used when running CLI.
-func CmdInspectRun(_ *cobra.Command, _ []string) {
+func CmdInspectRun(cmd *cobra.Command, _ []string) {
 	if InspectFlags.Web {
 		schemaCmd.PrintErrln("The Atlas UI is not available in this release.")
 		return
 	}
 	d, err := DefaultMux.OpenAtlas(InspectFlags.URL)
 	cobra.CheckErr(err)
-	inspectRun(d, InspectFlags.URL)
+	inspectRun(cmd.Context(), d, InspectFlags.URL)
 }
 
-func inspectRun(d *Driver, url string) {
-	ctx := context.Background()
+func inspectRun(ctx context.Context, d *Driver, url string) {
 	schemas := InspectFlags.Schema
-	if n, err := SchemaNameFromURL(url); n != "" {
+	if n, err := SchemaNameFromURL(ctx, url); n != "" {
 		cobra.CheckErr(err)
 		schemas = append(schemas, n)
 	}
diff --git a/cmd/action/mux.go b/cmd/action/mux.go
index bd8406b8..15ad8ff5 100644
--- a/cmd/action/mux.go
+++ b/cmd/action/mux.go
@@ -83,7 +83,7 @@ func urlParts(url string) (string, string, error) {
 }
 
 // SchemaNameFromURL parses the url the returns schema name
-func SchemaNameFromURL(url string) (string, error) {
+func SchemaNameFromURL(ctx context.Context, url string) (string, error) {
 	key, dsn, err := urlParts(url)
 	if err != nil {
 		return "", err
@@ -98,7 +98,7 @@ func SchemaNameFromURL(url string) (string, error) {
 	case "postgres":
 		return postgresSchema(dsn)
 	case "sqlite":
-		return schemaName(dsn)
+		return schemaName(ctx, dsn)
 	default:
 		return "", fmt.Errorf("unknown database type: %q", key)
 	}
@@ -119,7 +119,7 @@ func postgresSchema(dsn string) (string, error) {
 	return "", nil
 }
 
-func schemaName(dsn string) (string, error) {
+func schemaName(ctx context.Context, dsn string) (string, error) {
 	err := sqliteFileExists(dsn)
 	if err != nil {
 		return "", err
@@ -132,7 +132,7 @@ func schemaName(dsn string) (string, error) {
 	if err != nil {
 		return "", err
 	}
-	r, err := drv.InspectRealm(context.Background(), nil)
+	r, err := drv.InspectRealm(ctx, nil)
 	if err != nil {
 		return "", err
 	}
diff --git a/cmd/atlas/main.go b/cmd/atlas/main.go
index cdb277ce..bc0a0580 100644
--- a/cmd/atlas/main.go
+++ b/cmd/atlas/main.go
@@ -1,7 +1,9 @@
 package main
 
 import (
+	"context"
 	"os"
+	"os/signal"
 
 	"ariga.io/atlas/cmd/action"
 
@@ -11,8 +13,9 @@ import (
 )
 
 func main() {
+	ctx, _ := signal.NotifyContext(context.Background(), os.Interrupt, os.Kill)
 	action.RootCmd.SetOut(os.Stdout)
-	err := action.RootCmd.Execute()
+	err := action.RootCmd.ExecuteContext(ctx)
 	// Print error from command
 	if err != nil {
 		action.RootCmd.PrintErrln("Error:", err)
@@ -23,5 +26,4 @@ func main() {
 	if err != nil {
 		os.Exit(1)
 	}
-
 }


Test patch:
diff --git a/cmd/action/mux_test.go b/cmd/action/mux_test.go
index 56ac5e82..0545f846 100644
--- a/cmd/action/mux_test.go
+++ b/cmd/action/mux_test.go
@@ -5,6 +5,7 @@
 package action_test
 
 import (
+	"context"
 	"fmt"
 	"io/ioutil"
 	"log"
@@ -100,7 +101,7 @@ func Test_SQLiteFileDoestNotExist(t *testing.T) {
 	}
 	for _, tt := range tests {
 		t.Run(tt.url, func(t *testing.T) {
-			_, err := action.SchemaNameFromURL(tt.url)
+			_, err := action.SchemaNameFromURL(context.Background(), tt.url)
 			require.EqualError(t, err, tt.expected)
 		})
 	}
@@ -115,13 +116,13 @@ func Test_SQLiteFileExist(t *testing.T) {
 		r.NoError(err)
 	})
 	dsn := "sqlite://file://" + file.Name()
-	_, err = action.SchemaNameFromURL(dsn)
+	_, err = action.SchemaNameFromURL(context.Background(), dsn)
 	r.NoError(err)
 }
 
 func Test_SQLiteInMemory(t *testing.T) {
 	r := require.New(t)
-	_, err := action.SchemaNameFromURL("sqlite://file:test.db?cache=shared&mode=memory")
+	_, err := action.SchemaNameFromURL(context.Background(), "sqlite://file:test.db?cache=shared&mode=memory")
 	r.NoError(err)
 }
 
@@ -151,7 +152,7 @@ func Test_PostgresSchemaDSN(t *testing.T) {
 	}
 	for _, tt := range tests {
 		t.Run(tt.url, func(t *testing.T) {
-			schema, err := action.SchemaNameFromURL(tt.url)
+			schema, err := action.SchemaNameFromURL(context.Background(), tt.url)
 			require.Equal(t, tt.wantErr, err != nil)
 			require.Equal(t, tt.expected, schema)
 		})