From 932d0ebf7ed6a499cbafa5b280fb8167b3195d10 Mon Sep 17 00:00:00 2001 From: Brandon Martin Date: Thu, 12 Dec 2024 09:34:59 -0800 Subject: [PATCH] Scalar type representation (#165) ### What Return proper representation for scalar types. Still need to verify and/or fix serialization of types. ### How --------- Co-authored-by: Daniel Harvey --- crates/configuration/src/introspection.rs | 30 +++++++ crates/ndc-sqlserver/src/schema.rs | 8 +- ...e_stored_procedure_with_relationships.snap | 14 +-- ...uery_tests__types__select_value_types.snap | 4 +- .../snapshots/schema_tests__get_schema.snap | 87 +++++++++++++++++++ crates/query-engine/sql/src/sql/helpers.rs | 18 +++- .../translation/src/translation/helpers.rs | 2 +- .../src/translation/query/aggregates.rs | 60 ++++++++++--- .../translation/src/translation/query/mod.rs | 10 ++- .../translation/src/translation/query/root.rs | 11 ++- 10 files changed, 213 insertions(+), 31 deletions(-) diff --git a/crates/configuration/src/introspection.rs b/crates/configuration/src/introspection.rs index 0f6550cc..a2701c0d 100644 --- a/crates/configuration/src/introspection.rs +++ b/crates/configuration/src/introspection.rs @@ -1,4 +1,5 @@ //! Configuration and state for our connector. +use ndc_models::TypeRepresentation; use serde::Deserialize; #[derive(Deserialize, Debug)] @@ -64,3 +65,32 @@ pub struct IntrospectionForeignKeyColumn { pub struct IntrospectionSchema { pub name: String, } + +pub fn map_type_representation(sql_type: &str) -> Option { + match sql_type.to_lowercase().as_str() { + "bigint" => Some(TypeRepresentation::Int64), + "bit" => Some(TypeRepresentation::Boolean), + "decimal" | "numeric" | "money" | "smallmoney" => Some(TypeRepresentation::BigDecimal), + "int" => Some(TypeRepresentation::Int32), + "smallint" => Some(TypeRepresentation::Int16), + "tinyint" => Some(TypeRepresentation::Int16), + "float" => Some(TypeRepresentation::Float64), + "real" => Some(TypeRepresentation::Float32), + "date" => Some(TypeRepresentation::Date), + "datetime" | "datetime2" | "smalldatetime" => Some(TypeRepresentation::Timestamp), + "datetimeoffset" => Some(TypeRepresentation::TimestampTZ), + "time" | "timestamp" => Some(TypeRepresentation::String), + "char" | "varchar" | "text" | "nchar" | "nvarchar" | "ntext" => { + Some(TypeRepresentation::String) + } + "binary" | "varbinary" | "image" => Some(TypeRepresentation::String), + "uniqueidentifier" => Some(TypeRepresentation::UUID), + "xml" => Some(TypeRepresentation::String), + "json" => Some(TypeRepresentation::JSON), + // TODO: Add support for hierarchyid and sql_variant + // "geometry" => Some(TypeRepresentation::Geometry), + // "geography" => Some(TypeRepresentation::Geography), + // "hierarchyid" | "sql_variant" => XXX, + _ => None, + } +} diff --git a/crates/ndc-sqlserver/src/schema.rs b/crates/ndc-sqlserver/src/schema.rs index 8d8ecab7..612a7a98 100644 --- a/crates/ndc-sqlserver/src/schema.rs +++ b/crates/ndc-sqlserver/src/schema.rs @@ -12,9 +12,8 @@ use ndc_sdk::models::ObjectTypeName; use ndc_sdk::models::ProcedureName; use ndc_sdk::models::ScalarTypeName; use ndc_sdk::models::TypeRepresentation; -use query_engine_metadata::metadata; - use ndc_sqlserver_configuration as configuration; +use query_engine_metadata::metadata; use query_engine_metadata::metadata::stored_procedures::{ StoredProcedureArgumentInfo, StoredProcedures, }; @@ -308,8 +307,9 @@ pub fn get_schema( ( scalar_type.0.clone().into(), models::ScalarType { - // TODO(PY): Add representation for beta - representation: None, + representation: configuration::introspection::map_type_representation( + scalar_type.0.as_str(), + ), aggregate_functions: metadata .aggregate_functions .0 diff --git a/crates/ndc-sqlserver/tests/snapshots/mutation_tests__stored_procedures__execute_stored_procedure_with_relationships.snap b/crates/ndc-sqlserver/tests/snapshots/mutation_tests__stored_procedures__execute_stored_procedure_with_relationships.snap index d76bbcb8..aa531911 100644 --- a/crates/ndc-sqlserver/tests/snapshots/mutation_tests__stored_procedures__execute_stored_procedure_with_relationships.snap +++ b/crates/ndc-sqlserver/tests/snapshots/mutation_tests__stored_procedures__execute_stored_procedure_with_relationships.snap @@ -13,31 +13,31 @@ expression: result "rows": [ { "InvoiceId": 98, - "Total": 3.98 + "Total": "3.98" }, { "InvoiceId": 121, - "Total": 3.96 + "Total": "3.96" }, { "InvoiceId": 143, - "Total": 5.94 + "Total": "5.94" }, { "InvoiceId": 195, - "Total": 0.99 + "Total": "0.99" }, { "InvoiceId": 316, - "Total": 1.98 + "Total": "1.98" }, { "InvoiceId": 327, - "Total": 13.86 + "Total": "13.86" }, { "InvoiceId": 382, - "Total": 8.91 + "Total": "8.91" } ] } diff --git a/crates/ndc-sqlserver/tests/snapshots/query_tests__types__select_value_types.snap b/crates/ndc-sqlserver/tests/snapshots/query_tests__types__select_value_types.snap index e26a32ed..046adeb0 100644 --- a/crates/ndc-sqlserver/tests/snapshots/query_tests__types__select_value_types.snap +++ b/crates/ndc-sqlserver/tests/snapshots/query_tests__types__select_value_types.snap @@ -6,10 +6,10 @@ expression: result { "rows": [ { - "bigint": 1, + "bigint": "1", "int": 245, "float": 24.555, - "numeric": 24.555, + "numeric": "24.555", "char": "s", "text": "textual text", "date": "2021-12-21", diff --git a/crates/ndc-sqlserver/tests/snapshots/schema_tests__get_schema.snap b/crates/ndc-sqlserver/tests/snapshots/schema_tests__get_schema.snap index ed04fbe2..4c781638 100644 --- a/crates/ndc-sqlserver/tests/snapshots/schema_tests__get_schema.snap +++ b/crates/ndc-sqlserver/tests/snapshots/schema_tests__get_schema.snap @@ -5,6 +5,9 @@ expression: result { "scalar_types": { "bigint": { + "representation": { + "type": "int64" + }, "aggregate_functions": { "APPROX_COUNT_DISTINCT": { "result_type": { @@ -118,6 +121,9 @@ expression: result } }, "binary": { + "representation": { + "type": "string" + }, "aggregate_functions": { "APPROX_COUNT_DISTINCT": { "result_type": { @@ -183,6 +189,9 @@ expression: result } }, "bit": { + "representation": { + "type": "boolean" + }, "aggregate_functions": { "APPROX_COUNT_DISTINCT": { "result_type": { @@ -248,6 +257,9 @@ expression: result } }, "char": { + "representation": { + "type": "string" + }, "aggregate_functions": { "APPROX_COUNT_DISTINCT": { "result_type": { @@ -339,6 +351,9 @@ expression: result } }, "date": { + "representation": { + "type": "date" + }, "aggregate_functions": { "APPROX_COUNT_DISTINCT": { "result_type": { @@ -404,6 +419,9 @@ expression: result } }, "datetime": { + "representation": { + "type": "timestamp" + }, "aggregate_functions": { "APPROX_COUNT_DISTINCT": { "result_type": { @@ -481,6 +499,9 @@ expression: result } }, "datetime2": { + "representation": { + "type": "timestamp" + }, "aggregate_functions": { "APPROX_COUNT_DISTINCT": { "result_type": { @@ -546,6 +567,9 @@ expression: result } }, "datetimeoffset": { + "representation": { + "type": "timestamptz" + }, "aggregate_functions": { "APPROX_COUNT_DISTINCT": { "result_type": { @@ -611,6 +635,9 @@ expression: result } }, "decimal": { + "representation": { + "type": "bigdecimal" + }, "aggregate_functions": { "APPROX_COUNT_DISTINCT": { "result_type": { @@ -724,6 +751,9 @@ expression: result } }, "float": { + "representation": { + "type": "float64" + }, "aggregate_functions": { "APPROX_COUNT_DISTINCT": { "result_type": { @@ -1032,6 +1062,9 @@ expression: result } }, "image": { + "representation": { + "type": "string" + }, "aggregate_functions": { "COUNT_BIG": { "result_type": { @@ -1050,6 +1083,9 @@ expression: result } }, "int": { + "representation": { + "type": "int32" + }, "aggregate_functions": { "APPROX_COUNT_DISTINCT": { "result_type": { @@ -1167,6 +1203,9 @@ expression: result "comparison_operators": {} }, "money": { + "representation": { + "type": "bigdecimal" + }, "aggregate_functions": { "APPROX_COUNT_DISTINCT": { "result_type": { @@ -1280,6 +1319,9 @@ expression: result } }, "nchar": { + "representation": { + "type": "string" + }, "aggregate_functions": { "APPROX_COUNT_DISTINCT": { "result_type": { @@ -1359,6 +1401,9 @@ expression: result } }, "ntext": { + "representation": { + "type": "string" + }, "aggregate_functions": { "COUNT_BIG": { "result_type": { @@ -1391,6 +1436,9 @@ expression: result } }, "numeric": { + "representation": { + "type": "bigdecimal" + }, "aggregate_functions": { "APPROX_COUNT_DISTINCT": { "result_type": { @@ -1492,6 +1540,9 @@ expression: result } }, "nvarchar": { + "representation": { + "type": "string" + }, "aggregate_functions": { "APPROX_COUNT_DISTINCT": { "result_type": { @@ -1571,6 +1622,9 @@ expression: result } }, "real": { + "representation": { + "type": "float32" + }, "aggregate_functions": { "APPROX_COUNT_DISTINCT": { "result_type": { @@ -1684,6 +1738,9 @@ expression: result } }, "smalldatetime": { + "representation": { + "type": "timestamp" + }, "aggregate_functions": { "APPROX_COUNT_DISTINCT": { "result_type": { @@ -1749,6 +1806,9 @@ expression: result } }, "smallint": { + "representation": { + "type": "int16" + }, "aggregate_functions": { "APPROX_COUNT_DISTINCT": { "result_type": { @@ -1862,6 +1922,9 @@ expression: result } }, "smallmoney": { + "representation": { + "type": "bigdecimal" + }, "aggregate_functions": { "APPROX_COUNT_DISTINCT": { "result_type": { @@ -2099,6 +2162,9 @@ expression: result } }, "text": { + "representation": { + "type": "string" + }, "aggregate_functions": { "COUNT_BIG": { "result_type": { @@ -2143,6 +2209,9 @@ expression: result } }, "time": { + "representation": { + "type": "string" + }, "aggregate_functions": { "APPROX_COUNT_DISTINCT": { "result_type": { @@ -2208,6 +2277,9 @@ expression: result } }, "timestamp": { + "representation": { + "type": "string" + }, "aggregate_functions": { "APPROX_COUNT_DISTINCT": { "result_type": { @@ -2273,6 +2345,9 @@ expression: result } }, "tinyint": { + "representation": { + "type": "int16" + }, "aggregate_functions": { "APPROX_COUNT_DISTINCT": { "result_type": { @@ -2386,6 +2461,9 @@ expression: result } }, "uniqueidentifier": { + "representation": { + "type": "uuid" + }, "aggregate_functions": { "APPROX_COUNT_DISTINCT": { "result_type": { @@ -2463,6 +2541,9 @@ expression: result } }, "varbinary": { + "representation": { + "type": "string" + }, "aggregate_functions": { "APPROX_COUNT_DISTINCT": { "result_type": { @@ -2528,6 +2609,9 @@ expression: result } }, "varchar": { + "representation": { + "type": "string" + }, "aggregate_functions": { "APPROX_COUNT_DISTINCT": { "result_type": { @@ -2619,6 +2703,9 @@ expression: result } }, "xml": { + "representation": { + "type": "string" + }, "aggregate_functions": { "APPROX_COUNT_DISTINCT": { "result_type": { diff --git a/crates/query-engine/sql/src/sql/helpers.rs b/crates/query-engine/sql/src/sql/helpers.rs index fc3eea94..4e03dce2 100644 --- a/crates/query-engine/sql/src/sql/helpers.rs +++ b/crates/query-engine/sql/src/sql/helpers.rs @@ -50,12 +50,28 @@ pub fn make_column( table: TableReference, name: ColumnName, alias: ColumnAlias, + column_type: &ScalarType, ) -> (ColumnAlias, Expression) { ( alias, - Expression::ColumnReference(ColumnReference::TableColumn { table, name }), + cast_column( + column_type, + Expression::ColumnReference(ColumnReference::TableColumn { table, name }), + ), ) } + +/// Some types need to be stringified to match representations in `ndc-spec`, this function wraps those types in cast so they becomes strings in JSON. +pub fn cast_column(column_type: &ScalarType, expression: Expression) -> Expression { + match column_type.0.as_str() { + "bigint" | "decimal" | "numeric" | "money" | "smallmoney" => Expression::Cast { + expression: Box::new(expression), + r#type: ScalarType("varchar".to_string()), + }, + _ => expression, + } +} + /// Create column aliases using this function so we build everything in one place. pub fn make_column_alias(name: String) -> ColumnAlias { ColumnAlias { name } diff --git a/crates/query-engine/translation/src/translation/helpers.rs b/crates/query-engine/translation/src/translation/helpers.rs index a25fd89e..bdc2d44a 100644 --- a/crates/query-engine/translation/src/translation/helpers.rs +++ b/crates/query-engine/translation/src/translation/helpers.rs @@ -12,7 +12,7 @@ use query_engine_sql::sql; #[derive(Debug)] /// Static information from the query and metadata. pub struct Env<'a> { - metadata: &'a metadata::Metadata, + pub metadata: &'a metadata::Metadata, relationships: BTreeMap, } diff --git a/crates/query-engine/translation/src/translation/query/aggregates.rs b/crates/query-engine/translation/src/translation/query/aggregates.rs index f7230f8e..18e4fce3 100644 --- a/crates/query-engine/translation/src/translation/query/aggregates.rs +++ b/crates/query-engine/translation/src/translation/query/aggregates.rs @@ -2,13 +2,18 @@ use indexmap::IndexMap; -use crate::translation::error::Error; -use query_engine_sql::sql; +use crate::translation::{ + error::Error, + helpers::{CollectionOrProcedureInfo, Env}, +}; +use query_engine_sql::sql::{self, ast::ScalarType, helpers::cast_column}; /// Translate any aggregates we should include in the query into our SQL AST. pub fn translate( table: &sql::ast::TableReference, aggregates: &IndexMap, + collection_info: &CollectionOrProcedureInfo, + env: &Env, ) -> Result, Error> { aggregates .into_iter() @@ -45,18 +50,49 @@ pub fn translate( column: sql::helpers::make_column_alias(column.to_string()), }, ); - match function.as_str() { - "SUM" | "AVG" => sql::ast::Expression::FunctionCall { - function: sql::ast::Function::Unknown(function.to_string()), - args: vec![sql::ast::Expression::Cast { - expression: Box::new(column_ref_expression), - r#type: sql::ast::ScalarType("BIGINT".to_string()), - }], + let column_type = match collection_info { + CollectionOrProcedureInfo::Collection(ci) => match ci { + crate::translation::helpers::CollectionInfo::Table { info, .. } => { + info.columns.get(column.as_str()).map(|c| c.r#type.clone()) + } + crate::translation::helpers::CollectionInfo::NativeQuery { + info, + .. + } => info.columns.get(column.as_str()).map(|c| c.r#type.clone()), + _ => None, }, - _ => sql::ast::Expression::FunctionCall { - function: sql::ast::Function::Unknown(function.to_string()), - args: vec![column_ref_expression], + CollectionOrProcedureInfo::Procedure(_) => None, + }; + let function_type = column_type.and_then(|column_type| { + env.metadata + .aggregate_functions + .0 + .get(&column_type) + .and_then(|functions| functions.get(function.as_str())) + }); + match function.as_str() { + "SUM" | "AVG" => sql::ast::Expression::Cast { + expression: Box::new(sql::ast::Expression::FunctionCall { + function: sql::ast::Function::Unknown(function.to_string()), + args: vec![sql::ast::Expression::Cast { + expression: Box::new(column_ref_expression), + r#type: sql::ast::ScalarType("BIGINT".to_string()), + }], + }), + r#type: sql::ast::ScalarType("varchar".to_string()), }, + _ => { + let expression = sql::ast::Expression::FunctionCall { + function: sql::ast::Function::Unknown(function.to_string()), + args: vec![column_ref_expression], + }; + match function_type.map(|f| f.return_type.clone()) { + Some(scalar_type) => { + cast_column(&ScalarType(scalar_type.0), expression) + } + None => expression, + } + } } } ndc_models::Aggregate::StarCount {} => { diff --git a/crates/query-engine/translation/src/translation/query/mod.rs b/crates/query-engine/translation/src/translation/query/mod.rs index 8345d7c7..c8a6091d 100644 --- a/crates/query-engine/translation/src/translation/query/mod.rs +++ b/crates/query-engine/translation/src/translation/query/mod.rs @@ -89,8 +89,14 @@ pub fn translate_query( )?; // translate aggregate select. if there are no fields, make this a None - let aggregate_select = - root::translate_aggregate_query(env, state, current_table, from_clause, query)?; + let aggregate_select = root::translate_aggregate_query( + env, + state, + collection_info, + current_table, + from_clause, + query, + )?; match (row_select, aggregate_select) { (Some(rows), None) => Ok(sql::helpers::SelectSet::Rows(rows)), diff --git a/crates/query-engine/translation/src/translation/query/root.rs b/crates/query-engine/translation/src/translation/query/root.rs index 320d85d7..c0aa4629 100644 --- a/crates/query-engine/translation/src/translation/query/root.rs +++ b/crates/query-engine/translation/src/translation/query/root.rs @@ -11,6 +11,7 @@ use crate::translation::helpers::CollectionOrProcedureInfo; use crate::translation::helpers::{ CollectionInfo, Env, RootAndCurrentTables, State, TableNameAndReference, }; +use query_engine_sql::sql::ast::ScalarType; use super::relationships; use super::sorting; @@ -20,6 +21,7 @@ use query_engine_sql::sql; pub fn translate_aggregate_query( env: &Env, state: &mut State, + collection_info: &CollectionOrProcedureInfo, current_table: &TableNameAndReference, from_clause: &sql::ast::From, query: &ndc_models::Query, @@ -28,8 +30,12 @@ pub fn translate_aggregate_query( None => Ok(None), Some(aggregate_fields) => { // create all aggregate columns - let aggregate_columns = - aggregates::translate(¤t_table.reference, aggregate_fields)?; + let aggregate_columns = aggregates::translate( + ¤t_table.reference, + aggregate_fields, + collection_info, + env, + )?; // create the select clause and the joins, order by, where clauses. // We don't add the limit afterwards. @@ -189,6 +195,7 @@ pub fn translate_rows_query( current_table.reference.clone(), column_info.name, sql::helpers::make_column_alias(alias.to_string()), + &ScalarType(column_info.r#type.0), )) } ndc_models::Field::Relationship {