fix(cubesql): Improve Trino SQL push down compatibility (#9861)

Signed-off-by: Alex Qyoun-ae <[email protected]>

Author: MazterQyou

packages/cubejs-schema-compiler/src/adapter/PrestodbQuery.ts

@@ -136,6 +136,7 @@ export class PrestodbQuery extends BaseQuery {
     const templates = super.sqlTemplates();
     templates.functions.DATETRUNC = 'DATE_TRUNC({{ args_concat }})';
     templates.functions.DATEPART = 'DATE_PART({{ args_concat }})';
+    templates.functions.DATEDIFF = 'DATE_DIFF(\'{{ date_part }}\', {{ args[1] }}, {{ args[2] }})';
     templates.functions.CURRENTDATE = 'CURRENT_DATE';
     delete templates.functions.PERCENTILECONT;
     templates.statements.select = '{% if ctes %} WITH \n' +

rust/cubesql/cubesql/src/compile/date_parser.rs

@@ -2,9 +2,11 @@ use crate::compile::engine::df::scan::DataFusionError;
 use chrono::{NaiveDate, NaiveDateTime};
 
 pub fn parse_date_str(s: &str) -> Result<NaiveDateTime, DataFusionError> {
-    let parsed = NaiveDateTime::parse_from_str(s, "%Y-%m-%dT%H:%M:%S%.f")
+    let parsed = NaiveDateTime::parse_from_str(s, "%Y-%m-%d %H:%M:%S")
         .or_else(|_| NaiveDateTime::parse_from_str(s, "%Y-%m-%d %H:%M:%S%.f"))
         .or_else(|_| NaiveDateTime::parse_from_str(s, "%Y-%m-%dT%H:%M:%S"))
+        .or_else(|_| NaiveDateTime::parse_from_str(s, "%Y-%m-%d %H:%M:%S%.f UTC"))
+        .or_else(|_| NaiveDateTime::parse_from_str(s, "%Y-%m-%dT%H:%M:%S%.f"))
         .or_else(|_| NaiveDateTime::parse_from_str(s, "%Y-%m-%dT%H:%M:%S%.fZ"))
         .or_else(|_| {
             NaiveDate::parse_from_str(s, "%Y-%m-%d").map(|date| date.and_hms_opt(0, 0, 0).unwrap())

rust/cubesql/cubesql/src/compile/mod.rs

@@ -17267,4 +17267,51 @@ LIMIT {{ limit }}{% endif %}"#.to_string(),
         let sql = logical_plan.find_cube_scan_wrapped_sql().wrapped_sql.sql;
         assert!(sql.contains("CAST(1 AS VARCHAR)"));
     }
+
+    #[tokio::test]
+    async fn test_trino_datediff() {
+        if !Rewriter::sql_push_down_enabled() {
+            return;
+        }
+        init_testing_logger();
+
+        let query_plan = convert_select_to_query_plan_customized(
+            r#"
+            SELECT
+                KibanaSampleDataEcommerce.id,
+                KibanaSampleDataEcommerce.order_date,
+                KibanaSampleDataEcommerce.last_mod,
+                DATEDIFF(
+                    day,
+                    KibanaSampleDataEcommerce.order_date,
+                    KibanaSampleDataEcommerce.last_mod
+                ) as conv_date_diff,
+                COUNT(*)
+            FROM KibanaSampleDataEcommerce
+            WHERE (
+                KibanaSampleDataEcommerce.order_date > cast('2025-01-01T00:00:00.000' as timestamp)
+                    AND KibanaSampleDataEcommerce.order_date < cast('2025-01-01T23:59:59.999' as timestamp)
+                    AND KibanaSampleDataEcommerce.customer_gender = 'test'
+            )
+            GROUP BY 1, 2, 3, 4
+            "#
+            .to_string(),
+            DatabaseProtocol::PostgreSQL,
+            vec![(
+                "functions/DATEDIFF".to_string(),
+                "DATE_DIFF('{{ date_part }}', {{ args[1] }}, {{ args[2] }})".to_string(),
+            )],
+        )
+        .await;
+
+        let physical_plan = query_plan.as_physical_plan().await.unwrap();
+        println!(
+            "Physical plan: {}",
+            displayable(physical_plan.as_ref()).indent()
+        );
+
+        let logical_plan = query_plan.as_logical_plan();
+        let sql = logical_plan.find_cube_scan_wrapped_sql().wrapped_sql.sql;
+        assert!(sql.contains("DATE_DIFF('day', "));
+    }
 }

rust/cubesql/pg-srv/src/values/timestamp.rs

@@ -157,6 +157,7 @@ impl FromProtocolValue for TimestampValue {
         // more formats, so let's align this with parse_date_str function from cubesql crate.
         let parsed_datetime = NaiveDateTime::parse_from_str(as_str, "%Y-%m-%d %H:%M:%S")
             .or_else(|_| NaiveDateTime::parse_from_str(as_str, "%Y-%m-%d %H:%M:%S%.f"))
+            .or_else(|_| NaiveDateTime::parse_from_str(as_str, "%Y-%m-%d %H:%M:%S%.f UTC"))
             .or_else(|_| NaiveDateTime::parse_from_str(as_str, "%Y-%m-%dT%H:%M:%S"))
             .or_else(|_| NaiveDateTime::parse_from_str(as_str, "%Y-%m-%dT%H:%M:%S%.f"))
             .or_else(|_| NaiveDateTime::parse_from_str(as_str, "%Y-%m-%dT%H:%M:%S%.fZ"))