make sure use the time zone setting correctly

Signed-off-by: Andres Taylor <[email protected]>

go/vt/vtgate/executor_set_test.go

@@ -625,3 +625,21 @@ func TestExecutorSetAndSelect(t *testing.T) {
 		})
 	}
 }
+
+// TestTimeZone verifies that setting different time zones in the session
+// results in different outputs for the `now()` function.
+func TestTimeZone(t *testing.T) {
+	e, _, _, _, ctx := createExecutorEnv(t)
+
+	session := NewAutocommitSession(&vtgatepb.Session{TargetString: KsTestUnsharded, EnableSystemSettings: true})
+	session.SystemVariables = map[string]string{"time_zone": "'+08:00'"}
+	qr, err := e.Execute(ctx, nil, "TestExecutorSetAndSelect", session, "select now()", nil)
+
+	require.NoError(t, err)
+	session.SystemVariables["time_zone"] = "'+02:00'"
+
+	qrWith, err := e.Execute(ctx, nil, "TestExecutorSetAndSelect", session, "select now()", nil)
+	require.NoError(t, err)
+
+	assert.False(t, qr.Rows[0][0].Equal(qrWith.Rows[0][0]), "%v vs %v", qr.Rows[0][0].ToString(), qrWith.Rows[0][0].ToString())
+}

go/vt/vtgate/safe_session.go

@@ -25,8 +25,6 @@ import (
 
 	"google.golang.org/protobuf/proto"
 
-	"vitess.io/vitess/go/mysql/datetime"
-
 	"vitess.io/vitess/go/vt/sqlparser"
 	"vitess.io/vitess/go/vt/srvtopo"
 	"vitess.io/vitess/go/vt/sysvars"
@@ -561,16 +559,10 @@ func (session *SafeSession) HasSystemVariables() (found bool) {
 	return
 }
 
-func (session *SafeSession) TimeZone() *time.Location {
+func (session *SafeSession) TimeZone() string {
 	session.mu.Lock()
-	tz, ok := session.SystemVariables["time_zone"]
-	session.mu.Unlock()
-
-	if !ok {
-		return nil
-	}
-	loc, _ := datetime.ParseTimeZone(tz)
-	return loc
+	defer session.mu.Unlock()
+	return session.SystemVariables["time_zone"]
 }
 
 // ForeignKeyChecks returns the foreign_key_checks stored in system_variables map in the session.

go/vt/vtgate/safe_session_test.go

@@ -19,9 +19,7 @@ package vtgate
 import (
 	"reflect"
 	"testing"
-	"time"
 
-	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/require"
 
 	querypb "vitess.io/vitess/go/vt/proto/query"
@@ -66,35 +64,3 @@ func TestPrequeries(t *testing.T) {
 		t.Errorf("got %v but wanted %v", preQueries, want)
 	}
 }
-
-func TestTimeZone(t *testing.T) {
-	testCases := []struct {
-		tz   string
-		want string
-	}{
-		{
-			tz:   "Europe/Amsterdam",
-			want: "Europe/Amsterdam",
-		},
-		{
-			tz:   "+02:00",
-			want: "UTC+02:00",
-		},
-		{
-			tz:   "foo",
-			want: (*time.Location)(nil).String(),
-		},
-	}
-
-	for _, tc := range testCases {
-		t.Run(tc.tz, func(t *testing.T) {
-			session := NewSafeSession(&vtgatepb.Session{
-				SystemVariables: map[string]string{
-					"time_zone": tc.tz,
-				},
-			})
-
-			assert.Equal(t, tc.want, session.TimeZone().String())
-		})
-	}
-}

go/vt/vtgate/vcursor_impl.go

@@ -25,6 +25,9 @@ import (
 	"sync/atomic"
 	"time"
 
+	"vitess.io/vitess/go/mysql/datetime"
+	"vitess.io/vitess/go/vt/vtgate/evalengine"
+
 	"github.com/google/uuid"
 
 	"vitess.io/vitess/go/mysql/collations"
@@ -230,7 +233,36 @@ func (vc *vcursorImpl) Environment() *vtenv.Environment {
 }
 
 func (vc *vcursorImpl) TimeZone() *time.Location {
-	return vc.safeSession.TimeZone()
+	zone := vc.safeSession.TimeZone()
+	if zone == "" {
+		return nil
+	}
+
+	ast, err := vc.Environment().Parser().ParseExpr(zone)
+	if err != nil {
+		return nil
+	}
+
+	cfg := &evalengine.Config{
+		NoConstantFolding: true,
+		NoCompilation:     true,
+	}
+
+	expr, err := evalengine.Translate(ast, cfg)
+	if err != nil {
+		return nil
+	}
+
+	env := evalengine.EmptyExpressionEnv(vc.Environment())
+	result, err := env.Evaluate(expr)
+	if err != nil {
+		return nil
+	}
+
+	loc, _ := datetime.ParseTimeZone(result.Value(vc.collation).ToString())
+	// it's safe to ignore the error - if we get an error, loc will be nil,
+	// and this is exactly the behaviour we want anyway
+	return loc
 }
 
 func (vc *vcursorImpl) SQLMode() string {