fix(cubesql): Support EXTRACT(EPOCH) over timestamp diff on Snowflake (cube-js#11105)

Author: MazterQyou

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

@@ -114,6 +114,12 @@ export class SnowflakeQuery extends BaseQuery {
     templates.functions.BTRIM = 'TRIM({{ args_concat }})';
     templates.functions.STRING_AGG = 'LISTAGG({% if distinct %}DISTINCT {% endif %}{{ args_concat }})';
     templates.expressions.extract = 'EXTRACT({{ date_part }} FROM {{ expr }})';
+    // Snowflake can't EXTRACT(EPOCH FROM <interval>), so the epoch of a timestamp
+    // difference (left - right) is rendered as fractional seconds between them.
+    // TIMESTAMPDIFF is measured once at microsecond granularity (no per-second
+    // boundary rounding) and divided to seconds, matching Postgres' fractional
+    // EXTRACT(EPOCH FROM interval).
+    templates.expressions.extract_epoch_diff = 'TIMESTAMPDIFF(MICROSECOND, {{ right }}, {{ left }}) / 1000000';
     templates.expressions.interval = 'INTERVAL \'{{ interval }}\'';
     templates.expressions.timestamp_literal = '\'{{ value }}\'::timestamp_tz';
     templates.expressions.like = '{{ expr }} {% if negated %}NOT {% endif %}LIKE {{ pattern }}{% if default_escape %} ESCAPE \'\\\\\'{% endif %}';

rust/cubesql/cubesql/src/compile/engine/df/wrapper.rs

@@ -3063,6 +3063,22 @@ impl WrappedSelectNode {
         ))
     }
 
+    /// Returns the operands of a timestamp/date subtraction `left - right`,
+    /// peeling any cast wrapping the subtraction. Used to rewrite
+    /// `EXTRACT(EPOCH FROM (left - right))` for dialects that don't support
+    /// taking the epoch of an interval.
+    fn timestamp_diff_operands(expr: &Expr) -> Option<(&Expr, &Expr)> {
+        match expr {
+            Expr::BinaryExpr {
+                left,
+                op: Operator::Minus,
+                right,
+            } => Some((left, right)),
+            Expr::Cast { expr, .. } => Self::timestamp_diff_operands(expr),
+            _ => None,
+        }
+    }
+
     #[inline(never)]
     fn generate_sql_for_scalar_function<'ctx>(
         mut sql_query: SqlQuery,
@@ -3124,6 +3140,42 @@ impl WrappedSelectNode {
                                         date_part
                                     )));
                         }
+                        // Some dialects (e.g. Snowflake) can't EXTRACT(EPOCH FROM <interval>),
+                        // i.e. take the epoch of a timestamp difference `a - b`. When the
+                        // dialect provides a dedicated template, render the difference as a
+                        // diff in seconds instead.
+                        let sql_templates = sql_generator.get_sql_templates();
+                        if date_part.eq_ignore_ascii_case("epoch")
+                            && sql_templates.contains_template("expressions/extract_epoch_diff")
+                        {
+                            if let Some((left, right)) = Self::timestamp_diff_operands(&args[1]) {
+                                let (left_sql, query) = Self::generate_sql_for_expr(
+                                    sql_query,
+                                    sql_generator.clone(),
+                                    left.clone(),
+                                    push_to_cube_context,
+                                    subqueries,
+                                )?;
+                                let (right_sql, query) = Self::generate_sql_for_expr(
+                                    query,
+                                    sql_generator.clone(),
+                                    right.clone(),
+                                    push_to_cube_context,
+                                    subqueries,
+                                )?;
+                                return Ok((
+                                    sql_templates
+                                        .extract_epoch_diff_expr(left_sql, right_sql)
+                                        .map_err(|e| {
+                                            DataFusionError::Internal(format!(
+                                                "Can't generate SQL for scalar function: {}",
+                                                e
+                                            ))
+                                        })?,
+                                    query,
+                                ));
+                            }
+                        }
                         let (arg_sql, query) = Self::generate_sql_for_expr(
                             sql_query,
                             sql_generator.clone(),
@@ -3132,8 +3184,7 @@ impl WrappedSelectNode {
                             subqueries,
                         )?;
                         return Ok((
-                            sql_generator
-                                .get_sql_templates()
+                            sql_templates
                                 .extract_expr(date_part.to_string(), arg_sql)
                                 .map_err(|e| {
                                     DataFusionError::Internal(format!(

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

@@ -16964,6 +16964,51 @@ LIMIT {{ limit }}{% endif %}"#.to_string(),
         assert!(sql.contains("unix_timestamp"));
     }
 
+    #[tokio::test]
+    async fn test_extract_epoch_diff_pushdown() {
+        if !Rewriter::sql_push_down_enabled() {
+            return;
+        }
+        init_testing_logger();
+
+        // EXTRACT(EPOCH FROM (a - b)) — epoch of a timestamp difference.
+        let query = "
+            SELECT customer_gender,
+                   AVG(EXTRACT(EPOCH FROM (order_date - last_mod)) / 86400) AS avg_days
+            FROM KibanaSampleDataEcommerce
+            GROUP BY 1
+        ";
+
+        // Generic (no dedicated template) keeps EXTRACT(EPOCH FROM (a - b)).
+        let query_plan =
+            convert_select_to_query_plan(query.to_string(), DatabaseProtocol::PostgreSQL).await;
+        let logical_plan = query_plan.as_logical_plan();
+        let sql = logical_plan.find_cube_scan_wrapped_sql().wrapped_sql.sql;
+        assert!(sql.contains("EXTRACT(epoch"));
+
+        // Snowflake-style: epoch of a difference is rendered as a seconds diff.
+        let query_plan = convert_select_to_query_plan_customized(
+            query.to_string(),
+            DatabaseProtocol::PostgreSQL,
+            vec![
+                (
+                    "expressions/extract".to_string(),
+                    "EXTRACT({{ date_part }} FROM {{ expr }})".to_string(),
+                ),
+                (
+                    "expressions/extract_epoch_diff".to_string(),
+                    "TIMESTAMPDIFF(MICROSECOND, {{ right }}, {{ left }}) / 1000000".to_string(),
+                ),
+            ],
+        )
+        .await;
+
+        let logical_plan = query_plan.as_logical_plan();
+        let sql = logical_plan.find_cube_scan_wrapped_sql().wrapped_sql.sql;
+        assert!(!sql.to_uppercase().contains("EXTRACT(EPOCH"));
+        assert!(sql.contains("TIMESTAMPDIFF(MICROSECOND,"));
+    }
+
     #[tokio::test]
     async fn test_push_down_to_grouped_query_with_filters() {
         if !Rewriter::sql_push_down_enabled() {

rust/cubesql/cubesql/src/transport/service.rs

@@ -777,6 +777,20 @@ impl SqlTemplates {
         )
     }
 
+    /// Renders the epoch (in seconds) of a timestamp difference `left - right`.
+    /// Used for dialects (e.g. Snowflake) where `EXTRACT(EPOCH FROM (left - right))`
+    /// is invalid because EPOCH can't be extracted from an interval.
+    pub fn extract_epoch_diff_expr(
+        &self,
+        left: String,
+        right: String,
+    ) -> Result<String, CubeError> {
+        self.render_template(
+            "expressions/extract_epoch_diff",
+            context! { left => left, right => right },
+        )
+    }
+
     pub fn interval_any_expr(
         &self,
         interval: String,