date_trunc invalid units with infinite value

From c046465d4532e9b9086e0aadd13526e2a284b072 Mon Sep 17 00:00:00 2001
From: Joseph Koshakow <koshy44@gmail.com>
Date: Sun, 1 Dec 2024 14:55:35 -0500
Subject: [PATCH v1] Fix date_trunc units for infinite intervals

Previously, if an infinite value was passed to date_trunc, then the
same infinite value would always be returned regardless of the unit.
This commit updates the function so that an error is returned when an
invalid unit is passed to date_trunc with an infinite value.

This matches the behavior of date_trunc with a finite value and
date_part with an infinite value.
---
 src/backend/utils/adt/timestamp.c         | 121 ++++++++++++++++++----
 src/test/regress/expected/interval.out    |   9 ++
 src/test/regress/expected/timestamp.out   |   2 +
 src/test/regress/expected/timestamptz.out |   4 +
 src/test/regress/sql/interval.sql         |   8 ++
 src/test/regress/sql/timestamp.sql        |   2 +
 src/test/regress/sql/timestamptz.sql      |   4 +
 7 files changed, 131 insertions(+), 19 deletions(-)

diff --git a/src/backend/utils/adt/timestamp.c b/src/backend/utils/adt/timestamp.c
index e00cd3c969..a2489caa6b 100644
--- a/src/backend/utils/adt/timestamp.c
+++ b/src/backend/utils/adt/timestamp.c
@@ -4620,9 +4620,6 @@ timestamp_trunc(PG_FUNCTION_ARGS)
 	struct pg_tm tt,
 			   *tm = &tt;
 
-	if (TIMESTAMP_NOT_FINITE(timestamp))
-		PG_RETURN_TIMESTAMP(timestamp);
-
 	lowunits = downcase_truncate_identifier(VARDATA_ANY(units),
 											VARSIZE_ANY_EXHDR(units),
 											false);
@@ -4631,6 +4628,39 @@ timestamp_trunc(PG_FUNCTION_ARGS)
 
 	if (type == UNITS)
 	{
+		if (TIMESTAMP_NOT_FINITE(timestamp))
+		{
+			/*
+			 * Errors thrown here for invalid units should exactly match those
+			 * below, else there will be unexpected discrepancies between
+			 * finite- and infinite-input cases.
+			 */
+			switch (val)
+			{
+				case DTK_WEEK:
+				case DTK_MILLENNIUM:
+				case DTK_CENTURY:
+				case DTK_DECADE:
+				case DTK_YEAR:
+				case DTK_QUARTER:
+				case DTK_MONTH:
+				case DTK_DAY:
+				case DTK_HOUR:
+				case DTK_MINUTE:
+				case DTK_SECOND:
+				case DTK_MILLISEC:
+				case DTK_MICROSEC:
+					PG_RETURN_TIMESTAMP(timestamp);
+					break;
+				default:
+					ereport(ERROR,
+							(errcode(ERRCODE_FEATURE_NOT_SUPPORTED),
+							 errmsg("unit \"%s\" not supported for type %s",
+									lowunits, format_type_be(TIMESTAMPOID))));
+					result = 0;
+			}
+		}
+
 		if (timestamp2tm(timestamp, NULL, tm, &fsec, NULL, NULL) != 0)
 			ereport(ERROR,
 					(errcode(ERRCODE_DATETIME_VALUE_OUT_OF_RANGE),
@@ -4836,6 +4866,40 @@ timestamptz_trunc_internal(text *units, TimestampTz timestamp, pg_tz *tzp)
 
 	if (type == UNITS)
 	{
+		if (TIMESTAMP_NOT_FINITE(timestamp))
+		{
+			/*
+			 * Errors thrown here for invalid units should exactly match those
+			 * below, else there will be unexpected discrepancies between
+			 * finite- and infinite-input cases.
+			 */
+			switch (val)
+			{
+				case DTK_WEEK:
+				case DTK_MILLENNIUM:
+				case DTK_CENTURY:
+				case DTK_DECADE:
+				case DTK_YEAR:
+				case DTK_QUARTER:
+				case DTK_MONTH:
+				case DTK_DAY:
+				case DTK_HOUR:
+				case DTK_MINUTE:
+				case DTK_SECOND:
+				case DTK_MILLISEC:
+				case DTK_MICROSEC:
+					PG_RETURN_TIMESTAMPTZ(timestamp);
+					break;
+
+				default:
+					ereport(ERROR,
+							(errcode(ERRCODE_FEATURE_NOT_SUPPORTED),
+							 errmsg("unit \"%s\" not supported for type %s",
+									lowunits, format_type_be(TIMESTAMPTZOID))));
+					result = 0;
+			}
+		}
+
 		if (timestamp2tm(timestamp, &tz, tm, &fsec, NULL, tzp) != 0)
 			ereport(ERROR,
 					(errcode(ERRCODE_DATETIME_VALUE_OUT_OF_RANGE),
@@ -4966,9 +5030,6 @@ timestamptz_trunc(PG_FUNCTION_ARGS)
 	TimestampTz timestamp = PG_GETARG_TIMESTAMPTZ(1);
 	TimestampTz result;
 
-	if (TIMESTAMP_NOT_FINITE(timestamp))
-		PG_RETURN_TIMESTAMPTZ(timestamp);
-
 	result = timestamptz_trunc_internal(units, timestamp, session_timezone);
 
 	PG_RETURN_TIMESTAMPTZ(result);
@@ -4986,13 +5047,6 @@ timestamptz_trunc_zone(PG_FUNCTION_ARGS)
 	TimestampTz result;
 	pg_tz	   *tzp;
 
-	/*
-	 * timestamptz_zone() doesn't look up the zone for infinite inputs, so we
-	 * don't do so here either.
-	 */
-	if (TIMESTAMP_NOT_FINITE(timestamp))
-		PG_RETURN_TIMESTAMP(timestamp);
-
 	/*
 	 * Look up the requested timezone.
 	 */
@@ -5020,12 +5074,6 @@ interval_trunc(PG_FUNCTION_ARGS)
 
 	result = (Interval *) palloc(sizeof(Interval));
 
-	if (INTERVAL_NOT_FINITE(interval))
-	{
-		memcpy(result, interval, sizeof(Interval));
-		PG_RETURN_INTERVAL_P(result);
-	}
-
 	lowunits = downcase_truncate_identifier(VARDATA_ANY(units),
 											VARSIZE_ANY_EXHDR(units),
 											false);
@@ -5034,6 +5082,41 @@ interval_trunc(PG_FUNCTION_ARGS)
 
 	if (type == UNITS)
 	{
+		if (INTERVAL_NOT_FINITE(interval))
+		{
+			/*
+			 * Errors thrown here for invalid units should exactly match those
+			 * below, else there will be unexpected discrepancies between
+			 * finite- and infinite-input cases.
+			 */
+			switch (val)
+			{
+				case DTK_MILLENNIUM:
+				case DTK_CENTURY:
+				case DTK_DECADE:
+				case DTK_YEAR:
+				case DTK_QUARTER:
+				case DTK_MONTH:
+				case DTK_DAY:
+				case DTK_HOUR:
+				case DTK_MINUTE:
+				case DTK_SECOND:
+				case DTK_MILLISEC:
+				case DTK_MICROSEC:
+					memcpy(result, interval, sizeof(Interval));
+					PG_RETURN_INTERVAL_P(result);
+					break;
+
+				default:
+					ereport(ERROR,
+							(errcode(ERRCODE_FEATURE_NOT_SUPPORTED),
+							 errmsg("unit \"%s\" not supported for type %s",
+									lowunits, format_type_be(INTERVALOID)),
+							 (val == DTK_WEEK) ? errdetail("Months usually have fractional weeks.") : 0));
+					result = 0;
+			}
+		}
+
 		interval2itm(*interval, tm);
 		switch (val)
 		{
diff --git a/src/test/regress/expected/interval.out b/src/test/regress/expected/interval.out
index e5d919d0cf..a16e3ccdb2 100644
--- a/src/test/regress/expected/interval.out
+++ b/src/test/regress/expected/interval.out
@@ -2219,6 +2219,15 @@ SELECT i AS interval, date_trunc('hour', i)
  -infinity | -infinity
 (2 rows)
 
+SELECT i AS interval, date_trunc('week', i)
+    FROM INFINITE_INTERVAL_TBL
+    WHERE NOT isfinite(i);
+ERROR:  unit "week" not supported for type interval
+DETAIL:  Months usually have fractional weeks.
+SELECT i AS interval, date_trunc('ago', i)
+    FROM INFINITE_INTERVAL_TBL
+    WHERE NOT isfinite(i);
+ERROR:  unit "ago" not recognized for type interval
 SELECT i AS interval, justify_days(i), justify_hours(i), justify_interval(i)
     FROM INFINITE_INTERVAL_TBL
     WHERE NOT isfinite(i);
diff --git a/src/test/regress/expected/timestamp.out b/src/test/regress/expected/timestamp.out
index e287260051..6aaa19c8f4 100644
--- a/src/test/regress/expected/timestamp.out
+++ b/src/test/regress/expected/timestamp.out
@@ -591,6 +591,8 @@ SELECT date_trunc( 'week', timestamp '2004-02-29 15:44:17.71393' ) AS week_trunc
  Mon Feb 23 00:00:00 2004
 (1 row)
 
+SELECT date_trunc( 'ago', timestamp 'infinity' ) AS invalid_trunc;
+ERROR:  unit "ago" not recognized for type timestamp without time zone
 -- verify date_bin behaves the same as date_trunc for relevant intervals
 -- case 1: AD dates, origin < input
 SELECT
diff --git a/src/test/regress/expected/timestamptz.out b/src/test/regress/expected/timestamptz.out
index b437613ac8..a6dd45626c 100644
--- a/src/test/regress/expected/timestamptz.out
+++ b/src/test/regress/expected/timestamptz.out
@@ -701,6 +701,8 @@ SELECT date_trunc( 'week', timestamp with time zone '2004-02-29 15:44:17.71393'
  Mon Feb 23 00:00:00 2004 PST
 (1 row)
 
+SELECT date_trunc( 'ago', timestamp with time zone 'infinity' ) AS invalid_trunc;
+ERROR:  unit "ago" not recognized for type timestamp with time zone
 SELECT date_trunc('day', timestamp with time zone '2001-02-16 20:38:40+00', 'Australia/Sydney') as sydney_trunc;  -- zone name
          sydney_trunc         
 ------------------------------
@@ -719,6 +721,8 @@ SELECT date_trunc('day', timestamp with time zone '2001-02-16 20:38:40+00', 'VET
  Thu Feb 15 20:00:00 2001 PST
 (1 row)
 
+SELECT date_trunc('ago', timestamp with time zone 'infinity', 'GMT') AS invalid_zone_trunc;
+ERROR:  unit "ago" not recognized for type timestamp with time zone
 -- verify date_bin behaves the same as date_trunc for relevant intervals
 SELECT
   str,
diff --git a/src/test/regress/sql/interval.sql b/src/test/regress/sql/interval.sql
index 55054ae65d..43bc793925 100644
--- a/src/test/regress/sql/interval.sql
+++ b/src/test/regress/sql/interval.sql
@@ -776,6 +776,14 @@ SELECT i AS interval, date_trunc('hour', i)
     FROM INFINITE_INTERVAL_TBL
     WHERE NOT isfinite(i);
 
+SELECT i AS interval, date_trunc('week', i)
+    FROM INFINITE_INTERVAL_TBL
+    WHERE NOT isfinite(i);
+
+SELECT i AS interval, date_trunc('ago', i)
+    FROM INFINITE_INTERVAL_TBL
+    WHERE NOT isfinite(i);
+
 SELECT i AS interval, justify_days(i), justify_hours(i), justify_interval(i)
     FROM INFINITE_INTERVAL_TBL
     WHERE NOT isfinite(i);
diff --git a/src/test/regress/sql/timestamp.sql b/src/test/regress/sql/timestamp.sql
index 748469576d..55f80530ea 100644
--- a/src/test/regress/sql/timestamp.sql
+++ b/src/test/regress/sql/timestamp.sql
@@ -176,6 +176,8 @@ SELECT d1 - timestamp without time zone '1997-01-02' AS diff
 
 SELECT date_trunc( 'week', timestamp '2004-02-29 15:44:17.71393' ) AS week_trunc;
 
+SELECT date_trunc( 'ago', timestamp 'infinity' ) AS invalid_trunc;
+
 -- verify date_bin behaves the same as date_trunc for relevant intervals
 
 -- case 1: AD dates, origin < input
diff --git a/src/test/regress/sql/timestamptz.sql b/src/test/regress/sql/timestamptz.sql
index 6b91e7eddc..a92586c363 100644
--- a/src/test/regress/sql/timestamptz.sql
+++ b/src/test/regress/sql/timestamptz.sql
@@ -200,10 +200,14 @@ SELECT d1 - timestamp with time zone '1997-01-02' AS diff
    FROM TIMESTAMPTZ_TBL WHERE d1 BETWEEN '1902-01-01' AND '2038-01-01';
 
 SELECT date_trunc( 'week', timestamp with time zone '2004-02-29 15:44:17.71393' ) AS week_trunc;
+SELECT date_trunc( 'ago', timestamp with time zone 'infinity' ) AS invalid_trunc;
 
 SELECT date_trunc('day', timestamp with time zone '2001-02-16 20:38:40+00', 'Australia/Sydney') as sydney_trunc;  -- zone name
 SELECT date_trunc('day', timestamp with time zone '2001-02-16 20:38:40+00', 'GMT') as gmt_trunc;  -- fixed-offset abbreviation
 SELECT date_trunc('day', timestamp with time zone '2001-02-16 20:38:40+00', 'VET') as vet_trunc;  -- variable-offset abbreviation
+SELECT date_trunc('ago', timestamp with time zone 'infinity', 'GMT') AS invalid_zone_trunc;
+
+
 
 -- verify date_bin behaves the same as date_trunc for relevant intervals
 SELECT
-- 
2.43.0