Repository: src-d/go-mysql-server (go)

Issue:
Fix behaviour of query cancellation
Continuation of #369 

- Drop connection does not kill any query
- Kill query kills the query but does not drop the connection
- Kill drops the connection after killing all executing queries

This affects index creation and so on.

Pull request description:
*: correctly handle process management using processlist
Closes #371 

Summary of changes:

- Session manager, no longer tracks contexts created for queries, only sessions. Contexts are managed by the ProcessList as part of the Process.
- Pid is available from Context now, and connection id from the Session.
- trackProcess rule does no longer register the process, only attaches step information and wraps the node if needed to track completion.
- KILL correctly handled: query just kills query but keeps connection, connection drops connection and all associated queries.
- KILL connection is also tested now (using unsafe, though :( )
- DROP INDEX cannot drop indexes that are currently being created.

Patch:
diff --git a/engine.go b/engine.go
index 7a76da7..c73ba3e 100644
--- a/engine.go
+++ b/engine.go
@@ -7,6 +7,7 @@ import (
 	"gopkg.in/src-d/go-mysql-server.v0/sql/analyzer"
 	"gopkg.in/src-d/go-mysql-server.v0/sql/expression/function"
 	"gopkg.in/src-d/go-mysql-server.v0/sql/parse"
+	"gopkg.in/src-d/go-mysql-server.v0/sql/plan"
 )
 
 // Config for the Engine.
@@ -43,13 +44,12 @@ func NewDefault() *Engine {
 	return New(c, a, nil)
 }
 
-// Query executes a query without attaching to any context.
+// Query executes a query.
 func (e *Engine) Query(
 	ctx *sql.Context,
 	query string,
 ) (sql.Schema, sql.RowIter, error) {
-	span, ctx := ctx.WithQuery(query).
-		Span("query", opentracing.Tag{Key: "query", Value: query})
+	span, ctx := ctx.Span("query", opentracing.Tag{Key: "query", Value: query})
 	defer span.Finish()
 
 	logrus.WithField("query", query).Debug("executing query")
@@ -59,6 +59,22 @@ func (e *Engine) Query(
 		return nil, nil, err
 	}
 
+	var typ = sql.QueryProcess
+	if _, ok := parsed.(*plan.CreateIndex); ok {
+		typ = sql.CreateIndexProcess
+	}
+
+	ctx, err = e.Catalog.AddProcess(ctx, typ, query)
+	defer func() {
+		if err != nil && ctx != nil {
+			e.Catalog.Done(ctx.Pid())
+		}
+	}()
+
+	if err != nil {
+		return nil, nil, err
+	}
+
 	analyzed, err := e.Analyzer.Analyze(ctx, parsed)
 	if err != nil {
 		return nil, nil, err
diff --git a/server/context.go b/server/context.go
index d5ad191..04794dc 100644
--- a/server/context.go
+++ b/server/context.go
@@ -5,7 +5,6 @@ import (
 	"sync"
 
 	opentracing "github.com/opentracing/opentracing-go"
-	uuid "github.com/satori/go.uuid"
 	"gopkg.in/src-d/go-mysql-server.v0/sql"
 	"gopkg.in/src-d/go-vitess.v0/mysql"
 )
@@ -19,20 +18,19 @@ type DoneFunc func()
 
 // DefaultSessionBuilder is a SessionBuilder that returns a base session.
 func DefaultSessionBuilder(c *mysql.Conn, addr string) sql.Session {
-	return sql.NewSession(addr, c.User)
+	return sql.NewSession(addr, c.User, c.ConnectionID)
 }
 
 // SessionManager is in charge of creating new sessions for the given
 // connections and keep track of which sessions are in each connection, so
 // they can be cancelled is the connection is closed.
 type SessionManager struct {
-	addr            string
-	tracer          opentracing.Tracer
-	mu              *sync.Mutex
-	builder         SessionBuilder
-	sessions        map[uint32]sql.Session
-	sessionContexts map[uint32][]uuid.UUID
-	contexts        map[uuid.UUID]context.CancelFunc
+	addr     string
+	tracer   opentracing.Tracer
+	mu       *sync.Mutex
+	builder  SessionBuilder
+	sessions map[uint32]sql.Session
+	pid      uint64
 }
 
 // NewSessionManager creates a SessionManager with the given SessionBuilder.
@@ -42,16 +40,21 @@ func NewSessionManager(
 	addr string,
 ) *SessionManager {
 	return &SessionManager{
-		addr:            addr,
-		tracer:          tracer,
-		mu:              new(sync.Mutex),
-		builder:         builder,
-		sessions:        make(map[uint32]sql.Session),
-		sessionContexts: make(map[uint32][]uuid.UUID),
-		contexts:        make(map[uuid.UUID]context.CancelFunc),
+		addr:     addr,
+		tracer:   tracer,
+		mu:       new(sync.Mutex),
+		builder:  builder,
+		sessions: make(map[uint32]sql.Session),
 	}
 }
 
+func (s *SessionManager) nextPid() uint64 {
+	s.mu.Lock()
+	defer s.mu.Unlock()
+	s.pid++
+	return s.pid
+}
+
 // NewSession creates a Session for the given connection.
 func (s *SessionManager) NewSession(conn *mysql.Conn) {
 	s.mu.Lock()
@@ -60,32 +63,18 @@ func (s *SessionManager) NewSession(conn *mysql.Conn) {
 }
 
 // NewContext creates a new context for the session at the given conn.
-func (s *SessionManager) NewContext(conn *mysql.Conn) (*sql.Context, DoneFunc, error) {
-	ctx, cancel := context.WithCancel(context.Background())
+func (s *SessionManager) NewContext(conn *mysql.Conn) *sql.Context {
 	s.mu.Lock()
 	sess := s.sessions[conn.ConnectionID]
 	s.mu.Unlock()
-	context := sql.NewContext(ctx, sql.WithSession(sess), sql.WithTracer(s.tracer))
-	id := uuid.NewV4()
-
-	s.mu.Lock()
-	s.sessionContexts[conn.ConnectionID] = append(s.sessionContexts[conn.ConnectionID], id)
-	s.contexts[id] = cancel
-	s.mu.Unlock()
-
-	return context, func() {
-		s.mu.Lock()
-		defer s.mu.Unlock()
+	context := sql.NewContext(
+		context.Background(),
+		sql.WithSession(sess),
+		sql.WithTracer(s.tracer),
+		sql.WithPid(s.nextPid()),
+	)
 
-		delete(s.contexts, id)
-		ids := s.sessionContexts[conn.ConnectionID]
-		for i, sessID := range ids {
-			if sessID == id {
-				s.sessionContexts[conn.ConnectionID] = append(ids[:i], ids[i+1:]...)
-				break
-			}
-		}
-	}, nil
+	return context
 }
 
 // CloseConn closes the connection in the session manager and all its
@@ -93,10 +82,5 @@ func (s *SessionManager) NewContext(conn *mysql.Conn) (*sql.Context, DoneFunc, e
 func (s *SessionManager) CloseConn(conn *mysql.Conn) {
 	s.mu.Lock()
 	defer s.mu.Unlock()
-
-	for _, id := range s.sessionContexts[conn.ConnectionID] {
-		s.contexts[id]()
-		delete(s.contexts, id)
-	}
-	delete(s.sessionContexts, conn.ConnectionID)
+	delete(s.sessions, conn.ConnectionID)
 }
diff --git a/server/handler.go b/server/handler.go
index c6ef0c1..16d2992 100644
--- a/server/handler.go
+++ b/server/handler.go
@@ -70,14 +70,9 @@ func (h *Handler) ComQuery(
 	query string,
 	callback func(*sqltypes.Result) error,
 ) error {
-	ctx, done, err := h.sm.NewContext(c)
-	if err != nil {
-		return err
-	}
+	ctx := h.sm.NewContext(c)
 
-	defer done()
-
-	handled, err := h.handleKill(query)
+	handled, err := h.handleKill(c, query)
 	if err != nil {
 		return err
 	}
@@ -122,6 +117,10 @@ func (h *Handler) ComQuery(
 		r.RowsAffected++
 	}
 
+	if err := rows.Close(); err != nil {
+		return err
+	}
+
 	// Even if r.RowsAffected = 0, the callback must be
 	// called to update the state in the go-vitess' listener
 	// and avoid returning errors when the query doesn't
@@ -133,40 +132,40 @@ func (h *Handler) ComQuery(
 	return callback(r)
 }
 
-func (h *Handler) handleKill(query string) (bool, error) {
+func (h *Handler) handleKill(conn *mysql.Conn, query string) (bool, error) {
 	q := strings.ToLower(query)
 	s := regKillCmd.FindStringSubmatch(q)
 	if s == nil {
 		return false, nil
 	}
 
-	id, err := strconv.Atoi(s[2])
+	id, err := strconv.ParseUint(s[2], 10, 64)
 	if err != nil {
 		return false, err
 	}
 
-	logrus.Infof("handleKill: id %v", id)
-
-	h.mu.Lock()
-	c, ok := h.c[uint32(id)]
-	h.mu.Unlock()
-	if !ok {
-		return false, errConnectionNotFound.New(id)
-	}
-
-	h.sm.CloseConn(c)
-
 	// KILL CONNECTION and KILL should close the connection. KILL QUERY only
 	// cancels the query.
 	//
 	// https://dev.mysql.com/doc/refman/5.7/en/kill.html
 
-	if s[1] != "query" {
-		c.Close()
-
+	if s[1] == "query" {
+		logrus.Infof("kill query: id %v", id)
+		h.e.Catalog.Kill(id)
+	} else {
+		logrus.Infof("kill connection: id %v, pid: %v", conn.ConnectionID, id)
 		h.mu.Lock()
-		delete(h.c, uint32(id))
+		c, ok := h.c[conn.ConnectionID]
+		delete(h.c, conn.ConnectionID)
 		h.mu.Unlock()
+
+		if !ok {
+			return false, errConnectionNotFound.New(conn.ConnectionID)
+		}
+
+		h.e.Catalog.KillConnection(id)
+		h.sm.CloseConn(c)
+		c.Close()
 	}
 
 	return true, nil
diff --git a/sql/analyzer/process.go b/sql/analyzer/process.go
index e85f317..678c87b 100644
--- a/sql/analyzer/process.go
+++ b/sql/analyzer/process.go
@@ -5,6 +5,8 @@ import (
 	"gopkg.in/src-d/go-mysql-server.v0/sql/plan"
 )
 
+// trackProcess will wrap the query in a process node and add progress items
+// to the already existing process.
 func trackProcess(ctx *sql.Context, a *Analyzer, n sql.Node) (sql.Node, error) {
 	if !n.Resolved() {
 		return n, nil
@@ -14,18 +16,7 @@ func trackProcess(ctx *sql.Context, a *Analyzer, n sql.Node) (sql.Node, error) {
 		return n, nil
 	}
 
-	var query string
-	if x, ok := ctx.Value(sql.QueryKey).(string); ok {
-		query = x
-	}
-
-	var typ = sql.QueryProcess
-	if _, ok := n.(*plan.CreateIndex); ok {
-		typ = sql.CreateIndexProcess
-	}
-
 	processList := a.Catalog.ProcessList
-	pid := processList.AddProcess(ctx, typ, query)
 
 	var seen = make(map[string]struct{})
 	n, err := n.TransformUp(func(n sql.Node) (sql.Node, error) {
@@ -48,11 +39,11 @@ func trackProcess(ctx *sql.Context, a *Analyzer, n sql.Node) (sql.Node, error) {
 				}
 				total = count
 			}
-			processList.AddProgressItem(pid, name, total)
+			processList.AddProgressItem(ctx.Pid(), name, total)
 
 			seen[name] = struct{}{}
 			t := plan.NewProcessTable(n.Table, func() {
-				processList.UpdateProgress(pid, name, 1)
+				processList.UpdateProgress(ctx.Pid(), name, 1)
 			})
 
 			return plan.NewResolvedTable(t), nil
@@ -64,5 +55,11 @@ func trackProcess(ctx *sql.Context, a *Analyzer, n sql.Node) (sql.Node, error) {
 		return nil, err
 	}
 
-	return plan.NewQueryProcess(n, func() { processList.Done(pid) }), nil
+	// Don't wrap CreateIndex in a QueryProcess, as it is a CreateIndexProcess.
+	// CreateIndex will take care of marking the process as done on its own.
+	if _, ok := n.(*plan.CreateIndex); ok {
+		return n, nil
+	}
+
+	return plan.NewQueryProcess(n, func() { processList.Done(ctx.Pid()) }), nil
 }
diff --git a/sql/plan/create_index.go b/sql/plan/create_index.go
index a50b4ce..2cf09de 100644
--- a/sql/plan/create_index.go
+++ b/sql/plan/create_index.go
@@ -154,6 +154,7 @@ func (c *CreateIndex) RowIter(ctx *sql.Context) (sql.RowIter, error) {
 
 	createIndex := func() {
 		c.createIndex(ctx, log, driver, index, iter, created, ready)
+		c.Catalog.ProcessList.Done(ctx.Pid())
 	}
 
 	log.WithField("async", c.Async).Info("starting to save the index")
diff --git a/sql/plan/drop_index.go b/sql/plan/drop_index.go
index 77258c4..8b6066a 100644
--- a/sql/plan/drop_index.go
+++ b/sql/plan/drop_index.go
@@ -12,6 +12,9 @@ var (
 	ErrTableNotValid = errors.NewKind("table is not valid")
 	// ErrTableNotNameable is returned when the table is not nameable.
 	ErrTableNotNameable = errors.NewKind("can't get name from table")
+	// ErrIndexNotAvailable is returned when trying to delete an index that is
+	// still not ready for usage.
+	ErrIndexNotAvailable = errors.NewKind("index %q is still not ready for usage and can't be deleted")
 )
 
 // DropIndex is a node to drop an index.
@@ -59,6 +62,10 @@ func (d *DropIndex) RowIter(ctx *sql.Context) (sql.RowIter, error) {
 	}
 	d.Catalog.ReleaseIndex(index)
 
+	if !d.Catalog.CanUseIndex(index) {
+		return nil, ErrIndexNotAvailable.New(d.Name)
+	}
+
 	done, err := d.Catalog.DeleteIndex(db.Name(), d.Name, true)
 	if err != nil {
 		return nil, err
diff --git a/sql/processlist.go b/sql/processlist.go
index 57eec66..fcc1100 100644
--- a/sql/processlist.go
+++ b/sql/processlist.go
@@ -1,9 +1,12 @@
 package sql
 
 import (
+	"context"
 	"fmt"
 	"sync"
 	"time"
+
+	"gopkg.in/src-d/go-errors.v1"
 )
 
 // Progress between done items and total items.
@@ -44,14 +47,19 @@ func (p ProcessType) String() string {
 
 // Process represents a process in the SQL server.
 type Process struct {
-	Pid       uint64
-	User      string
-	Type      ProcessType
-	Query     string
-	Progress  map[string]Progress
-	StartedAt time.Time
+	Pid        uint64
+	Connection uint32
+	User       string
+	Type       ProcessType
+	Query      string
+	Progress   map[string]Progress
+	StartedAt  time.Time
+	Kill       context.CancelFunc
 }
 
+// Done needs to be called when this process has finished.
+func (p *Process) Done() { p.Kill() }
+
 // Seconds returns the number of seconds this process has been running.
 func (p *Process) Seconds() uint64 {
 	return uint64(time.Since(p.StartedAt) / time.Second)
@@ -61,7 +69,6 @@ func (p *Process) Seconds() uint64 {
 // status.
 type ProcessList struct {
 	mu    sync.RWMutex
-	pid   uint64
 	procs map[uint64]*Process
 }
 
@@ -72,28 +79,41 @@ func NewProcessList() *ProcessList {
 	}
 }
 
+// ErrPidAlreadyUsed is returned when the pid is already registered.
+var ErrPidAlreadyUsed = errors.NewKind("pid %d is already in use")
+
 // AddProcess adds a new process to the list given a process type and a query.
 // Steps is a map between the name of the items that need to be completed and
 // the total amount in these items. -1 means unknown.
+// It returns a new context that should be passed around from now on. That
+// context will be cancelled if the process is killed.
 func (pl *ProcessList) AddProcess(
 	ctx *Context,
 	typ ProcessType,
 	query string,
-) (pid uint64) {
+) (*Context, error) {
 	pl.mu.Lock()
 	defer pl.mu.Unlock()
 
-	pl.pid++
-	pl.procs[pl.pid] = &Process{
-		Pid:       pl.pid,
-		Type:      typ,
-		Query:     query,
-		Progress:  make(map[string]Progress),
-		User:      ctx.Session.User(),
-		StartedAt: time.Now(),
+	if _, ok := pl.procs[ctx.Pid()]; ok {
+		return nil, ErrPidAlreadyUsed.New(ctx.Pid())
 	}
 
-	return pl.pid
+	newCtx, cancel := context.WithCancel(ctx)
+	ctx = ctx.WithContext(newCtx)
+
+	pl.procs[ctx.Pid()] = &Process{
+		Pid:        ctx.Pid(),
+		Connection: ctx.ID(),
+		Type:       typ,
+		Query:      query,
+		Progress:   make(map[string]Progress),
+		User:       ctx.Session.User(),
+		StartedAt:  time.Now(),
+		Kill:       cancel,
+	}
+
+	return ctx, nil
 }
 
 // UpdateProgress updates the progress of the item with the given name for the
@@ -117,7 +137,7 @@ func (pl *ProcessList) UpdateProgress(pid uint64, name string, delta int64) {
 }
 
 // AddProgressItem adds a new item to track progress from to the proces with
-// the given pid.
+// the given pid. If the pid does not exist, it will do nothing.
 func (pl *ProcessList) AddProgressItem(pid uint64, name string, total int64) {
 	pl.mu.Lock()
 	defer pl.mu.Unlock()
@@ -135,10 +155,42 @@ func (pl *ProcessList) AddProgressItem(pid uint64, name string, total int64) {
 	}
 }
 
+// Kill terminates a process if it exists.
+func (pl *ProcessList) Kill(pid uint64) {
+	pl.Done(pid)
+}
+
+// KillConnection kills all processes that have the same connection as the one
+// of the process with the given process id. If the process does not exist, it
+// will do nothing.
+func (pl *ProcessList) KillConnection(pid uint64) {
+	pl.mu.Lock()
+	defer pl.mu.Unlock()
+
+	proc, ok := pl.procs[pid]
+	if !ok {
+		return
+	}
+
+	conn := proc.Connection
+	for pid, proc := range pl.procs {
+		if proc.Connection == conn {
+			proc.Kill()
+			delete(pl.procs, pid)
+		}
+	}
+}
+
 // Done removes the finished process with the given pid from the process list.
+// If the process does not exist, it will do nothing.
 func (pl *ProcessList) Done(pid uint64) {
 	pl.mu.Lock()
 	defer pl.mu.Unlock()
+
+	if proc, ok := pl.procs[pid]; ok {
+		proc.Done()
+	}
+
 	delete(pl.procs, pid)
 }
 
diff --git a/sql/session.go b/sql/session.go
index 8d6ebea..b7d6df8 100644
--- a/sql/session.go
+++ b/sql/session.go
@@ -28,10 +28,13 @@ type Session interface {
 	Set(key string, typ Type, value interface{})
 	// Get session configuration.
 	Get(key string) (Type, interface{})
+	// ID returns the unique ID of the connection.
+	ID() uint32
 }
 
 // BaseSession is the basic session type.
 type BaseSession struct {
+	id     uint32
 	addr   string
 	user   string
 	mu     sync.RWMutex
@@ -63,6 +66,9 @@ func (s *BaseSession) Get(key string) (Type, interface{}) {
 	return v.typ, v.value
 }
 
+// ID implements the Session interface.
+func (s *BaseSession) ID() uint32 { return s.id }
+
 type typedValue struct {
 	typ   Type
 	value interface{}
@@ -79,8 +85,9 @@ func defaultSessionConfig() map[string]typedValue {
 }
 
 // NewSession creates a new session with data.
-func NewSession(address string, user string) Session {
+func NewSession(address string, user string, id uint32) Session {
 	return &BaseSession{
+		id:     id,
 		addr:   address,
 		user:   user,
 		config: defaultSessionConfig(),
@@ -96,6 +103,7 @@ func NewBaseSession() Session {
 type Context struct {
 	context.Context
 	Session
+	pid    uint64
 	tracer opentracing.Tracer
 }
 
@@ -116,6 +124,13 @@ func WithTracer(t opentracing.Tracer) ContextOption {
 	}
 }
 
+// WithPid adds the given pid to the context.
+func WithPid(pid uint64) ContextOption {
+	return func(ctx *Context) {
+		ctx.pid = pid
+	}
+}
+
 // NewContext creates a new query context. Options can be passed to configure
 // the context. If some aspect of the context is not configure, the default
 // value will be used.
@@ -124,7 +139,7 @@ func NewContext(
 	ctx context.Context,
 	opts ...ContextOption,
 ) *Context {
-	c := &Context{ctx, NewBaseSession(), opentracing.NoopTracer{}}
+	c := &Context{ctx, NewBaseSession(), 0, opentracing.NoopTracer{}}
 	for _, opt := range opts {
 		opt(c)
 	}
@@ -134,6 +149,9 @@ func NewContext(
 // NewEmptyContext returns a default context with default values.
 func NewEmptyContext() *Context { return NewContext(context.TODO()) }
 
+// Pid returns the process id associated with this context.
+func (c *Context) Pid() uint64 { return c.pid }
+
 // Span creates a new tracing span with the given context.
 // It will return the span and a new context that should be passed to all
 // childrens of this span.
@@ -148,14 +166,12 @@ func (c *Context) Span(
 	span := c.tracer.StartSpan(opName, opts...)
 	ctx := opentracing.ContextWithSpan(c.Context, span)
 
-	return span, &Context{ctx, c.Session, c.tracer}
+	return span, &Context{ctx, c.Session, c.Pid(), c.tracer}
 }
 
-// WithQuery adds the query to the context.
-func (c *Context) WithQuery(query string) *Context {
-	nc := *c
-	nc.Context = context.WithValue(c.Context, QueryKey, query)
-	return &nc
+// WithContext returns a new context with the given underlying context.
+func (c *Context) WithContext(ctx context.Context) *Context {
+	return &Context{ctx, c.Session, c.Pid(), c.tracer}
 }
 
 // NewSpanIter creates a RowIter executed in the given span.


Test patch:
diff --git a/engine_test.go b/engine_test.go
index d3372de..ac62b17 100644
--- a/engine_test.go
+++ b/engine_test.go
@@ -336,7 +336,7 @@ func TestOrderByColumns(t *testing.T) {
 	require := require.New(t)
 	e := newEngine(t)
 
-	_, iter, err := e.Query(sql.NewEmptyContext(), "SELECT s, i FROM mytable ORDER BY 2 DESC")
+	_, iter, err := e.Query(newCtx(), "SELECT s, i FROM mytable ORDER BY 2 DESC")
 	require.NoError(err)
 
 	rows, err := sql.RowIterToRows(iter)
@@ -401,7 +401,7 @@ func TestAmbiguousColumnResolution(t *testing.T) {
 	e.AddDatabase(db)
 
 	q := `SELECT f.a, bar.b, f.b FROM foo f INNER JOIN bar ON f.a = bar.c`
-	ctx := sql.NewEmptyContext()
+	ctx := newCtx()
 
 	_, rows, err := e.Query(ctx, q)
 	require.NoError(err)
@@ -490,7 +490,7 @@ func TestNaturalJoin(t *testing.T) {
 	e := sqle.NewDefault()
 	e.AddDatabase(db)
 
-	_, iter, err := e.Query(sql.NewEmptyContext(), `SELECT * FROM t1 NATURAL JOIN t2`)
+	_, iter, err := e.Query(newCtx(), `SELECT * FROM t1 NATURAL JOIN t2`)
 	require.NoError(err)
 
 	rows, err := sql.RowIterToRows(iter)
@@ -542,7 +542,7 @@ func TestNaturalJoinEqual(t *testing.T) {
 	e := sqle.NewDefault()
 	e.AddDatabase(db)
 
-	_, iter, err := e.Query(sql.NewEmptyContext(), `SELECT * FROM t1 NATURAL JOIN t2`)
+	_, iter, err := e.Query(newCtx(), `SELECT * FROM t1 NATURAL JOIN t2`)
 	require.NoError(err)
 
 	rows, err := sql.RowIterToRows(iter)
@@ -589,7 +589,7 @@ func TestNaturalJoinDisjoint(t *testing.T) {
 	e := sqle.NewDefault()
 	e.AddDatabase(db)
 
-	_, iter, err := e.Query(sql.NewEmptyContext(), `SELECT * FROM t1 NATURAL JOIN t2`)
+	_, iter, err := e.Query(newCtx(), `SELECT * FROM t1 NATURAL JOIN t2`)
 	require.NoError(err)
 
 	rows, err := sql.RowIterToRows(iter)
@@ -661,7 +661,7 @@ func TestInnerNestedInNaturalJoins(t *testing.T) {
 	e := sqle.NewDefault()
 	e.AddDatabase(db)
 
-	_, iter, err := e.Query(sql.NewEmptyContext(), `SELECT * FROM table1 INNER JOIN table2 ON table1.i = table2.i2 NATURAL JOIN table3`)
+	_, iter, err := e.Query(newCtx(), `SELECT * FROM table1 INNER JOIN table2 ON table1.i = table2.i2 NATURAL JOIN table3`)
 	require.NoError(err)
 
 	rows, err := sql.RowIterToRows(iter)
@@ -681,7 +681,7 @@ func TestInnerNestedInNaturalJoins(t *testing.T) {
 func testQuery(t *testing.T, e *sqle.Engine, q string, expected []sql.Row) {
 	t.Run(q, func(t *testing.T) {
 		require := require.New(t)
-		session := sql.NewEmptyContext()
+		session := newCtx()
 
 		_, iter, err := e.Query(session, q)
 		require.NoError(err)
@@ -765,7 +765,7 @@ const expectedTree = `Offset(2)
 
 func TestPrintTree(t *testing.T) {
 	require := require.New(t)
-	node, err := parse.Parse(sql.NewEmptyContext(), `
+	node, err := parse.Parse(newCtx(), `
 		SELECT t.foo, bar.baz
 		FROM tbl t
 		INNER JOIN bar
@@ -782,7 +782,7 @@ func TestStarPanic197(t *testing.T) {
 	require := require.New(t)
 	e := newEngine(t)
 
-	ctx := sql.NewEmptyContext()
+	ctx := newCtx()
 	_, iter, err := e.Query(ctx, `SELECT * FROM mytable GROUP BY i, s`)
 	require.NoError(err)
 
@@ -802,13 +802,13 @@ func TestIndexes(t *testing.T) {
 	e.Catalog.RegisterIndexDriver(pilosa.NewIndexDriver(tmpDir))
 
 	_, _, err = e.Query(
-		sql.NewEmptyContext(),
+		newCtx(),
 		"CREATE INDEX myidx ON mytable USING pilosa (i) WITH (async = false)",
 	)
 	require.NoError(t, err)
 
 	_, _, err = e.Query(
-		sql.NewEmptyContext(),
+		newCtx(),
 		"CREATE INDEX myidx_multi ON mytable USING pilosa (i, s) WITH (async = false)",
 	)
 	require.NoError(t, err)
@@ -920,7 +920,7 @@ func TestCreateIndex(t *testing.T) {
 	require.NoError(os.MkdirAll(tmpDir, 0644))
 	e.Catalog.RegisterIndexDriver(pilosa.NewIndexDriver(tmpDir))
 
-	_, iter, err := e.Query(sql.NewEmptyContext(), "CREATE INDEX myidx ON mytable USING pilosa (i)")
+	_, iter, err := e.Query(newCtx(), "CREATE INDEX myidx ON mytable USING pilosa (i)")
 	require.NoError(err)
 	rows, err := sql.RowIterToRows(iter)
 	require.NoError(err)
@@ -961,7 +961,7 @@ func TestOrderByGroupBy(t *testing.T) {
 	e.AddDatabase(db)
 
 	_, iter, err := e.Query(
-		sql.NewEmptyContext(),
+		newCtx(),
 		"SELECT team, COUNT(*) FROM members GROUP BY team ORDER BY 2",
 	)
 	require.NoError(err)
@@ -1038,18 +1038,18 @@ func TestReadOnly(t *testing.T) {
 	a.CurrentDatabase = "mydb"
 	e := sqle.New(catalog, a, nil)
 
-	_, _, err := e.Query(sql.NewEmptyContext(), `SELECT i FROM mytable`)
+	_, _, err := e.Query(newCtx(), `SELECT i FROM mytable`)
 	require.NoError(err)
 
-	_, _, err = e.Query(sql.NewEmptyContext(), `CREATE INDEX foo ON mytable USING pilosa (i, s)`)
+	_, _, err = e.Query(newCtx(), `CREATE INDEX foo ON mytable USING pilosa (i, s)`)
 	require.Error(err)
 	require.True(analyzer.ErrQueryNotAllowed.Is(err))
 
-	_, _, err = e.Query(sql.NewEmptyContext(), `DROP INDEX foo ON mytable`)
+	_, _, err = e.Query(newCtx(), `DROP INDEX foo ON mytable`)
 	require.Error(err)
 	require.True(analyzer.ErrQueryNotAllowed.Is(err))
 
-	_, _, err = e.Query(sql.NewEmptyContext(), `INSERT INTO foo (i, s) VALUES(42, 'yolo')`)
+	_, _, err = e.Query(newCtx(), `INSERT INTO foo (i, s) VALUES(42, 'yolo')`)
 	require.Error(err)
 	require.True(analyzer.ErrQueryNotAllowed.Is(err))
 }
@@ -1059,10 +1059,14 @@ func TestSessionVariables(t *testing.T) {
 
 	e := newEngine(t)
 
-	ctx := sql.NewEmptyContext()
+	session := sql.NewBaseSession()
+	ctx := sql.NewContext(context.Background(), sql.WithSession(session), sql.WithPid(1))
+
 	_, _, err := e.Query(ctx, `set autocommit=1, sql_mode = concat(@@sql_mode,',STRICT_TRANS_TABLES')`)
 	require.NoError(err)
 
+	ctx = sql.NewContext(context.Background(), sql.WithSession(session), sql.WithPid(2))
+
 	_, iter, err := e.Query(ctx, `SELECT @@autocommit, @@session.sql_mode`)
 	require.NoError(err)
 
@@ -1076,6 +1080,13 @@ func insertRows(t *testing.T, table sql.Inserter, rows ...sql.Row) {
 	t.Helper()
 
 	for _, r := range rows {
-		require.NoError(t, table.Insert(sql.NewEmptyContext(), r))
+		require.NoError(t, table.Insert(newCtx(), r))
 	}
 }
+
+var pid uint64
+
+func newCtx() *sql.Context {
+	pid++
+	return sql.NewContext(context.Background(), sql.WithPid(pid))
+}
diff --git a/server/handler_test.go b/server/handler_test.go
index 7010720..34b0d6b 100644
--- a/server/handler_test.go
+++ b/server/handler_test.go
@@ -1,7 +1,10 @@
 package server
 
 import (
+	"net"
+	"reflect"
 	"testing"
+	"unsafe"
 
 	"gopkg.in/src-d/go-mysql-server.v0"
 	"gopkg.in/src-d/go-mysql-server.v0/mem"
@@ -130,11 +133,25 @@ func TestHandlerOutput(t *testing.T) {
 }
 
 func newConn(id uint32) *mysql.Conn {
-	return &mysql.Conn{
+	conn := &mysql.Conn{
 		ConnectionID: id,
 	}
+
+	// Set conn so it does not panic when we close it
+	val := reflect.ValueOf(conn).Elem()
+	field := val.FieldByName("conn")
+	field = reflect.NewAt(field.Type(), unsafe.Pointer(field.UnsafeAddr())).Elem()
+	field.Set(reflect.ValueOf(new(mockConn)))
+
+	return conn
+}
+
+type mockConn struct {
+	net.Conn
 }
 
+func (c *mockConn) Close() error { return nil }
+
 func TestHandlerKill(t *testing.T) {
 	require := require.New(t)
 	e := setupMemDB(require)
@@ -153,15 +170,12 @@ func TestHandlerKill(t *testing.T) {
 	require.Len(handler.c, 0)
 
 	conn1 := newConn(1)
-
 	handler.NewConnection(conn1)
-
-	require.Len(handler.c, 1)
-	c, ok := handler.c[1]
-	require.True(ok)
-	require.Equal(conn1, c)
-
 	conn2 := newConn(2)
+	handler.NewConnection(conn2)
+
+	require.Len(handler.sm.sessions, 2)
+	require.Len(handler.c, 2)
 
 	err := handler.ComQuery(conn2, "KILL QUERY 1", func(res *sqltypes.Result) error {
 		return nil
@@ -169,11 +183,17 @@ func TestHandlerKill(t *testing.T) {
 
 	require.NoError(err)
 
-	require.Len(handler.c, 1)
-	c, ok = handler.c[1]
-	require.True(ok)
-	require.Equal(conn1, c)
+	require.Len(handler.sm.sessions, 2)
+	require.Len(handler.c, 2)
+	require.Equal(conn1, handler.c[1])
+	require.Equal(conn2, handler.c[2])
+
+	err = handler.ComQuery(conn2, "KILL 1", func(res *sqltypes.Result) error {
+		return nil
+	})
+	require.NoError(err)
 
-	// Cannot test KILL CONNECTION as the connection can not be mocked. Calling
-	// mysql.Conn.Close panics.
+	require.Len(handler.sm.sessions, 1)
+	require.Len(handler.c, 1)
+	require.Equal(conn1, handler.c[1])
 }
diff --git a/sql/analyzer/process_test.go b/sql/analyzer/process_test.go
index f725df9..7a866b5 100644
--- a/sql/analyzer/process_test.go
+++ b/sql/analyzer/process_test.go
@@ -1,7 +1,9 @@
 package analyzer
 
 import (
+	"context"
 	"testing"
+	"time"
 
 	"github.com/stretchr/testify/require"
 	"gopkg.in/src-d/go-mysql-server.v0/mem"
@@ -22,7 +24,9 @@ func TestTrackProcess(t *testing.T) {
 		expression.NewLiteral(int64(1), sql.Int64),
 	)
 
-	ctx := sql.NewEmptyContext().WithQuery("SELECT foo")
+	ctx := sql.NewContext(context.Background(), sql.WithPid(1))
+	ctx, err := catalog.AddProcess(ctx, sql.QueryProcess, "SELECT foo")
+	require.NoError(err)
 
 	result, err := rule.Apply(ctx, a, node)
 	require.NoError(err)
@@ -58,6 +62,12 @@ func TestTrackProcess(t *testing.T) {
 	require.NoError(err)
 
 	require.Len(catalog.Processes(), 0)
+
+	select {
+	case <-ctx.Done():
+	case <-time.After(5 * time.Millisecond):
+		t.Errorf("expecting context to be cancelled")
+	}
 }
 
 func withoutProcessTracking(a *Analyzer) *Analyzer {
diff --git a/sql/plan/drop_index_test.go b/sql/plan/drop_index_test.go
index 02bd79a..f5746b5 100644
--- a/sql/plan/drop_index_test.go
+++ b/sql/plan/drop_index_test.go
@@ -52,3 +52,48 @@ func TestDeleteIndex(t *testing.T) {
 	require.Equal([]string{"idx"}, driver.deleted)
 	require.Nil(catalog.Index("foo", "idx"))
 }
+
+func TestDeleteIndexNotReady(t *testing.T) {
+	require := require.New(t)
+
+	table := mem.NewTable("foo", sql.Schema{
+		{Name: "a", Source: "foo"},
+		{Name: "b", Source: "foo"},
+		{Name: "c", Source: "foo"},
+	})
+
+	driver := new(mockDriver)
+	catalog := sql.NewCatalog()
+	catalog.RegisterIndexDriver(driver)
+	db := mem.NewDatabase("foo")
+	db.AddTable("foo", table)
+	catalog.Databases = append(catalog.Databases, db)
+
+	var expressions = []sql.Expression{
+		expression.NewGetFieldWithTable(0, sql.Int64, "foo", "c", true),
+		expression.NewGetFieldWithTable(1, sql.Int64, "foo", "a", true),
+	}
+
+	done, ready, err := catalog.AddIndex(&mockIndex{id: "idx", db: "foo", table: "foo", exprs: expressions})
+	require.NoError(err)
+
+	idx := catalog.Index("foo", "idx")
+	require.NotNil(idx)
+	catalog.ReleaseIndex(idx)
+
+	di := NewDropIndex("idx", NewResolvedTable(table))
+	di.Catalog = catalog
+	di.CurrentDatabase = "foo"
+
+	_, err = di.RowIter(sql.NewEmptyContext())
+	require.Error(err)
+	require.True(ErrIndexNotAvailable.Is(err))
+
+	time.Sleep(50 * time.Millisecond)
+
+	require.Equal(([]string)(nil), driver.deleted)
+	require.NotNil(catalog.Index("foo", "idx"))
+
+	close(done)
+	<-ready
+}
diff --git a/sql/plan/processlist_test.go b/sql/plan/processlist_test.go
index e206f95..83bea41 100644
--- a/sql/plan/processlist_test.go
+++ b/sql/plan/processlist_test.go
@@ -13,15 +13,20 @@ func TestShowProcessList(t *testing.T) {
 
 	n := NewShowProcessList()
 	p := sql.NewProcessList()
-	sess := sql.NewSession("0.0.0.0:1234", "foo")
-	ctx := sql.NewContext(context.Background(), sql.WithSession(sess))
+	sess := sql.NewSession("0.0.0.0:1234", "foo", 1)
+	ctx := sql.NewContext(context.Background(), sql.WithPid(1), sql.WithSession(sess))
 
-	pid := p.AddProcess(ctx, sql.QueryProcess, "SELECT foo")
-	p.AddProgressItem(pid, "a", 5)
-	p.AddProgressItem(pid, "b", 6)
+	ctx, err := p.AddProcess(ctx, sql.QueryProcess, "SELECT foo")
+	require.NoError(err)
+
+	p.AddProgressItem(ctx.Pid(), "a", 5)
+	p.AddProgressItem(ctx.Pid(), "b", 6)
+
+	ctx = sql.NewContext(context.Background(), sql.WithPid(2), sql.WithSession(sess))
+	ctx, err = p.AddProcess(ctx, sql.CreateIndexProcess, "SELECT bar")
+	require.NoError(err)
 
-	pid = p.AddProcess(ctx, sql.CreateIndexProcess, "SELECT bar")
-	p.AddProgressItem(pid, "foo", 2)
+	p.AddProgressItem(ctx.Pid(), "foo", 2)
 
 	p.UpdateProgress(1, "a", 3)
 	p.UpdateProgress(1, "a", 1)
diff --git a/sql/processlist_test.go b/sql/processlist_test.go
index ca57ddc..3c96739 100644
--- a/sql/processlist_test.go
+++ b/sql/processlist_test.go
@@ -12,33 +12,40 @@ func TestProcessList(t *testing.T) {
 	require := require.New(t)
 
 	p := NewProcessList()
-	sess := NewSession("0.0.0.0:1234", "foo")
-	ctx := NewContext(context.Background(), WithSession(sess))
-	pid := p.AddProcess(ctx, QueryProcess, "SELECT foo")
+	sess := NewSession("0.0.0.0:1234", "foo", 1)
+	ctx := NewContext(context.Background(), WithPid(1), WithSession(sess))
+	ctx, err := p.AddProcess(ctx, QueryProcess, "SELECT foo")
+	require.NoError(err)
 
-	require.Equal(uint64(1), pid)
+	require.Equal(uint64(1), ctx.Pid())
 	require.Len(p.procs, 1)
 
-	p.AddProgressItem(pid, "a", 5)
-	p.AddProgressItem(pid, "b", 6)
+	p.AddProgressItem(ctx.Pid(), "a", 5)
+	p.AddProgressItem(ctx.Pid(), "b", 6)
 
 	expectedProcess := &Process{
-		Pid:  1,
-		Type: QueryProcess,
+		Pid:        1,
+		Connection: 1,
+		Type:       QueryProcess,
 		Progress: map[string]Progress{
 			"a": Progress{0, 5},
 			"b": Progress{0, 6},
 		},
 		User:      "foo",
 		Query:     "SELECT foo",
-		StartedAt: p.procs[pid].StartedAt,
+		StartedAt: p.procs[ctx.Pid()].StartedAt,
 	}
-	require.Equal(expectedProcess, p.procs[pid])
+	require.NotNil(p.procs[ctx.Pid()].Kill)
+	p.procs[ctx.Pid()].Kill = nil
+	require.Equal(expectedProcess, p.procs[ctx.Pid()])
 
-	pid = p.AddProcess(ctx, CreateIndexProcess, "SELECT bar")
-	p.AddProgressItem(pid, "foo", 2)
+	ctx = NewContext(context.Background(), WithPid(2), WithSession(sess))
+	ctx, err = p.AddProcess(ctx, CreateIndexProcess, "SELECT bar")
+	require.NoError(err)
 
-	require.Equal(uint64(2), pid)
+	p.AddProgressItem(ctx.Pid(), "foo", 2)
+
+	require.Equal(uint64(2), ctx.Pid())
 	require.Len(p.procs, 2)
 
 	p.UpdateProgress(1, "a", 3)
@@ -52,10 +59,15 @@ func TestProcessList(t *testing.T) {
 
 	var expected []Process
 	for _, p := range p.procs {
-		expected = append(expected, *p)
+		np := *p
+		np.Kill = nil
+		expected = append(expected, np)
 	}
 
 	result := p.Processes()
+	for i := range result {
+		result[i].Kill = nil
+	}
 
 	sortByPid(expected)
 	sortByPid(result)
diff --git a/sql/session_test.go b/sql/session_test.go
index f50f82e..76740ce 100644
--- a/sql/session_test.go
+++ b/sql/session_test.go
@@ -11,7 +11,7 @@ import (
 func TestSessionConfig(t *testing.T) {
 	require := require.New(t)
 
-	sess := NewSession("foo", "bar")
+	sess := NewSession("foo", "bar", 1)
 
 	typ, v := sess.Get("foo")
 	require.Equal(Null, typ)