Skip to content

Commit a5693d6

Browse files
tdraierfrankaloia
tdraier
authored and
frankaloia
committed
[core] Include query in error (#11400)
* Include query in error * merge * Fix duplicate * fixed test
1 parent 103b1f0 commit a5693d6

File tree

8 files changed

+81
-63
lines changed

8 files changed

+81
-63
lines changed

‎core/bin/core_api.rs

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -3718,7 +3718,7 @@ async fn databases_query_run(
37183718
"The query returned too many rows",
37193719
None,
37203720
),
3721-
Err(QueryDatabaseError::ExecutionError(s)) => {
3721+
Err(QueryDatabaseError::ExecutionError(s, _)) => {
37223722
error_response(StatusCode::BAD_REQUEST, "query_execution_error", &s, None)
37233723
}
37243724
Err(e) => error_response(

‎core/src/blocks/database.rs

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -116,9 +116,10 @@ impl Block for Database {
116116
}),
117117
Err(e) => match &e {
118118
// If the actual query failed, we don't fail the block, instead we return a block result with an error inside.
119-
QueryDatabaseError::ExecutionError(s) => Ok(BlockResult {
119+
QueryDatabaseError::ExecutionError(s, q) => Ok(BlockResult {
120120
value: json!({
121121
"error": s,
122+
"query": q,
122123
}),
123124
meta: None,
124125
}),

‎core/src/databases/database.rs

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -25,7 +25,7 @@ pub enum QueryDatabaseError {
2525
#[error("Result is too large: {0}")]
2626
ResultTooLarge(String),
2727
#[error("Query execution error: {0}")]
28-
ExecutionError(String),
28+
ExecutionError(String, Option<String>),
2929
}
3030

3131
#[derive(Serialize)]

‎core/src/databases/remote_databases/bigquery.rs

Lines changed: 15 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -261,7 +261,10 @@ impl BigQueryRemoteDatabase {
261261
ResponseError {
262262
error: NestedResponseError { message, code, .. },
263263
},
264-
} => QueryDatabaseError::ExecutionError(format!("{} (code={})", message, code)),
264+
} => QueryDatabaseError::ExecutionError(
265+
format!("{} (code={})", message, code),
266+
Some(query.to_string()),
267+
),
265268
_ => QueryDatabaseError::GenericError(anyhow!("Error inserting job: {}", e)),
266269
})?;
267270

@@ -312,9 +315,10 @@ impl RemoteDatabase for BigQueryRemoteDatabase {
312315
let plan = self.get_query_plan(query).await?;
313316

314317
if !plan.is_select_query {
315-
Err(QueryDatabaseError::ExecutionError(format!(
316-
"Query is not a SELECT query"
317-
)))?
318+
Err(QueryDatabaseError::ExecutionError(
319+
format!("Query is not a SELECT query"),
320+
Some(query.to_string()),
321+
))?
318322
}
319323

320324
let used_tables: HashSet<&str> = plan
@@ -331,10 +335,13 @@ impl RemoteDatabase for BigQueryRemoteDatabase {
331335
.collect::<Vec<_>>();
332336

333337
if !used_forbidden_tables.is_empty() {
334-
Err(QueryDatabaseError::ExecutionError(format!(
335-
"Query uses tables that are not allowed: {}",
336-
used_forbidden_tables.join(", ")
337-
)))?
338+
Err(QueryDatabaseError::ExecutionError(
339+
format!(
340+
"Query uses tables that are not allowed: {}",
341+
used_forbidden_tables.join(", ")
342+
),
343+
Some(query.to_string()),
344+
))?
338345
}
339346

340347
self.execute_query(query).await

‎core/src/databases/remote_databases/salesforce/salesforce.rs

Lines changed: 35 additions & 29 deletions
Original file line numberDiff line numberDiff line change
@@ -222,10 +222,10 @@ impl SalesforceRemoteDatabase {
222222
QueryDatabaseError::GenericError(anyhow!("Error getting response text: {}", e))
223223
})?;
224224

225-
return Err(QueryDatabaseError::ExecutionError(format!(
226-
"Query failed: {}",
227-
error_text
228-
)));
225+
return Err(QueryDatabaseError::ExecutionError(
226+
format!("Query failed: {}", error_text),
227+
Some(query.to_string()),
228+
));
229229
}
230230

231231
let response_text = response.text().await.map_err(|e| {
@@ -335,10 +335,10 @@ impl SalesforceRemoteDatabase {
335335
QueryDatabaseError::GenericError(anyhow!("Error getting response text: {}", e))
336336
})?;
337337

338-
return Err(QueryDatabaseError::ExecutionError(format!(
339-
"Failed to describe object {}: {}",
340-
sobject, error_text
341-
)));
338+
return Err(QueryDatabaseError::ExecutionError(
339+
format!("Failed to describe object {}: {}", sobject, error_text),
340+
None,
341+
));
342342
}
343343

344344
let response_text = response.text().await.map_err(|e| {
@@ -537,10 +537,13 @@ impl SalesforceRemoteDatabase {
537537

538538
// Check if primary object is allowed, same as for relationships
539539
if !allowed_objects.contains(&primary_object.to_lowercase()) {
540-
return Err(QueryDatabaseError::ExecutionError(format!(
541-
"Primary object '{}' is not allowed. You don't have access to this object.",
542-
primary_object
543-
)));
540+
return Err(QueryDatabaseError::ExecutionError(
541+
format!(
542+
"Primary object '{}' is not allowed. You don't have access to this object.",
543+
primary_object
544+
),
545+
None,
546+
));
544547
}
545548

546549
// Check for objects that aren't directly allowed
@@ -611,14 +614,17 @@ impl SalesforceRemoteDatabase {
611614

612615
// If we still have unknown objects, they're not allowed
613616
if !unknown_objects.is_empty() {
614-
return Err(QueryDatabaseError::ExecutionError(format!(
615-
"Query uses tables/relationships that are not allowed: {}",
616-
unknown_objects
617-
.iter()
618-
.map(|s| s.as_str())
619-
.collect::<Vec<_>>()
620-
.join(", ")
621-
)));
617+
return Err(QueryDatabaseError::ExecutionError(
618+
format!(
619+
"Query uses tables/relationships that are not allowed: {}",
620+
unknown_objects
621+
.iter()
622+
.map(|s| s.as_str())
623+
.collect::<Vec<_>>()
624+
.join(", ")
625+
),
626+
None,
627+
));
622628
}
623629

624630
Ok(())
@@ -638,15 +644,15 @@ impl RemoteDatabase for SalesforceRemoteDatabase {
638644
) -> Result<(Vec<QueryResult>, TableSchema, String), QueryDatabaseError> {
639645
// Parse the JSON query
640646
let parsed_query = serde_json::from_str::<StructuredQuery>(query).map_err(|e| {
641-
QueryDatabaseError::ExecutionError(format!("Failed to parse JSON query: {}", e))
647+
QueryDatabaseError::ExecutionError(format!("Failed to parse JSON query: {}", e), None)
642648
})?;
643649

644650
// Validate the structured query
645651
if let Err(e) = parsed_query.validate() {
646-
return Err(QueryDatabaseError::ExecutionError(format!(
647-
"Invalid structured query: {}",
648-
e
649-
)));
652+
return Err(QueryDatabaseError::ExecutionError(
653+
format!("Invalid structured query: {}", e),
654+
None,
655+
));
650656
}
651657

652658
// Extract all objects referenced in the query
@@ -676,10 +682,10 @@ impl RemoteDatabase for SalesforceRemoteDatabase {
676682

677683
// Convert the structured query to SOQL
678684
let soql_query = convert_to_soql(&parsed_query).map_err(|e| {
679-
QueryDatabaseError::ExecutionError(format!(
680-
"Error converting JSON query to SOQL: {}",
681-
e
682-
))
685+
QueryDatabaseError::ExecutionError(
686+
format!("Error converting JSON query to SOQL: {}", e),
687+
None,
688+
)
683689
})?;
684690

685691
// Execute the SOQL query

‎core/src/databases/remote_databases/salesforce/sandbox/convert.rs

Lines changed: 2 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -29,7 +29,7 @@ pub fn convert_to_soql(query: &StructuredQuery) -> Result<String, SoqlError> {
2929

3030
// Add regular fields
3131
if !query.fields.is_empty() {
32-
select_parts.push(query.fields.join(", "));
32+
select_parts.extend(query.fields.iter().cloned());
3333
}
3434

3535
// Add parent fields if present
@@ -71,12 +71,7 @@ pub fn convert_to_soql(query: &StructuredQuery) -> Result<String, SoqlError> {
7171
if let Some(fields) = group_by_fields {
7272
for field in fields {
7373
// If the field is not in the fields or in the aggregates, add it to the select parts
74-
if !query.fields.contains(field)
75-
&& !query
76-
.aggregates
77-
.iter()
78-
.any(|agg| agg.field.as_str() == field)
79-
{
74+
if !select_parts.contains(field) {
8075
select_parts.push(field.clone());
8176
}
8277
}

‎core/src/databases/remote_databases/snowflake.rs

Lines changed: 24 additions & 15 deletions
Original file line numberDiff line numberDiff line change
@@ -225,13 +225,16 @@ impl SnowflakeRemoteDatabase {
225225
) -> Result<(Vec<QueryResult>, TableSchema, String), QueryDatabaseError> {
226226
let executor = match session.execute(query).await {
227227
Ok(executor) => Ok(executor),
228-
Err(snowflake_connector_rs::Error::TimedOut) => Err(
229-
QueryDatabaseError::ExecutionError("Query execution timed out".to_string()),
230-
),
231-
Err(e) => Err(QueryDatabaseError::ExecutionError(format!(
232-
"Error executing query: {}",
233-
e
234-
))),
228+
Err(snowflake_connector_rs::Error::TimedOut) => {
229+
Err(QueryDatabaseError::ExecutionError(
230+
"Query execution timed out".to_string(),
231+
Some(query.to_string()),
232+
))
233+
}
234+
Err(e) => Err(QueryDatabaseError::ExecutionError(
235+
format!("Error executing query: {}", e),
236+
Some(query.to_string()),
237+
)),
235238
}?;
236239

237240
let mut query_result_rows: usize = 0;
@@ -310,10 +313,13 @@ impl SnowflakeRemoteDatabase {
310313
.collect::<Vec<_>>();
311314

312315
if !used_forbidden_tables.is_empty() {
313-
Err(QueryDatabaseError::ExecutionError(format!(
314-
"Query uses tables that are not allowed: {}",
315-
used_forbidden_tables.join(", ")
316-
)))?
316+
Err(QueryDatabaseError::ExecutionError(
317+
format!(
318+
"Query uses tables that are not allowed: {}",
319+
used_forbidden_tables.join(", ")
320+
),
321+
Some(query.to_string()),
322+
))?
317323
}
318324

319325
let used_forbidden_operations = plan
@@ -331,10 +337,13 @@ impl SnowflakeRemoteDatabase {
331337
.collect::<Vec<_>>();
332338

333339
if !used_forbidden_operations.is_empty() {
334-
Err(QueryDatabaseError::ExecutionError(format!(
335-
"Query contains forbidden operations: {}",
336-
used_forbidden_operations.join(", ")
337-
)))?
340+
Err(QueryDatabaseError::ExecutionError(
341+
format!(
342+
"Query contains forbidden operations: {}",
343+
used_forbidden_operations.join(", ")
344+
),
345+
Some(query.to_string()),
346+
))?
338347
}
339348

340349
Ok(())

‎core/src/databases/transient_database.rs

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -21,7 +21,7 @@ impl From<SqliteWorkerError> for QueryDatabaseError {
2121
match &e {
2222
SqliteWorkerError::TooManyResultRows => QueryDatabaseError::TooManyResultRows,
2323
SqliteWorkerError::QueryExecutionError(msg) => {
24-
QueryDatabaseError::ExecutionError(msg.clone())
24+
QueryDatabaseError::ExecutionError(msg.clone(), None)
2525
}
2626
_ => QueryDatabaseError::GenericError(e.into()),
2727
}

0 commit comments

Comments
 (0)