diff --git a/crates/sqlquack-semantics/src/duckdb/profile/_helpers/duck_types.rs b/crates/sqlquack-semantics/src/duckdb/profile/_helpers/duck_types.rs index ec07e354a..709ef549a 100644 --- a/crates/sqlquack-semantics/src/duckdb/profile/_helpers/duck_types.rs +++ b/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 { 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 { 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()), } } diff --git a/crates/sqlquack-semantics/src/duckdb/profile/_helpers/typing.rs b/crates/sqlquack-semantics/src/duckdb/profile/_helpers/typing.rs index 3b2a4f07c..59c1c4077 100644 --- a/crates/sqlquack-semantics/src/duckdb/profile/_helpers/typing.rs +++ b/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) -> Option { + 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" diff --git a/crates/sqlquack-semantics/src/duckdb/profile/constants.rs b/crates/sqlquack-semantics/src/duckdb/profile/constants.rs index 6cc35ce43..ca4e489f7 100644 --- a/crates/sqlquack-semantics/src/duckdb/profile/constants.rs +++ b/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. diff --git a/crates/sqlquack-semantics/src/duckdb/profile/tests/test_output_types.rs b/crates/sqlquack-semantics/src/duckdb/profile/tests/test_output_types.rs index 7b0c438f1..977f611d5 100644 --- a/crates/sqlquack-semantics/src/duckdb/profile/tests/test_output_types.rs +++ b/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 = 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 + ); + } +}