Skip to content

Commit 9c4aae9

Browse files
JasonShinJasonShin
committed
refactor
1 parent 03e6d5f commit 9c4aae9

5 files changed

Lines changed: 338 additions & 335 deletions

File tree

‎src/ts_generator/sql_parser/expressions/function_handlers/json_functions.rs‎

Lines changed: 203 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,203 @@
1+
use crate::common::lazy::DB_SCHEMA;
2+
use crate::core::connection::DBConn;
3+
use crate::ts_generator::errors::TsGeneratorError;
4+
use crate::ts_generator::sql_parser::expressions::translate_data_type::translate_value;
5+
use crate::ts_generator::sql_parser::expressions::translate_table_with_joins::translate_table_from_expr;
6+
use crate::ts_generator::sql_parser::quoted_strings::DisplayIndent;
7+
use crate::ts_generator::types::ts_query::{TsFieldType, TsQuery};
8+
use sqlparser::ast::{Expr, FunctionArg, FunctionArgExpr, TableWithJoins, Value};
9+
10+
/// Extract key name from a function argument (should be a string literal)
11+
fn extract_key_name(arg: &FunctionArg) -> Option<String> {
12+
match arg {
13+
FunctionArg::Unnamed(FunctionArgExpr::Expr(Expr::Value(val))) => match &val.value {
14+
Value::SingleQuotedString(s) | Value::DoubleQuotedString(s) => Some(s.clone()),
15+
_ => None,
16+
},
17+
_ => None,
18+
}
19+
}
20+
21+
/// Extract expression from a function argument
22+
fn extract_expr_from_arg(arg: &FunctionArg) -> Option<&Expr> {
23+
match arg {
24+
FunctionArg::Unnamed(FunctionArgExpr::Expr(expr)) => Some(expr),
25+
FunctionArg::Named {
26+
arg: FunctionArgExpr::Expr(expr),
27+
..
28+
} => Some(expr),
29+
_ => None,
30+
}
31+
}
32+
33+
/// Infer the TypeScript type from an SQL expression
34+
pub async fn infer_type_from_expr(
35+
expr: &Expr,
36+
single_table_name: &Option<&str>,
37+
table_with_joins: &Option<Vec<TableWithJoins>>,
38+
db_conn: &DBConn,
39+
) -> Option<(TsFieldType, bool)> {
40+
match expr {
41+
Expr::Identifier(ident) => {
42+
let column_name = DisplayIndent(ident).to_string();
43+
if let Some(table_name) = single_table_name {
44+
let table_details = DB_SCHEMA.lock().await.fetch_table(&vec![table_name], db_conn).await;
45+
46+
if let Some(table_details) = table_details {
47+
if let Some(field) = table_details.get(&column_name) {
48+
Some((field.field_type.to_owned(), field.is_nullable))
49+
} else {
50+
Some((TsFieldType::Any, false))
51+
}
52+
} else {
53+
Some((TsFieldType::Any, false))
54+
}
55+
} else {
56+
Some((TsFieldType::Any, false))
57+
}
58+
}
59+
Expr::CompoundIdentifier(idents) if idents.len() == 2 => {
60+
let column_name = DisplayIndent(&idents[1]).to_string();
61+
if let Ok(table_name) = translate_table_from_expr(table_with_joins, expr) {
62+
let table_details = DB_SCHEMA
63+
.lock()
64+
.await
65+
.fetch_table(&vec![table_name.as_str()], db_conn)
66+
.await;
67+
68+
if let Some(table_details) = table_details {
69+
if let Some(field) = table_details.get(&column_name) {
70+
Some((field.field_type.to_owned(), field.is_nullable))
71+
} else {
72+
Some((TsFieldType::Any, false))
73+
}
74+
} else {
75+
Some((TsFieldType::Any, false))
76+
}
77+
} else {
78+
Some((TsFieldType::Any, false))
79+
}
80+
}
81+
Expr::Value(val) => {
82+
if let Some(ts_field_type) = translate_value(&val.value) {
83+
Some((ts_field_type, false))
84+
} else {
85+
Some((TsFieldType::Any, false))
86+
}
87+
}
88+
_ => Some((TsFieldType::Any, false)),
89+
}
90+
}
91+
92+
/// Process key-value pairs from JSON build object arguments
93+
pub async fn process_json_build_object_args(
94+
args: &[FunctionArg],
95+
single_table_name: &Option<&str>,
96+
table_with_joins: &Option<Vec<TableWithJoins>>,
97+
db_conn: &DBConn,
98+
) -> Option<Vec<(String, TsFieldType, bool)>> {
99+
if args.len() % 2 != 0 {
100+
// Invalid number of arguments
101+
return None;
102+
}
103+
104+
let mut object_fields = vec![];
105+
106+
// Process key-value pairs
107+
for i in (0..args.len()).step_by(2) {
108+
let key_arg = &args[i];
109+
let value_arg = &args[i + 1];
110+
111+
// Extract key name
112+
let key_name = extract_key_name(key_arg)?;
113+
114+
// Extract value expression
115+
let value_expr = extract_expr_from_arg(value_arg)?;
116+
117+
// Infer value type
118+
let (value_type, is_nullable) = infer_type_from_expr(value_expr, single_table_name, table_with_joins, db_conn).await?;
119+
120+
object_fields.push((key_name, value_type, is_nullable));
121+
}
122+
123+
Some(object_fields)
124+
}
125+
126+
/// Handle JSON build functions (jsonb_build_object, json_build_object, etc.)
127+
pub async fn handle_json_build_function(
128+
function_name: &str,
129+
args: &[FunctionArg],
130+
single_table_name: &Option<&str>,
131+
table_with_joins: &Option<Vec<TableWithJoins>>,
132+
db_conn: &DBConn,
133+
alias: &str,
134+
ts_query: &mut TsQuery,
135+
is_selection: bool,
136+
expr_for_logging: Option<&str>,
137+
) -> Result<(), TsGeneratorError> {
138+
let expr_log = expr_for_logging.unwrap_or("");
139+
140+
// Handle jsonb_build_object / json_build_object
141+
if function_name.to_uppercase() == "JSONB_BUILD_OBJECT" || function_name.to_uppercase() == "JSON_BUILD_OBJECT" {
142+
if let Some(object_fields) = process_json_build_object_args(args, single_table_name, table_with_joins, db_conn).await
143+
{
144+
let object_type = TsFieldType::StructuredObject(object_fields);
145+
return ts_query.insert_result(Some(alias), &[object_type], is_selection, false, expr_log);
146+
}
147+
}
148+
149+
// For other build functions or on failure, return Any
150+
ts_query.insert_result(Some(alias), &[TsFieldType::Any], is_selection, false, expr_log)
151+
}
152+
153+
/// Handle JSON aggregation functions (jsonb_agg, json_agg, etc.)
154+
pub async fn handle_json_agg_function(
155+
args: &[FunctionArg],
156+
single_table_name: &Option<&str>,
157+
table_with_joins: &Option<Vec<TableWithJoins>>,
158+
db_conn: &DBConn,
159+
alias: &str,
160+
ts_query: &mut TsQuery,
161+
is_selection: bool,
162+
expr_for_logging: Option<&str>,
163+
) -> Result<(), TsGeneratorError> {
164+
use super::super::functions::is_json_build_function;
165+
use sqlparser::ast::FunctionArguments;
166+
167+
let expr_log = expr_for_logging.unwrap_or("");
168+
169+
// jsonb_agg typically takes a single expression
170+
if args.len() != 1 {
171+
return ts_query.insert_result(Some(alias), &[TsFieldType::Any], is_selection, false, expr_log);
172+
}
173+
174+
let arg_expr = extract_expr_from_arg(&args[0]);
175+
176+
if let Some(arg_expr) = arg_expr {
177+
// Check if the argument is a jsonb_build_object function
178+
if let Expr::Function(inner_func) = arg_expr {
179+
let inner_func_name = inner_func.name.to_string();
180+
if is_json_build_function(inner_func_name.as_str()) {
181+
// Extract arguments from the inner function
182+
let inner_args = match &inner_func.args {
183+
FunctionArguments::List(arg_list) => &arg_list.args,
184+
_ => {
185+
return ts_query.insert_result(Some(alias), &[TsFieldType::Any], is_selection, false, expr_log);
186+
}
187+
};
188+
189+
// Process the inner jsonb_build_object
190+
if let Some(object_fields) =
191+
process_json_build_object_args(inner_args, single_table_name, table_with_joins, db_conn).await
192+
{
193+
let object_type = TsFieldType::StructuredObject(object_fields);
194+
let array_type = TsFieldType::Array(Box::new(object_type));
195+
return ts_query.insert_result(Some(alias), &[array_type], is_selection, false, expr_log);
196+
}
197+
}
198+
}
199+
}
200+
201+
// If we can't infer the type, return Any
202+
ts_query.insert_result(Some(alias), &[TsFieldType::Any], is_selection, false, expr_log)
203+
}

‎src/ts_generator/sql_parser/expressions/function_handlers/mod.rs‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,2 @@
1+
pub mod json_functions;
2+
pub mod polymorphic_functions;

‎src/ts_generator/sql_parser/expressions/function_handlers/polymorphic_functions.rs‎

Lines changed: 93 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,93 @@
1+
use sqlparser::ast::{Expr, Function, FunctionArg, FunctionArgExpr, FunctionArguments, TableWithJoins};
2+
use crate::common::lazy::DB_SCHEMA;
3+
use crate::ts_generator::errors::TsGeneratorError;
4+
use crate::ts_generator::sql_parser::expressions::translate_data_type::translate_value;
5+
use crate::ts_generator::sql_parser::expressions::translate_table_with_joins::translate_table_from_expr;
6+
use crate::ts_generator::sql_parser::quoted_strings::DisplayIndent;
7+
use crate::ts_generator::types::ts_query::{TsFieldType, TsQuery};
8+
9+
pub async fn handle_polymorphic_functions(
10+
ts_query: &mut TsQuery,
11+
single_table_name: &Option<&str>,
12+
table_with_joins: &Option<Vec<TableWithJoins>>,
13+
alias: &str,
14+
is_selection: bool,
15+
expr_for_logging: &str,
16+
func_obj: &Function,
17+
db_conn: &crate::core::connection::DBConn,
18+
) -> Result<(), TsGeneratorError> {
19+
// In sqlparser 0.59.0, args is a FunctionArguments enum
20+
// Extract the first argument from the appropriate variant
21+
let first_arg = match &func_obj.args {
22+
FunctionArguments::List(arg_list) => arg_list.args.first(),
23+
FunctionArguments::None => None,
24+
FunctionArguments::Subquery(_) => None, // Can't infer type from subquery easily
25+
};
26+
27+
if let Some(first_arg) = first_arg {
28+
let first_expr = match first_arg {
29+
FunctionArg::Unnamed(FunctionArgExpr::Expr(expr)) => Some(expr),
30+
FunctionArg::Named {
31+
arg: FunctionArgExpr::Expr(expr),
32+
..
33+
} => Some(expr),
34+
_ => None,
35+
};
36+
37+
if let Some(arg_expr) = first_expr {
38+
// Try to infer type from the first argument
39+
match arg_expr {
40+
Expr::Identifier(ident) => {
41+
let column_name = DisplayIndent(ident).to_string();
42+
if let Some(table_name) = single_table_name {
43+
let table_details = &DB_SCHEMA.lock().await.fetch_table(&vec![table_name], db_conn).await;
44+
45+
if let Some(table_details) = table_details {
46+
if let Some(field) = table_details.get(&column_name) {
47+
return ts_query.insert_result(
48+
Some(alias),
49+
&[field.field_type.to_owned()],
50+
is_selection,
51+
false, // IFNULL/COALESCE removes nullability
52+
expr_for_logging,
53+
);
54+
}
55+
}
56+
}
57+
}
58+
Expr::CompoundIdentifier(idents) if idents.len() == 2 => {
59+
let column_name = DisplayIndent(&idents[1]).to_string();
60+
if let Ok(table_name) = translate_table_from_expr(table_with_joins, arg_expr) {
61+
let table_details = &DB_SCHEMA
62+
.lock()
63+
.await
64+
.fetch_table(&vec![table_name.as_str()], db_conn)
65+
.await;
66+
67+
if let Some(table_details) = table_details {
68+
if let Some(field) = table_details.get(&column_name) {
69+
return ts_query.insert_result(
70+
Some(alias),
71+
&[field.field_type.to_owned()],
72+
is_selection,
73+
false, // IFNULL/COALESCE removes nullability
74+
expr_for_logging,
75+
);
76+
}
77+
}
78+
}
79+
}
80+
Expr::Value(val) => {
81+
// If first arg is a literal value, infer from that
82+
if let Some(ts_field_type) = translate_value(&val.value) {
83+
return ts_query.insert_result(Some(alias), &[ts_field_type], is_selection, false, expr_for_logging);
84+
}
85+
}
86+
_ => {}
87+
}
88+
}
89+
}
90+
91+
// Fallback to Any if we couldn't infer the type
92+
ts_query.insert_result(Some(alias), &[TsFieldType::Any], is_selection, false, expr_for_logging)
93+
}

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

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -4,6 +4,8 @@ pub mod translate_expr;
44
pub mod translate_table_with_joins;
55
pub mod translate_wildcard_expr;
66

7+
pub mod function_handlers;
8+
79
#[cfg(test)]
810
#[path = "./functions.test.rs"]
911
mod functions_test;

0 commit comments

Comments
 (0)