Do not cancel context for async queries on success

Signed-off-by: Carlos <[email protected]>

engine.go

@@ -153,6 +153,18 @@ func (e *Engine) Query(
 	return analyzed.Schema(), iter, nil
 }
 
+// Async returns true if the query is async. If there are any errors with the
+// query it returns false
+func (e *Engine) Async(ctx *sql.Context, query string) bool {
+	parsed, err := parse.Parse(ctx, query)
+	if err != nil {
+		return false
+	}
+
+	asyncNode, ok := parsed.(sql.AsyncNode)
+	return ok && asyncNode.IsAsync()
+}
+
 // AddDatabase adds the given database to the catalog.
 func (e *Engine) AddDatabase(db sql.Database) {
 	e.Catalog.AddDatabase(db)

server/handler.go

@@ -113,9 +113,13 @@ func (h *Handler) ComQuery(
 	callback func(*sqltypes.Result) error,
 ) (err error) {
 	ctx := h.sm.NewContextWithQuery(c, query)
-	newCtx, cancel := context.WithCancel(ctx)
-	defer cancel()
-	ctx = ctx.WithContext(newCtx)
+
+	if !h.e.Async(ctx, query) {
+		newCtx, cancel := context.WithCancel(ctx)
+		ctx = ctx.WithContext(newCtx)
+
+		defer cancel()
+	}
 
 	handled, err := h.handleKill(c, query)
 	if err != nil {

sql/core.go

@@ -120,6 +120,12 @@ type OpaqueNode interface {
 	Opaque() bool
 }
 
+// AsyncNode is a node that can be executed asynchronously.
+type AsyncNode interface {
+	// IsAsync reports whether the node is async or not.
+	IsAsync() bool
+}
+
 // Expressioner is a node that contains expressions.
 type Expressioner interface {
 	// Expressions returns the list of expressions contained by the node.

sql/plan/create_index.go

@@ -282,6 +282,11 @@ func (c *CreateIndex) WithChildren(children ...sql.Node) (sql.Node, error) {
 	return &nc, nil
 }
 
+// IsAsync implements the AsyncNode interface.
+func (c *CreateIndex) IsAsync() bool {
+	return c.Async
+}
+
 // getColumnsAndPrepareExpressions extracts the unique columns required by all
 // those expressions and fixes the indexes of the GetFields in the expressions
 // to match a row with only the returned columns in that same order.