generate_series_with_timezone.patch

commit 3bc2fc7a56ecc68b00230d37d6aec97853d499f0
Author: Przemyslaw Sztoch <psztoch@finn.pl>
Date:   Tue Jun 14 20:14:50 2022 +0200

    timestamptz plus interval with timezone

diff --git a/src/backend/utils/adt/timestamp.c b/src/backend/utils/adt/timestamp.c
index f70f829d83..c25a0db1ad 100644
--- a/src/backend/utils/adt/timestamp.c
+++ b/src/backend/utils/adt/timestamp.c
@@ -69,6 +69,7 @@ typedef struct
 	TimestampTz finish;
 	Interval	step;
 	int			step_sign;
+	char		tzname[TZ_STRLEN_MAX + 1];
 } generate_series_timestamptz_fctx;
 
 
@@ -3003,83 +3004,124 @@ timestamptz_pl_interval(PG_FUNCTION_ARGS)
 {
 	TimestampTz timestamp = PG_GETARG_TIMESTAMPTZ(0);
 	Interval   *span = PG_GETARG_INTERVAL_P(1);
-	TimestampTz result;
+	pg_tz      *attimezone = NULL;
 	int			tz;
 
 	if (TIMESTAMP_NOT_FINITE(timestamp))
-		result = timestamp;
-	else
-	{
-		if (span->month != 0)
-		{
-			struct pg_tm tt,
-					   *tm = &tt;
-			fsec_t		fsec;
-
-			if (timestamp2tm(timestamp, &tz, tm, &fsec, NULL, NULL) != 0)
-				ereport(ERROR,
-						(errcode(ERRCODE_DATETIME_VALUE_OUT_OF_RANGE),
-						 errmsg("timestamp out of range")));
+		PG_RETURN_TIMESTAMP(timestamp);
 
-			tm->tm_mon += span->month;
-			if (tm->tm_mon > MONTHS_PER_YEAR)
-			{
-				tm->tm_year += (tm->tm_mon - 1) / MONTHS_PER_YEAR;
-				tm->tm_mon = ((tm->tm_mon - 1) % MONTHS_PER_YEAR) + 1;
-			}
-			else if (tm->tm_mon < 1)
-			{
-				tm->tm_year += tm->tm_mon / MONTHS_PER_YEAR - 1;
-				tm->tm_mon = tm->tm_mon % MONTHS_PER_YEAR + MONTHS_PER_YEAR;
-			}
+	if (PG_NARGS() > 2)
+	{
+		text	   *zone = PG_GETARG_TEXT_PP(2);
+		char		tzname[TZ_STRLEN_MAX + 1];
+		char	   *lowzone;
+		int			type,
+					val;
+		pg_tz      *tzp;
+		/*
+		 * Look up the requested timezone (see notes in timestamptz_zone()).
+		 */
+		text_to_cstring_buffer(zone, tzname, sizeof(tzname));
 
-			/* adjust for end of month boundary problems... */
-			if (tm->tm_mday > day_tab[isleap(tm->tm_year)][tm->tm_mon - 1])
-				tm->tm_mday = (day_tab[isleap(tm->tm_year)][tm->tm_mon - 1]);
+		/* DecodeTimezoneAbbrev requires lowercase input */
+		lowzone = downcase_truncate_identifier(tzname,
+											   strlen(tzname),
+											   false);
 
-			tz = DetermineTimeZoneOffset(tm, session_timezone);
+		type = DecodeTimezoneAbbrev(0, lowzone, &val, &tzp);
 
-			if (tm2timestamp(tm, fsec, &tz, &timestamp) != 0)
-				ereport(ERROR,
-						(errcode(ERRCODE_DATETIME_VALUE_OUT_OF_RANGE),
-						 errmsg("timestamp out of range")));
+		if (type == TZ || type == DTZ)
+		{
+			/* fixed-offset abbreviation, get a pg_tz descriptor for that */
+			tzp = pg_tzset_offset(-val);
 		}
-
-		if (span->day != 0)
+		else if (type == DYNTZ)
 		{
-			struct pg_tm tt,
-					   *tm = &tt;
-			fsec_t		fsec;
-			int			julian;
-
-			if (timestamp2tm(timestamp, &tz, tm, &fsec, NULL, NULL) != 0)
+			/* dynamic-offset abbreviation, use its referenced timezone */
+		}
+		else
+		{
+			/* try it as a full zone name */
+			tzp = pg_tzset(tzname);
+			if (!tzp)
 				ereport(ERROR,
-						(errcode(ERRCODE_DATETIME_VALUE_OUT_OF_RANGE),
-						 errmsg("timestamp out of range")));
+						(errcode(ERRCODE_INVALID_PARAMETER_VALUE),
+						 errmsg("time zone \"%s\" not recognized", tzname)));
+		}
+		attimezone = tzp;
+	}
 
-			/* Add days by converting to and from Julian */
-			julian = date2j(tm->tm_year, tm->tm_mon, tm->tm_mday) + span->day;
-			j2date(julian, &tm->tm_year, &tm->tm_mon, &tm->tm_mday);
+	/* Use session timezone if caller asks for default */
+	if (attimezone == NULL)
+		attimezone = session_timezone;
+	
+	if (span->month != 0)
+	{
+		struct pg_tm tt,
+				   *tm = &tt;
+		fsec_t		fsec;
 
-			tz = DetermineTimeZoneOffset(tm, session_timezone);
+		if (timestamp2tm(timestamp, &tz, tm, &fsec, NULL, attimezone) != 0)
+			ereport(ERROR,
+					(errcode(ERRCODE_DATETIME_VALUE_OUT_OF_RANGE),
+					 errmsg("timestamp out of range")));
 
-			if (tm2timestamp(tm, fsec, &tz, &timestamp) != 0)
-				ereport(ERROR,
-						(errcode(ERRCODE_DATETIME_VALUE_OUT_OF_RANGE),
-						 errmsg("timestamp out of range")));
+		tm->tm_mon += span->month;
+		if (tm->tm_mon > MONTHS_PER_YEAR)
+		{
+			tm->tm_year += (tm->tm_mon - 1) / MONTHS_PER_YEAR;
+			tm->tm_mon = ((tm->tm_mon - 1) % MONTHS_PER_YEAR) + 1;
+		}
+		else if (tm->tm_mon < 1)
+		{
+			tm->tm_year += tm->tm_mon / MONTHS_PER_YEAR - 1;
+			tm->tm_mon = tm->tm_mon % MONTHS_PER_YEAR + MONTHS_PER_YEAR;
 		}
 
-		timestamp += span->time;
+		/* adjust for end of month boundary problems... */
+		if (tm->tm_mday > day_tab[isleap(tm->tm_year)][tm->tm_mon - 1])
+			tm->tm_mday = (day_tab[isleap(tm->tm_year)][tm->tm_mon - 1]);
 
-		if (!IS_VALID_TIMESTAMP(timestamp))
+		tz = DetermineTimeZoneOffset(tm, attimezone);
+
+		if (tm2timestamp(tm, fsec, &tz, &timestamp) != 0)
 			ereport(ERROR,
 					(errcode(ERRCODE_DATETIME_VALUE_OUT_OF_RANGE),
 					 errmsg("timestamp out of range")));
+	}
 
-		result = timestamp;
+	if (span->day != 0)
+	{
+		struct pg_tm tt,
+				   *tm = &tt;
+		fsec_t		fsec;
+		int			julian;
+
+		if (timestamp2tm(timestamp, &tz, tm, &fsec, NULL, attimezone) != 0)
+			ereport(ERROR,
+					(errcode(ERRCODE_DATETIME_VALUE_OUT_OF_RANGE),
+					 errmsg("timestamp out of range")));
+
+		/* Add days by converting to and from Julian */
+		julian = date2j(tm->tm_year, tm->tm_mon, tm->tm_mday) + span->day;
+		j2date(julian, &tm->tm_year, &tm->tm_mon, &tm->tm_mday);
+
+		tz = DetermineTimeZoneOffset(tm, attimezone);
+
+		if (tm2timestamp(tm, fsec, &tz, &timestamp) != 0)
+			ereport(ERROR,
+					(errcode(ERRCODE_DATETIME_VALUE_OUT_OF_RANGE),
+					 errmsg("timestamp out of range")));
 	}
 
-	PG_RETURN_TIMESTAMP(result);
+	timestamp += span->time;
+
+	if (!IS_VALID_TIMESTAMP(timestamp))
+		ereport(ERROR,
+				(errcode(ERRCODE_DATETIME_VALUE_OUT_OF_RANGE),
+				 errmsg("timestamp out of range")));
+
+	PG_RETURN_TIMESTAMP(timestamp);
 }
 
 Datum
@@ -5888,6 +5930,15 @@ generate_series_timestamptz(PG_FUNCTION_ARGS)
 					(errcode(ERRCODE_INVALID_PARAMETER_VALUE),
 					 errmsg("step size cannot equal zero")));
 
+		if (PG_NARGS() > 3)
+		{
+			text *zone = PG_GETARG_TEXT_PP(3);
+			text_to_cstring_buffer(zone, fctx->tzname, sizeof(fctx->tzname));
+		}
+		else
+		{
+			fctx->tzname[0] = 0;
+		}
 		funcctx->user_fctx = fctx;
 		MemoryContextSwitchTo(oldcontext);
 	}
@@ -5906,9 +5957,20 @@ generate_series_timestamptz(PG_FUNCTION_ARGS)
 		timestamp_cmp_internal(result, fctx->finish) >= 0)
 	{
 		/* increment current in preparation for next iteration */
-		fctx->current = DatumGetTimestampTz(DirectFunctionCall2(timestamptz_pl_interval,
-																TimestampTzGetDatum(fctx->current),
-																PointerGetDatum(&fctx->step)));
+		if (fctx->tzname[0] == 0) {
+			fctx->current = DatumGetTimestampTz(DirectFunctionCall2(timestamptz_pl_interval,
+																	TimestampTzGetDatum(fctx->current),
+																	PointerGetDatum(&fctx->step)));
+		}
+		else
+		{
+			text *tzname_text = cstring_to_text(fctx->tzname);
+
+			fctx->current = DatumGetTimestampTz(DirectFunctionCall3(timestamptz_pl_interval,
+																	TimestampTzGetDatum(fctx->current),
+																	PointerGetDatum(&fctx->step),
+																	PointerGetDatum(tzname_text)));
+		}
 
 		/* do when there is more left to send */
 		SRF_RETURN_NEXT(funcctx, TimestampTzGetDatum(result));
diff --git a/src/include/catalog/pg_proc.dat b/src/include/catalog/pg_proc.dat
index 87aa571a33..f7532e422c 100644
--- a/src/include/catalog/pg_proc.dat
+++ b/src/include/catalog/pg_proc.dat
@@ -2429,7 +2429,7 @@
   proargtypes => 'text timestamptz', prosrc => 'timestamptz_trunc' },
 { oid => '1284',
   descr => 'truncate timestamp with time zone to specified units in specified time zone',
-  proname => 'date_trunc', provolatile => 's', prorettype => 'timestamptz',
+  proname => 'date_trunc', prorettype => 'timestamptz',
   proargtypes => 'text timestamptz text', prosrc => 'timestamptz_trunc_zone' },
 { oid => '1218', descr => 'truncate interval to specified units',
   proname => 'date_trunc', prorettype => 'interval',
@@ -11885,4 +11885,15 @@
   prorettype => 'bytea', proargtypes => 'pg_brin_minmax_multi_summary',
   prosrc => 'brin_minmax_multi_summary_send' },
 
+# timestamptz plus interval with timezone patch
+{ oid => '8800',
+  descr => 'add interval to timestamp with time zone in specified time zone',
+  proname => 'date_add',
+  prorettype => 'timestamptz', proargtypes => 'timestamptz interval text',
+  prosrc => 'timestamptz_pl_interval' },
+{ oid => '8801', descr => 'non-persistent series generator',
+  proname => 'generate_series', prorows => '1000', proretset => 't',
+  prorettype => 'timestamptz',
+  proargtypes => 'timestamptz timestamptz interval text',
+  prosrc => 'generate_series_timestamptz' },
 ]