fix(semantics): type DuckDB JSON casts, to_json, AT TIME ZONE and TO_TIMESTAMP

crates/sqlquack-semantics/src/duckdb/profile/_helpers/duck_types.rs

@@ -4,7 +4,7 @@ use sqlquack_syntax::DataType;
 
 use crate::duckdb::profile::constants::{
     DUCKDB_DECIMAL_LIMIT, DUCKDB_DEFAULT_DECIMAL, DUCKDB_INTEGER_ALIASES, DUCKDB_INTEGER_SPELLINGS,
-    DUCKDB_WIDEST_INTEGER,
+    DUCKDB_JSON_TYPE, DUCKDB_TIMESTAMP_PRECISION, DUCKDB_WIDEST_INTEGER,
 };
 
 /// What a DuckDB type is, as far as the type rules distinguish it.
@@ -181,6 +181,10 @@ pub(crate) fn cast_type(data_type: &DataType) -> Option<DuckType> {
         DataType::Binary { .. } | DataType::VarBinary { .. } | DataType::Blob => DuckType::Binary,
         DataType::Date => DuckType::Date,
         DataType::Time { timezone, .. } => DuckType::Time(*timezone),
+        DataType::Timestamp {
+            precision: Some(precision),
+            ..
+        } if *precision != DUCKDB_TIMESTAMP_PRECISION => return None,
         DataType::Timestamp { timezone, .. } => DuckType::Timestamp(*timezone),
         DataType::Interval { .. } => DuckType::Interval,
         DataType::Uuid => DuckType::Uuid,
@@ -222,6 +226,7 @@ pub(crate) fn cast_spelling(data_type: &DataType) -> Option<String> {
         DataType::Array { element_type, .. } | DataType::List { element_type } => {
             Some(format!("{}[]", cast_spelling(element_type)?))
         }
+        DataType::Json => Some(DUCKDB_JSON_TYPE.to_owned()),
         _ => cast_type(data_type).map(|target| target.spelling()),
     }
 }

crates/sqlquack-semantics/src/duckdb/profile/_helpers/typing.rs

@@ -11,8 +11,8 @@ use crate::duckdb::profile::_helpers::duck_types::{
 };
 use crate::duckdb::profile::constants::{
     DUCKDB_BIGINT_BITS, DUCKDB_DECIMAL_LIMIT, DUCKDB_EPOCH_PART, DUCKDB_INTEGER_BITS,
-    DUCKDB_JSON_TYPES, DUCKDB_JULIAN_PART, DUCKDB_SMALL_DECIMAL, DUCKDB_WIDEST_INTEGER,
-    ONE_ARGUMENT, TWO_ARGUMENTS,
+    DUCKDB_JSON_TYPE, DUCKDB_JSON_TYPES, DUCKDB_JULIAN_PART, DUCKDB_SMALL_DECIMAL,
+    DUCKDB_WIDEST_INTEGER, ONE_ARGUMENT, TWO_ARGUMENTS,
 };
 use crate::lineage::models::Lineage;
 
@@ -113,6 +113,16 @@ pub(crate) fn duckdb_type(
     if let Some(same) = passed_through(expression, &operands) {
         return Some(same);
     }
+    let to_json: bool = match expression {
+        Expression::ToJson(_) => true,
+        Expression::Function(function) => {
+            !function.quoted && function.name.eq_ignore_ascii_case("to_json")
+        }
+        _ => false,
+    };
+    if to_json {
+        return Some(Some(Rc::from(DUCKDB_JSON_TYPE)));
+    }
     let owned: Owned = operator_type(expression, &operands)
         .or_else(|| aggregate_type(expression, &operands))
         .or_else(|| scalar_type(expression, &operands))
@@ -290,6 +300,15 @@ fn scalar_type(expression: &Expression, operands: &Operands<'_, '_>) -> Owned {
         Expression::Uuid(_) => Some(Some(DuckType::Uuid)),
         Expression::Extract(function) => Some(extracted_type(&function.field)),
         Expression::Interval(_) => Some(Some(DuckType::Interval)),
+        Expression::AtTimeZone(function) => Some(zoned_type(operands.typed(&function.this))),
+        _ => None,
+    }
+}
+
+/// AT TIME ZONE's type: a TIMESTAMP gains the zone, a TIMESTAMPTZ loses it; others unknown.
+fn zoned_type(value: Option<DuckType>) -> Option<DuckType> {
+    match value? {
+        DuckType::Timestamp(zone) => Some(DuckType::Timestamp(!zone)),
         _ => None,
     }
 }
@@ -356,6 +375,11 @@ fn named_type(name: &str, arguments: &[Expression], operands: &Operands<'_, '_>)
         "CURRENT_TIME" => Some(Some(DuckType::Time(true))),
         "LOCALTIME" => Some(Some(DuckType::Time(false))),
         "LOCALTIMESTAMP" => Some(Some(DuckType::Timestamp(false))),
+        "TO_TIMESTAMP" if arguments.len() == ONE_ARGUMENT => Some(Some(DuckType::Timestamp(true))),
+        "TO_TIMESTAMP" => Some(None),
+        "TIMEZONE" if arguments.len() == TWO_ARGUMENTS => {
+            Some(zoned_type(operands.typed(&arguments[1])))
+        }
         "CURRENT_DATE" | "TODAY" => Some(Some(DuckType::Date)),
         "USER" | "CURRENT_USER" | "SESSION_USER" | "CURRENT_ROLE" | "CURRENT_SCHEMA"
         | "CURRENT_CATALOG" | "CURRENT_DATABASE"

crates/sqlquack-semantics/src/duckdb/profile/constants.rs

@@ -36,6 +36,10 @@ pub(crate) const DUCKDB_INTEGER_SPELLINGS: [(u16, bool, &str); 10] = [
     (64, true, "UBIGINT"),
     (128, true, "UHUGEINT"),
 ];
+/// The precision of DuckDB's TIMESTAMP, microseconds; any other precision is another type.
+pub(crate) const DUCKDB_TIMESTAMP_PRECISION: u32 = 6;
+/// The type of a JSON value: a cast to JSON or what to_json returns.
+pub(crate) const DUCKDB_JSON_TYPE: &str = "JSON";
 /// Semi-structured DuckDB types CONCAT does not read as text.
 pub(crate) const DUCKDB_JSON_TYPES: [&str; 2] = ["JSON", "JSONB"];
 /// DuckDB's Julian day date part, a DOUBLE.

crates/sqlquack-semantics/src/duckdb/profile/tests/test_output_types.rs

@@ -1,8 +1,11 @@
 use std::collections::HashMap;
 
+use crate::dialects::constants::DUCKDB;
 use crate::duckdb::profile::tests::helpers::duckdb_output_type;
 use crate::duckdb::profile::tests::helpers::replayed_engine_types;
 use crate::duckdb::profile::tests::test_types::DuckdbTypeTestCase;
+use crate::lineage::main::statement_lineage::statement_lineage;
+use crate::lineage::models::LineageContext;
 use crate::lineage::tests::helpers::typed_column;
 use crate::lineage::tests::test_types::{EngineTypeReplayTestCase, TypedShopSchema};
 
@@ -16,8 +19,8 @@ fn given_engine_recorded_duckdb_types_when_typing_outputs_then_none_is_wrong() {
         fixture: ENGINE_TYPES,
         replayed: replayed_engine_types,
         expected_wrong: &[],
-        expected_minimum_typed: 2556,
-        expected_maximum_untyped: 2,
+        expected_minimum_typed: 2568,
+        expected_maximum_untyped: 1,
     }];
     for test_case in test_cases {
         let replay = (test_case.replayed)(test_case.fixture);
@@ -175,3 +178,121 @@ fn given_duckdb_functions_when_typing_outputs_then_each_has_the_engine_type() {
         );
     }
 }
+
+#[test]
+fn given_json_and_time_zone_expressions_when_typing_outputs_then_both_parsers_give_the_engine_type()
+{
+    // Each type is DuckDB 2.0.0's typeof over the same columns.
+    let test_cases = [
+        DuckdbTypeTestCase {
+            description: "a cast to JSON is JSON",
+            sql: "SELECT CAST(name AS JSON) AS v FROM t",
+            expected_type: Some("JSON"),
+        },
+        DuckdbTypeTestCase {
+            description: "a :: cast to JSON is JSON",
+            sql: "SELECT name::JSON AS v FROM t",
+            expected_type: Some("JSON"),
+        },
+        DuckdbTypeTestCase {
+            description: "a cast to a JSON list is JSON[]",
+            sql: "SELECT CAST(name AS JSON[]) AS v FROM t",
+            expected_type: Some("JSON[]"),
+        },
+        DuckdbTypeTestCase {
+            description: "to_json gives JSON",
+            sql: "SELECT to_json(quantity) AS v FROM t",
+            expected_type: Some("JSON"),
+        },
+        DuckdbTypeTestCase {
+            description: "a JSON cast over to_json of a JSON COALESCE is JSON",
+            sql: "SELECT CAST(to_json(COALESCE(name::JSON, '[]'::JSON)) AS JSON) AS v FROM t",
+            expected_type: Some("JSON"),
+        },
+        DuckdbTypeTestCase {
+            description: "TO_TIMESTAMP of a number is TIMESTAMPTZ",
+            sql: "SELECT TO_TIMESTAMP(quantity / POWER(10, 3)) AS v FROM t",
+            expected_type: Some("TIMESTAMPTZ"),
+        },
+        DuckdbTypeTestCase {
+            description: "AT TIME ZONE gives a TIMESTAMP its zone",
+            sql: "SELECT created_at AT TIME ZONE 'UTC' AS v FROM t",
+            expected_type: Some("TIMESTAMPTZ"),
+        },
+        DuckdbTypeTestCase {
+            description: "AT TIME ZONE takes a TIMESTAMPTZ's zone away",
+            sql: "SELECT seen_at AT TIME ZONE 'UTC' AS v FROM t",
+            expected_type: Some("TIMESTAMP"),
+        },
+        DuckdbTypeTestCase {
+            description: "timezone() is AT TIME ZONE",
+            sql: "SELECT timezone('UTC', seen_at) AS v FROM t",
+            expected_type: Some("TIMESTAMP"),
+        },
+        DuckdbTypeTestCase {
+            description: "a COALESCE of a TO_TIMESTAMP at a zone and a TIMESTAMP is TIMESTAMP",
+            sql: "SELECT MIN(COALESCE(TO_TIMESTAMP(quantity / POWER(10, 3)) AT TIME ZONE 'UTC', created_at)) AS v FROM t",
+            expected_type: Some("TIMESTAMP"),
+        },
+        DuckdbTypeTestCase {
+            description: "a cast to TIMESTAMP(3), DuckDB's TIMESTAMP_MS, stays untyped",
+            sql: "SELECT CAST(name AS TIMESTAMP(3)) AS v FROM t",
+            expected_type: None,
+        },
+        DuckdbTypeTestCase {
+            description: "AT TIME ZONE over a TIMESTAMP(3) cast, DuckDB's TIMESTAMPTZ_NS, stays untyped",
+            sql: "SELECT CAST(name AS TIMESTAMP(3)) AT TIME ZONE 'UTC' AS v FROM t",
+            expected_type: None,
+        },
+        DuckdbTypeTestCase {
+            description: "AT TIME ZONE over a TIMESTAMP(6) cast is TIMESTAMPTZ",
+            sql: "SELECT CAST(name AS TIMESTAMP(6)) AT TIME ZONE 'UTC' AS v FROM t",
+            expected_type: Some("TIMESTAMPTZ"),
+        },
+        DuckdbTypeTestCase {
+            description: "AT TIME ZONE over a DATE stays untyped",
+            sql: "SELECT CAST(created_at AS DATE) AT TIME ZONE 'UTC' AS v FROM t",
+            expected_type: None,
+        },
+    ];
+    let schema: TypedShopSchema = TypedShopSchema {
+        closed: HashMap::from([(
+            "t".to_owned(),
+            vec![
+                typed_column("quantity", "INTEGER"),
+                typed_column("name", "VARCHAR"),
+                typed_column("created_at", "TIMESTAMP"),
+                typed_column("seen_at", "TIMESTAMPTZ"),
+            ],
+        )]),
+        open: HashMap::new(),
+        functions: HashMap::new(),
+        rules: HashMap::new(),
+        declared: Vec::new(),
+        arities: HashMap::new(),
+    };
+    for test_case in test_cases {
+        let statement = sqlquack_syntax_duckdb::parse_strict(test_case.sql)
+            .unwrap_or_else(|failure| panic!("{}: {failure:?}", test_case.description))
+            .statements
+            .remove(0);
+        let context = LineageContext {
+            profile: &DUCKDB,
+            schema: &schema,
+            sql: Some(test_case.sql),
+        };
+        let native: Option<String> = statement_lineage(&statement, context).outputs[0]
+            .data_type
+            .as_deref()
+            .map(str::to_owned);
+        assert_eq!(
+            (
+                duckdb_output_type(&schema, test_case.sql).as_deref(),
+                native.as_deref()
+            ),
+            (test_case.expected_type, test_case.expected_type),
+            "{}",
+            test_case.description
+        );
+    }
+}