From db345a2ce76c11ae2d0b9eee5f2a94846201b627 Mon Sep 17 00:00:00 2001
From: Finn Bear <finnbearlabs@gmail.com>
Date: Fri, 12 May 2023 10:59:58 -0700
Subject: [PATCH] Bugfix - time::ceil function edge case. (#1976)

---
 lib/src/fnc/time.rs   | 33 ++++++++++++++++++++++-----------
 lib/tests/function.rs | 14 ++++++++++++--
 2 files changed, 34 insertions(+), 13 deletions(-)

diff --git a/lib/src/fnc/time.rs b/lib/src/fnc/time.rs
index 17e8547e..aa19b766 100644
--- a/lib/src/fnc/time.rs
+++ b/lib/src/fnc/time.rs
@@ -3,20 +3,31 @@ use crate::sql::datetime::Datetime;
 use crate::sql::duration::Duration;
 use crate::sql::value::Value;
 use chrono::offset::TimeZone;
-use chrono::Datelike;
-use chrono::DurationRound;
-use chrono::Local;
-use chrono::Timelike;
-use chrono::Utc;
+use chrono::{DateTime, Datelike, DurationRound, Local, Timelike, Utc};
 
 pub fn ceil((val, duration): (Datetime, Duration)) -> Result<Value, Error> {
 	match chrono::Duration::from_std(*duration) {
-		Ok(d) => match val.duration_trunc(d).ok().and_then(|floor| floor.checked_add_signed(d)) {
-			Some(v) => Ok(v.into()),
-			_ => Err(Error::InvalidArguments {
-				name: String::from("time::ceil"),
-				message: String::from("The second argument must be a duration, and must be able to be represented as nanoseconds."),
-			}),
+		Ok(d) => {
+			let floor_to_ceil = |floor: DateTime<Utc>| -> Option<DateTime<Utc>> {
+				if floor == *val {
+					Some(floor)
+				} else {
+					floor.checked_add_signed(d)
+				}
+			};
+
+			let result = val
+				.duration_trunc(d)
+				.ok()
+				.and_then(floor_to_ceil);
+
+			match result {
+				Some(v) => Ok(v.into()),
+				_ => Err(Error::InvalidArguments {
+					name: String::from("time::ceil"),
+					message: String::from("The second argument must be a duration, and must be able to be represented as nanoseconds."),
+				}),
+			}
 		},
 		_ => Err(Error::InvalidArguments {
 			name: String::from("time::ceil"),
diff --git a/lib/tests/function.rs b/lib/tests/function.rs
index a6e96cf8..a63ffdc7 100644
--- a/lib/tests/function.rs
+++ b/lib/tests/function.rs
@@ -3559,11 +3559,12 @@ async fn function_time_ceil() -> Result<(), Error> {
 	let sql = r#"
 		RETURN time::ceil("1987-06-22T08:30:45Z", 1w);
 		RETURN time::ceil("1987-06-22T08:30:45Z", 1y);
+		RETURN time::ceil("2023-05-11T03:09:00Z", 1s);
 	"#;
 	let dbs = Datastore::new("memory").await?;
 	let ses = Session::for_kv().with_ns("test").with_db("test");
 	let res = &mut dbs.execute(&sql, &ses, None, false).await?;
-	assert_eq!(res.len(), 2);
+	assert_eq!(res.len(), 3);
 	//
 	let tmp = res.remove(0).result?;
 	let val = Value::parse("'1987-06-25T00:00:00Z'");
@@ -3573,6 +3574,10 @@ async fn function_time_ceil() -> Result<(), Error> {
 	let val = Value::parse("'1987-12-28T00:00:00Z'");
 	assert_eq!(tmp, val);
 	//
+	let tmp = res.remove(0).result?;
+	let val = Value::parse("'2023-05-11T03:09:00Z'");
+	assert_eq!(tmp, val);
+	//
 	Ok(())
 }
 
@@ -3602,11 +3607,12 @@ async fn function_time_floor() -> Result<(), Error> {
 	let sql = r#"
 		RETURN time::floor("1987-06-22T08:30:45Z", 1w);
 		RETURN time::floor("1987-06-22T08:30:45Z", 1y);
+		RETURN time::floor("2023-05-11T03:09:00Z", 1s);
 	"#;
 	let dbs = Datastore::new("memory").await?;
 	let ses = Session::for_kv().with_ns("test").with_db("test");
 	let res = &mut dbs.execute(&sql, &ses, None, false).await?;
-	assert_eq!(res.len(), 2);
+	assert_eq!(res.len(), 3);
 	//
 	let tmp = res.remove(0).result?;
 	let val = Value::parse("'1987-06-18T00:00:00Z'");
@@ -3616,6 +3622,10 @@ async fn function_time_floor() -> Result<(), Error> {
 	let val = Value::parse("'1986-12-28T00:00:00Z'");
 	assert_eq!(tmp, val);
 	//
+	let tmp = res.remove(0).result?;
+	let val = Value::parse("'2023-05-11T03:09:00Z'");
+	assert_eq!(tmp, val);
+	//
 	Ok(())
 }