Merge #96837 #96873

96837: sql: correctly set location on `writeTextDatum` r=rafiss a=otan

In c4f38dd, we made `fmtPgwireText` more accurate by correctly using the session location where possible and panicking if a location was not set. Unfortunately this was forgotten in a certain place on the writeTextDatum path when writing arrays out.

fixes #96877
fixes #96853

Release note: None

96873: kv/kvserver: skip TestDelegateSnapshot r=arulajmani a=andrewbaptist

Refs: #96841

Reason: flaky test

Release justification: non-production code changes

Release note: None

Co-authored-by: Oliver Tan <otan@cockroachlabs.com>
Co-authored-by: Andrew Baptist <baptist@cockroachlabs.com>

pkg/kv/kvserver/replica_learner_test.go

@@ -356,7 +356,7 @@ func TestAddReplicaWithReceiverThrottling(t *testing.T) {
 func TestDelegateSnapshot(t *testing.T) {
 	defer leaktest.AfterTest(t)()
 	defer log.Scope(t).Close(t)
-	skip.UnderStress(t, "Occasionally fails until 87553 is resolved")
+	skip.WithIssue(t, 96841, "Occasionally fails until 87553 is resolved")
 
 	ctx := context.Background()
 

pkg/sql/logictest/testdata/logic_test/timestamp

@@ -656,6 +656,24 @@ SELECT to_char(d, 'FF1 FF2 FF3 FF4 FF5 FF6  ff1 ff2 ff3 ff4 ff5 ff6  MS US')
 7 78 789 7890 78901 789010  7 78 789 7890 78901 789010  789 789010
 7 78 789 7890 78901 789012  7 78 789 7890 78901 789012  789 789012
 
+query T
+SET timezone = '00:00';
+SELECT ARRAY[
+  '2018-11-02 12:34:56.78901234'::timestamptz,
+  '2018-11-02 12:34:56.78901234-07'::timestamptz
+]
+----
+{"2018-11-02 12:34:56.789012+00","2018-11-02 19:34:56.789012+00"}
+
+query T
+SET timezone = 'Australia/Sydney';
+SELECT ARRAY[
+  '2018-11-02 12:34:56.78901234'::timestamptz,
+  '2018-11-02 12:34:56.78901234-07'::timestamptz
+]
+----
+{"2018-11-02 12:34:56.789012+11","2018-11-03 06:34:56.789012+11"}
+
 query TT
 SET timezone = '00:00';
 SELECT to_char(now(), 'OF') as of_t, to_char(now(), 'TZH:TZM') as "TZH:TZM";

pkg/sql/pgwire/conn.go

@@ -1442,7 +1442,7 @@ func (c *tenantEgressCounter) GetBatchNetworkEgress(
 			// Use the default values for the DataConversionConfig and location.
 			// See the comment in getRowNetworkEgress for why the writeText variant
 			// is used here instead of writeBinary.
-			c.buf.writeTextColumnarElement(ctx, &c.vecs, vecIdx, rowIdx, conv, nil /* sessionLoc */)
+			c.buf.writeTextColumnarElement(ctx, &c.vecs, vecIdx, rowIdx, conv, time.UTC)
 			egress += int64(c.buf.Len())
 			c.buf.reset()
 		}

pkg/sql/pgwire/types.go

@@ -157,8 +157,13 @@ func writeTextDatumNotNull(
 	sessionLoc *time.Location,
 	t *types.T,
 ) {
+
 	oldDCC := b.textFormatter.SetDataConversionConfig(conv)
-	defer b.textFormatter.SetDataConversionConfig(oldDCC)
+	oldLoc := b.textFormatter.SetLocation(sessionLoc)
+	defer func() {
+		b.textFormatter.SetDataConversionConfig(oldDCC)
+		b.textFormatter.SetLocation(oldLoc)
+	}()
 	switch v := tree.UnwrapDOidWrapper(d).(type) {
 	case *tree.DBitArray:
 		b.textFormatter.FormatNode(v)
@@ -293,7 +298,11 @@ func (b *writeBuffer) writeTextColumnarElement(
 	sessionLoc *time.Location,
 ) {
 	oldDCC := b.textFormatter.SetDataConversionConfig(conv)
-	defer b.textFormatter.SetDataConversionConfig(oldDCC)
+	oldLoc := b.textFormatter.SetLocation(sessionLoc)
+	defer func() {
+		b.textFormatter.SetDataConversionConfig(oldDCC)
+		b.textFormatter.SetLocation(oldLoc)
+	}()
 	typ := vecs.Vecs[vecIdx].Type()
 	if log.V(2) {
 		log.Infof(ctx, "pgwire writing TEXT columnar element of type: %s", typ)

pkg/sql/sem/tree/format.go

@@ -361,6 +361,13 @@ func (ctx *FmtCtx) SetDataConversionConfig(
 	return old
 }
 
+// SetLocation sets the location on ctx and returns the old one.
+func (ctx *FmtCtx) SetLocation(loc *time.Location) *time.Location {
+	old := ctx.location
+	ctx.location = loc
+	return old
+}
+
 // WithReformatTableNames modifies FmtCtx to to substitute the printing of table
 // names using the provided function, calls fn, then restores the original table
 // formatting.