Skip to content

Commit 96d459a

Browse files
JasonShinJasonShin
committed
Support JSON
1 parent a1d0241 commit 96d459a

8 files changed

Lines changed: 334 additions & 53 deletions

File tree

‎src/ts_generator/sql_parser/expressions/functions.rs‎

Lines changed: 24 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -170,3 +170,27 @@ pub static TYPE_POLYMORPHIC_FUNCTIONS: &[&str] = &[
170170
pub fn is_type_polymorphic_function(func_name: &str) -> bool {
171171
TYPE_POLYMORPHIC_FUNCTIONS.contains(&func_name.to_uppercase().as_str())
172172
}
173+
174+
// JSON/JSONB functions that build objects/arrays
175+
pub static JSON_BUILD_FUNCTIONS: &[&str] = &[
176+
"JSONB_BUILD_OBJECT",
177+
"JSON_BUILD_OBJECT",
178+
"JSONB_BUILD_ARRAY",
179+
"JSON_BUILD_ARRAY",
180+
];
181+
182+
// JSON/JSONB aggregation functions
183+
pub static JSON_AGG_FUNCTIONS: &[&str] = &[
184+
"JSONB_AGG",
185+
"JSON_AGG",
186+
"JSON_OBJECT_AGG",
187+
"JSONB_OBJECT_AGG",
188+
];
189+
190+
pub fn is_json_build_function(func_name: &str) -> bool {
191+
JSON_BUILD_FUNCTIONS.contains(&func_name.to_uppercase().as_str())
192+
}
193+
194+
pub fn is_json_agg_function(func_name: &str) -> bool {
195+
JSON_AGG_FUNCTIONS.contains(&func_name.to_uppercase().as_str())
196+
}

‎src/ts_generator/sql_parser/expressions/translate_expr.rs‎

Lines changed: 286 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,4 +1,4 @@
1-
use super::functions::{is_date_function, is_numeric_function, is_type_polymorphic_function};
1+
use super::functions::{is_date_function, is_json_agg_function, is_json_build_function, is_numeric_function, is_type_polymorphic_function};
22
use crate::common::lazy::DB_SCHEMA;
33
use crate::common::logger::{error, warning};
44
use crate::core::connection::DBConn;
@@ -707,6 +707,291 @@ pub async fn translate_expr(
707707
return ts_query.insert_result(Some(alias), &[TsFieldType::Any], is_selection, false, expr_for_logging);
708708
}
709709

710+
// Handle JSON build functions (jsonb_build_object, json_build_object, etc.)
711+
if is_json_build_function(function_name_str) {
712+
use sqlparser::ast::{FunctionArg, FunctionArgExpr, FunctionArguments};
713+
714+
let args = match &func_obj.args {
715+
FunctionArguments::List(arg_list) => &arg_list.args,
716+
_ => {
717+
// If no arguments or subquery, return Any
718+
return ts_query.insert_result(Some(alias), &[TsFieldType::Any], is_selection, false, expr_for_logging);
719+
}
720+
};
721+
722+
// jsonb_build_object takes key-value pairs
723+
// e.g., jsonb_build_object('id', id, 'name', name)
724+
if function_name_str.to_uppercase() == "JSONB_BUILD_OBJECT" || function_name_str.to_uppercase() == "JSON_BUILD_OBJECT" {
725+
if args.len() % 2 != 0 {
726+
// Invalid number of arguments
727+
return ts_query.insert_result(Some(alias), &[TsFieldType::Any], is_selection, false, expr_for_logging);
728+
}
729+
730+
let mut object_fields = vec![];
731+
732+
// Process key-value pairs
733+
for i in (0..args.len()).step_by(2) {
734+
let key_arg = &args[i];
735+
let value_arg = &args[i + 1];
736+
737+
// Extract key name (should be a string literal)
738+
let key_name = match key_arg {
739+
FunctionArg::Unnamed(FunctionArgExpr::Expr(Expr::Value(val))) => {
740+
match &val.value {
741+
Value::SingleQuotedString(s) | Value::DoubleQuotedString(s) => Some(s.clone()),
742+
_ => None,
743+
}
744+
}
745+
_ => None,
746+
};
747+
748+
if key_name.is_none() {
749+
// If we can't extract the key, return Any
750+
return ts_query.insert_result(Some(alias), &[TsFieldType::Any], is_selection, false, expr_for_logging);
751+
}
752+
753+
let key_name = key_name.unwrap();
754+
755+
// Extract value type from the expression
756+
let value_expr = match value_arg {
757+
FunctionArg::Unnamed(FunctionArgExpr::Expr(expr)) => Some(expr),
758+
FunctionArg::Named {
759+
arg: FunctionArgExpr::Expr(expr),
760+
..
761+
} => Some(expr),
762+
_ => None,
763+
};
764+
765+
if let Some(value_expr) = value_expr {
766+
let value_type = match value_expr {
767+
Expr::Identifier(ident) => {
768+
let column_name = DisplayIndent(ident).to_string();
769+
if let Some(table_name) = single_table_name {
770+
let table_details = &DB_SCHEMA.lock().await.fetch_table(&vec![table_name], db_conn).await;
771+
772+
if let Some(table_details) = table_details {
773+
if let Some(field) = table_details.get(&column_name) {
774+
Some((field.field_type.to_owned(), field.is_nullable))
775+
} else {
776+
Some((TsFieldType::Any, false))
777+
}
778+
} else {
779+
Some((TsFieldType::Any, false))
780+
}
781+
} else {
782+
Some((TsFieldType::Any, false))
783+
}
784+
}
785+
Expr::CompoundIdentifier(idents) if idents.len() == 2 => {
786+
let column_name = DisplayIndent(&idents[1]).to_string();
787+
if let Ok(table_name) = translate_table_from_expr(table_with_joins, value_expr) {
788+
let table_details = &DB_SCHEMA
789+
.lock()
790+
.await
791+
.fetch_table(&vec![table_name.as_str()], db_conn)
792+
.await;
793+
794+
if let Some(table_details) = table_details {
795+
if let Some(field) = table_details.get(&column_name) {
796+
Some((field.field_type.to_owned(), field.is_nullable))
797+
} else {
798+
Some((TsFieldType::Any, false))
799+
}
800+
} else {
801+
Some((TsFieldType::Any, false))
802+
}
803+
} else {
804+
Some((TsFieldType::Any, false))
805+
}
806+
}
807+
Expr::Value(val) => {
808+
if let Some(ts_field_type) = translate_value(&val.value) {
809+
Some((ts_field_type, false))
810+
} else {
811+
Some((TsFieldType::Any, false))
812+
}
813+
}
814+
_ => Some((TsFieldType::Any, false)),
815+
};
816+
817+
if let Some((value_type, is_nullable)) = value_type {
818+
object_fields.push((key_name, value_type, is_nullable));
819+
} else {
820+
// If we can't infer the type, return Any
821+
return ts_query.insert_result(Some(alias), &[TsFieldType::Any], is_selection, false, expr_for_logging);
822+
}
823+
} else {
824+
// If we can't extract the value expression, return Any
825+
return ts_query.insert_result(Some(alias), &[TsFieldType::Any], is_selection, false, expr_for_logging);
826+
}
827+
}
828+
829+
// Build the StructuredObject type with the inferred fields
830+
let object_type = TsFieldType::StructuredObject(object_fields);
831+
return ts_query.insert_result(Some(alias), &[object_type], is_selection, false, expr_for_logging);
832+
}
833+
834+
// For build_array functions, we'd need different logic
835+
// For now, return Any
836+
return ts_query.insert_result(Some(alias), &[TsFieldType::Any], is_selection, false, expr_for_logging);
837+
}
838+
839+
// Handle JSON aggregation functions (jsonb_agg, json_agg, etc.)
840+
if is_json_agg_function(function_name_str) {
841+
use sqlparser::ast::{FunctionArg, FunctionArgExpr, FunctionArguments};
842+
843+
let args = match &func_obj.args {
844+
FunctionArguments::List(arg_list) => &arg_list.args,
845+
_ => {
846+
// If no arguments or subquery, return Any
847+
return ts_query.insert_result(Some(alias), &[TsFieldType::Any], is_selection, false, expr_for_logging);
848+
}
849+
};
850+
851+
// jsonb_agg typically takes a single expression
852+
// e.g., jsonb_agg(jsonb_build_object('id', id, 'name', name))
853+
if args.len() == 1 {
854+
let first_arg = &args[0];
855+
let arg_expr = match first_arg {
856+
FunctionArg::Unnamed(FunctionArgExpr::Expr(expr)) => Some(expr),
857+
FunctionArg::Named {
858+
arg: FunctionArgExpr::Expr(expr),
859+
..
860+
} => Some(expr),
861+
_ => None,
862+
};
863+
864+
if let Some(arg_expr) = arg_expr {
865+
// Check if the argument is a jsonb_build_object function
866+
if let Expr::Function(inner_func) = arg_expr {
867+
let inner_func_name = inner_func.name.to_string();
868+
if is_json_build_function(inner_func_name.as_str()) {
869+
// Recursively infer the type of jsonb_build_object
870+
let inner_args = match &inner_func.args {
871+
FunctionArguments::List(arg_list) => &arg_list.args,
872+
_ => {
873+
return ts_query.insert_result(Some(alias), &[TsFieldType::Any], is_selection, false, expr_for_logging);
874+
}
875+
};
876+
877+
if inner_args.len() % 2 != 0 {
878+
// Invalid number of arguments
879+
return ts_query.insert_result(Some(alias), &[TsFieldType::Any], is_selection, false, expr_for_logging);
880+
}
881+
882+
let mut object_fields = vec![];
883+
884+
// Process key-value pairs
885+
for i in (0..inner_args.len()).step_by(2) {
886+
let key_arg = &inner_args[i];
887+
let value_arg = &inner_args[i + 1];
888+
889+
// Extract key name
890+
let key_name = match key_arg {
891+
FunctionArg::Unnamed(FunctionArgExpr::Expr(Expr::Value(val))) => {
892+
match &val.value {
893+
Value::SingleQuotedString(s) | Value::DoubleQuotedString(s) => Some(s.clone()),
894+
_ => None,
895+
}
896+
}
897+
_ => None,
898+
};
899+
900+
if key_name.is_none() {
901+
return ts_query.insert_result(Some(alias), &[TsFieldType::Any], is_selection, false, expr_for_logging);
902+
}
903+
904+
let key_name = key_name.unwrap();
905+
906+
// Extract value type
907+
let value_expr = match value_arg {
908+
FunctionArg::Unnamed(FunctionArgExpr::Expr(expr)) => Some(expr),
909+
FunctionArg::Named {
910+
arg: FunctionArgExpr::Expr(expr),
911+
..
912+
} => Some(expr),
913+
_ => None,
914+
};
915+
916+
if let Some(value_expr) = value_expr {
917+
let value_type = match value_expr {
918+
Expr::Identifier(ident) => {
919+
let column_name = DisplayIndent(ident).to_string();
920+
if let Some(table_name) = single_table_name {
921+
let table_details = &DB_SCHEMA.lock().await.fetch_table(&vec![table_name], db_conn).await;
922+
923+
if let Some(table_details) = table_details {
924+
if let Some(field) = table_details.get(&column_name) {
925+
Some((field.field_type.to_owned(), field.is_nullable))
926+
} else {
927+
Some((TsFieldType::Any, false))
928+
}
929+
} else {
930+
Some((TsFieldType::Any, false))
931+
}
932+
} else {
933+
Some((TsFieldType::Any, false))
934+
}
935+
}
936+
Expr::CompoundIdentifier(idents) if idents.len() == 2 => {
937+
let column_name = DisplayIndent(&idents[1]).to_string();
938+
if let Ok(table_name) = translate_table_from_expr(table_with_joins, value_expr) {
939+
let table_details = &DB_SCHEMA
940+
.lock()
941+
.await
942+
.fetch_table(&vec![table_name.as_str()], db_conn)
943+
.await;
944+
945+
if let Some(table_details) = table_details {
946+
if let Some(field) = table_details.get(&column_name) {
947+
Some((field.field_type.to_owned(), field.is_nullable))
948+
} else {
949+
Some((TsFieldType::Any, false))
950+
}
951+
} else {
952+
Some((TsFieldType::Any, false))
953+
}
954+
} else {
955+
Some((TsFieldType::Any, false))
956+
}
957+
}
958+
Expr::Value(val) => {
959+
if let Some(ts_field_type) = translate_value(&val.value) {
960+
Some((ts_field_type, false))
961+
} else {
962+
Some((TsFieldType::Any, false))
963+
}
964+
}
965+
_ => Some((TsFieldType::Any, false)),
966+
};
967+
968+
if let Some((value_type, is_nullable)) = value_type {
969+
object_fields.push((key_name, value_type, is_nullable));
970+
} else {
971+
return ts_query.insert_result(Some(alias), &[TsFieldType::Any], is_selection, false, expr_for_logging);
972+
}
973+
} else {
974+
return ts_query.insert_result(Some(alias), &[TsFieldType::Any], is_selection, false, expr_for_logging);
975+
}
976+
}
977+
978+
// Build the Array of StructuredObject type
979+
let object_type = TsFieldType::StructuredObject(object_fields);
980+
let array_type = TsFieldType::Array(Box::new(object_type));
981+
return ts_query.insert_result(Some(alias), &[array_type], is_selection, false, expr_for_logging);
982+
}
983+
}
984+
985+
// If not jsonb_build_object, try to infer the type of the expression
986+
// For now, return Any
987+
return ts_query.insert_result(Some(alias), &[TsFieldType::Any], is_selection, false, expr_for_logging);
988+
}
989+
}
990+
991+
// If we can't infer the type, return Any
992+
return ts_query.insert_result(Some(alias), &[TsFieldType::Any], is_selection, false, expr_for_logging);
993+
}
994+
710995
// Handle other function types
711996
if is_string_function(function_name_str) {
712997
ts_query.insert_result(

‎src/ts_generator/types/ts_query.rs‎

Lines changed: 16 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -16,6 +16,8 @@ pub enum TsFieldType {
1616
Number,
1717
Boolean,
1818
Object,
19+
// Structured object with named fields: Vec<(field_name, field_type, is_nullable)>
20+
StructuredObject(Vec<(String, TsFieldType, bool)>),
1921
Date,
2022
Null,
2123
Enum(Vec<String>),
@@ -34,6 +36,20 @@ impl fmt::Display for TsFieldType {
3436
TsFieldType::Number => write!(f, "number"),
3537
TsFieldType::String => write!(f, "string"),
3638
TsFieldType::Object => write!(f, "object"),
39+
TsFieldType::StructuredObject(fields) => {
40+
let field_strings: Vec<String> = fields
41+
.iter()
42+
.map(|(field_name, field_type, is_nullable)| {
43+
let type_str = if *is_nullable {
44+
format!("{} | null", field_type)
45+
} else {
46+
field_type.to_string()
47+
};
48+
format!("{}: {}", field_name, type_str)
49+
})
50+
.collect();
51+
write!(f, "{{ {} }}", field_strings.join("; "))
52+
}
3753
TsFieldType::Date => write!(f, "Date"),
3854
TsFieldType::Any => write!(f, "any"),
3955
TsFieldType::Null => write!(f, "null"),

‎tests/demo/postgres/jsonb_operations.queries.ts‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -2,7 +2,7 @@ export type JsonbBuildObjectBasicParams = [];
22

33
export interface IJsonbBuildObjectBasicResult {
44
id: number;
5-
itemJson: any;
5+
itemJson: { id: number; name: string; rarity: string | null };
66
}
77

88
export interface IJsonbBuildObjectBasicQuery {
@@ -13,7 +13,7 @@ export interface IJsonbBuildObjectBasicQuery {
1313
export type JsonbAggregationParams = [];
1414

1515
export interface IJsonbAggregationResult {
16-
items: any;
16+
items: Array<{ id: number; name: string }>;
1717
rarity: string | null;
1818
}
1919

‎tests/demo/postgres/jsonb_operations.snapshot.ts‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -2,7 +2,7 @@ export type JsonbBuildObjectBasicParams = [];
22

33
export interface IJsonbBuildObjectBasicResult {
44
id: number;
5-
itemJson: any;
5+
itemJson: { id: number; name: string; rarity: string | null };
66
}
77

88
export interface IJsonbBuildObjectBasicQuery {
@@ -13,7 +13,7 @@ export interface IJsonbBuildObjectBasicQuery {
1313
export type JsonbAggregationParams = [];
1414

1515
export interface IJsonbAggregationResult {
16-
items: any;
16+
items: Array<{ id: number; name: string }>;
1717
rarity: string | null;
1818
}
1919

0 commit comments

Comments
 (0)