Saving parsed query to the context

Sometimes it's nice to be able to access the parsed query, say in a `Table`'s callbacks; for example, if one implements a table backed by a system that allows fast access for very specific queries, it might not be possible (or at least not easy) to take advantage of that using one of the pre-defined `*Table` interfaces. In that case, the most straightforward way to implement this would be to return the relevant subset of rows from `PartitionRows`, depending on the query.

This patch allows to access the already parsed query, instead of parsing it all over again.

server/handler.go

@@ -392,6 +392,8 @@ func (h *Handler) doQuery(
 		}
 	}
 
+	sqlCtx.SetParsedQuery(parsed)
+
 	more := remainder != ""
 
 	var queryStr string

sql/session.go

@@ -27,6 +27,8 @@ import (
 	"go.opentelemetry.io/otel/attribute"
 	"go.opentelemetry.io/otel/trace"
 	"golang.org/x/sync/errgroup"
+
+	"github.com/dolthub/vitess/go/vt/sqlparser"
 )
 
 type key uint
@@ -239,6 +241,7 @@ type Context struct {
 	services    Services
 	pid         uint64
 	query       string
+	parsedQuery sqlparser.Statement
 	queryTime   time.Time
 	tracer      trace.Tracer
 	rootSpan    trace.Span
@@ -399,6 +402,20 @@ func (c *Context) WithQuery(q string) *Context {
 	return &nc
 }
 
+// ParsedQuery returns the parsed query associated with this context.
+// May return nil.
+func (c *Context) ParsedQuery() sqlparser.Statement {
+	if c == nil {
+		return nil
+	}
+	return c.parsedQuery
+}
+
+// SetParsedQuery adds the given parsed query to the context.
+func (c *Context) SetParsedQuery(parsed sqlparser.Statement) {
+	c.parsedQuery = parsed
+}
+
 // QueryTime returns the time.Time when the context associated with this query was created
 func (c *Context) QueryTime() time.Time {
 	if c == nil {