From 1c370d4aafc325586e1cafac3c762d82ac3dc208 Mon Sep 17 00:00:00 2001
From: Daylon Wilkins <daylon@liquidata.co>
Date: Tue, 22 Apr 2025 05:49:49 -0700
Subject: [PATCH] Added interpretation awareness

---
 processlist.go                   | 12 ++++++++++++
 sql/rowexec/transaction_iters.go |  2 +-
 sql/session.go                   | 19 +++++++++++++++++++
 3 files changed, 32 insertions(+), 1 deletion(-)

diff --git a/processlist.go b/processlist.go
index 40561ba994..6ea2b95904 100644
--- a/processlist.go
+++ b/processlist.go
@@ -113,6 +113,9 @@ func (pl *ProcessList) BeginQuery(
 	ctx *sql.Context,
 	query string,
 ) (*sql.Context, error) {
+	if ctx.IsInterpreted() {
+		return ctx, nil
+	}
 	pl.mu.Lock()
 	defer pl.mu.Unlock()
 
@@ -144,6 +147,9 @@ func (pl *ProcessList) BeginQuery(
 }
 
 func (pl *ProcessList) EndQuery(ctx *sql.Context) {
+	if ctx.IsInterpreted() {
+		return
+	}
 	pl.mu.Lock()
 	defer pl.mu.Unlock()
 	id := ctx.Session.ID()
@@ -177,6 +183,9 @@ func (pl *ProcessList) EndQuery(ctx *sql.Context) {
 // bracketed with EndOperation(). Should certainly be used for any
 // Handler callbacks which may access the database, like Prepare.
 func (pl *ProcessList) BeginOperation(ctx *sql.Context) (*sql.Context, error) {
+	if ctx.IsInterpreted() {
+		return ctx, nil
+	}
 	pl.mu.Lock()
 	defer pl.mu.Unlock()
 	id := ctx.Session.ID()
@@ -193,6 +202,9 @@ func (pl *ProcessList) BeginOperation(ctx *sql.Context) (*sql.Context, error) {
 }
 
 func (pl *ProcessList) EndOperation(ctx *sql.Context) {
+	if ctx.IsInterpreted() {
+		return
+	}
 	pl.mu.Lock()
 	defer pl.mu.Unlock()
 	id := ctx.Session.ID()
diff --git a/sql/rowexec/transaction_iters.go b/sql/rowexec/transaction_iters.go
index 69c45d0d5e..db69cf5327 100644
--- a/sql/rowexec/transaction_iters.go
+++ b/sql/rowexec/transaction_iters.go
@@ -78,7 +78,7 @@ type TransactionCommittingIter struct {
 
 func AddTransactionCommittingIter(ctx *sql.Context, qFlags *sql.QueryFlags, iter sql.RowIter) (sql.RowIter, error) {
 	// TODO: This is a bit of a hack. Need to figure out better relationship between new transaction node and warnings.
-	if qFlags != nil && qFlags.IsSet(sql.QFlagShowWarnings) {
+	if (qFlags != nil && qFlags.IsSet(sql.QFlagShowWarnings)) || ctx.IsInterpreted() {
 		return iter, nil
 	}
 
diff --git a/sql/session.go b/sql/session.go
index fcca905a45..e5cca1021a 100644
--- a/sql/session.go
+++ b/sql/session.go
@@ -255,6 +255,7 @@ type Context struct {
 	queryTime   time.Time
 	tracer      trace.Tracer
 	rootSpan    trace.Span
+	interpreted bool
 	Version     AnalyzerVersion
 }
 
@@ -328,6 +329,17 @@ func RunWithNowFunc(nowFunc func() time.Time, fn func() error) error {
 	return fn()
 }
 
+// RunInterpreted modifies the context such that all calls to Context.IsInterpreted will return `true`. It is safe to
+// recursively call this.
+func RunInterpreted[T any](ctx *Context, f func(ctx *Context) (T, error)) (T, error) {
+	current := ctx.interpreted
+	ctx.interpreted = true
+	defer func() {
+		ctx.interpreted = current
+	}()
+	return f(ctx)
+}
+
 func swapNowFunc(newNowFunc func() time.Time) func() time.Time {
 	ctxNowFuncMutex.Lock()
 	defer ctxNowFuncMutex.Unlock()
@@ -388,6 +400,13 @@ func (c *Context) ApplyOpts(opts ...ContextOption) {
 // NewEmptyContext returns a default context with default values.
 func NewEmptyContext() *Context { return NewContext(context.TODO()) }
 
+// IsInterpreted returns `true` when this is being called from within RunInterpreted. In such cases, GMS will choose to
+// handle logic differently, as running from within an interpreted function requires different considerations than
+// running in a standard environment.
+func (c *Context) IsInterpreted() bool {
+	return c.interpreted
+}
+
 // Pid returns the process id associated with this context.
 func (c *Context) Pid() uint64 {
 	if c == nil {