Add detailed diagnostics for validate_sql API (#1123)

Author: mwylde

crates/arroyo-api/src/lib.rs

@@ -356,6 +356,9 @@ impl IntoResponse for HttpError {
         GlobalUdf,
         GlobalUdfCollection,
         BadData,
+        SqlDiagnostic,
+        SqlSpan,
+        SqlLocation
     )),
     tags(
         (name = "ping", description = "Ping endpoint"),

crates/arroyo-api/src/pipelines.rs

@@ -28,7 +28,7 @@ use arroyo_datastream::logical::{
     ChainedLogicalOperator, LogicalNode, LogicalProgram, OperatorChain, OperatorName,
 };
 use arroyo_formats::ser::ArrowSerializer;
-use arroyo_planner::{ArroyoSchemaProvider, CompiledSql, SqlConfig};
+use arroyo_planner::{ArroyoSchemaProvider, CompiledSql, PlannerError, SqlConfig};
 use arroyo_rpc::formats::Format;
 use arroyo_rpc::grpc::rpc::compiler_grpc_client::CompilerGrpcClient;
 use arroyo_rpc::public_ids::{IdTypes, generate_id};
@@ -65,7 +65,7 @@ async fn compile_sql(
     auth_data: &AuthData,
     validate_only: bool,
     db: &DatabaseSource,
-) -> Result<CompiledSql, ErrorResp> {
+) -> Result<Result<CompiledSql, PlannerError>, ErrorResp> {
     let mut schema_provider = ArroyoSchemaProvider::new();
 
     let global_udfs = fetch_get_udfs(&db.client().await?, &auth_data.organization_id)
@@ -174,15 +174,14 @@ async fn compile_sql(
         schema_provider.add_connection_profile(profile);
     }
 
-    arroyo_planner::parse_and_get_program(
+    Ok(arroyo_planner::parse_and_get_program(
         &query,
         schema_provider,
         SqlConfig {
             default_parallelism: parallelism,
         },
     )
-    .await
-    .map_err(|err| bad_request(err.to_string()))
+    .await)
 }
 
 fn set_parallelism(program: &mut LogicalProgram, parallelism: usize) {
@@ -569,15 +568,15 @@ pub async fn validate_query(
         true,
         &state.database,
     )
-    .await
+    .await?
     {
         Ok(CompiledSql { program, .. }) => QueryValidationResult {
             graph: Some(program.try_into().map_err(log_and_map)?),
             errors: vec![],
         },
         Err(e) => QueryValidationResult {
             graph: None,
-            errors: vec![e.message],
+            errors: e.diagnostics,
         },
     };
 
@@ -660,7 +659,7 @@ async fn create_pipeline_inner(
         false,
         &state.database,
     )
-    .await?;
+    .await??;
 
     let pipeline_id = create_pipeline_int(
         pipeline_post.name,
@@ -718,7 +717,7 @@ pub async fn create_preview_pipeline(
         false,
         &state.database,
     )
-    .await?;
+    .await??;
 
     let pipeline_id = create_pipeline_int(
         format!("preview_{}", to_millis(SystemTime::now())),

crates/arroyo-api/src/rest_utils.rs

@@ -1,4 +1,5 @@
 use crate::{AuthData, cloud};
+use arroyo_planner::PlannerError;
 use arroyo_rpc::log_event;
 use axum::Json;
 use axum::extract::rejection::JsonRejection;
@@ -67,6 +68,20 @@ pub fn map_delete_err(name: &str, user: &str, error: DbError) -> ErrorResp {
     }
 }
 
+impl From<PlannerError> for ErrorResp {
+    fn from(value: PlannerError) -> Self {
+        let mut message = "Failed to plan query:".to_string();
+        for d in &value.diagnostics {
+            message.push_str(&format!("\n * {}", d.message));
+        }
+
+        ErrorResp {
+            status_code: StatusCode::BAD_REQUEST,
+            message,
+        }
+    }
+}
+
 impl From<DbError> for ErrorResp {
     fn from(value: DbError) -> Self {
         match value {

crates/arroyo-planner/src/lib.rs

@@ -64,6 +64,7 @@ use arroyo_connectors::connectors;
 use arroyo_datastream::logical::LogicalProgram;
 use arroyo_datastream::optimizers::ChainingOptimizer;
 use arroyo_operator::connector::Connection;
+use arroyo_rpc::api_types::pipelines::{SqlDiagnostic, SqlLocation, SqlSpan};
 use arroyo_rpc::df::ArroyoSchema;
 use arroyo_rpc::{TIMESTAMP_FIELD, duration_from_sql};
 use arroyo_udf_host::ParsedUdfFile;
@@ -548,17 +549,52 @@ impl Default for SqlConfig {
     }
 }
 
+#[derive(Debug)]
+pub struct PlannerError {
+    pub diagnostics: Vec<SqlDiagnostic>,
+}
+
 pub async fn parse_and_get_program(
     query: &str,
     schema_provider: ArroyoSchemaProvider,
     config: SqlConfig,
-) -> Result<CompiledSql> {
+) -> Result<CompiledSql, PlannerError> {
     let query = query.to_string();
 
     if query.trim().is_empty() {
-        return plan_err!("Query is empty");
-    }
-    parse_and_get_arrow_program(query, schema_provider, config).await
+        return Err(PlannerError {
+            diagnostics: vec![SqlDiagnostic::message("Query is empty!")],
+        });
+    }
+    parse_and_get_arrow_program(query, schema_provider, config)
+        .await
+        .map_err(|e| {
+            let diagnostics = e
+                .iter()
+                .map(|error| {
+                    let diagnostic = error.diagnostic();
+                    SqlDiagnostic {
+                        message: diagnostic
+                            .map(|diagnostic| diagnostic.message.clone())
+                            .unwrap_or_else(|| error.to_string()),
+                        span: diagnostic.and_then(|diagnostic| {
+                            diagnostic.span.map(|span| SqlSpan {
+                                start: SqlLocation {
+                                    line: span.start.line,
+                                    column: span.start.column,
+                                },
+                                end: SqlLocation {
+                                    line: span.end.line,
+                                    column: span.end.column,
+                                },
+                            })
+                        }),
+                    }
+                })
+                .collect();
+
+            PlannerError { diagnostics }
+        })
 }
 
 #[derive(Debug, Clone, PartialEq, Eq, Hash)]

crates/arroyo-planner/src/tables.rs

@@ -57,8 +57,7 @@ use datafusion::{
     },
 };
 use itertools::Itertools;
-use sqlparser::ast;
-use sqlparser::ast::TableConstraint;
+use sqlparser::ast::{self, TableConstraint};
 use std::sync::Arc;
 use std::{collections::HashMap, time::Duration};
 use tracing::warn;
@@ -140,7 +139,10 @@ fn produce_optimized_plan(
     statement: &Statement,
     schema_provider: &ArroyoSchemaProvider,
 ) -> Result<LogicalPlan> {
-    let sql_to_rel = SqlToRel::new(schema_provider);
+    let sql_to_rel = SqlToRel::new_with_options(
+        schema_provider,
+        datafusion::sql::planner::ParserOptions::new().with_collect_spans(true),
+    );
 
     let plan = sql_to_rel.sql_statement_to_plan(statement.clone())?;
 

crates/arroyo-planner/src/test/plan_tests.rs

@@ -47,16 +47,21 @@ async fn validate_query(path: &Path) {
     if fail {
         let err = result.unwrap_err();
         if let Some(error_message) = error_message {
-            let err_s = err.to_string();
-            let err: Vec<_> = err_s.split_whitespace().collect();
+            let err: Vec<_> = err
+                .diagnostics
+                .first()
+                .unwrap()
+                .message
+                .split_whitespace()
+                .collect();
             let err = err.join(" ");
             assert!(
                 err.contains(error_message),
                 "expected error message '{error_message}' not found; instead got '{err}'"
             );
         }
     } else if let Err(e) = result {
-        println!("{e}");
-        panic!("{}", e);
+        println!("{e:?}");
+        panic!("{e:?}");
     }
 }

crates/arroyo-planner/src/test/queries/virtual_bad_schema.sql

@@ -1,4 +1,4 @@
---fail=Schema error: No field named notfield. Valid fields are input.length.
+--fail='notfield' not found
 create table input (
     length JSON,
     diff INT GENERATED ALWAYS AS (notfield) STORED

crates/arroyo-planner/src/types.rs

@@ -1,15 +1,14 @@
-use std::{sync::Arc, time::SystemTime};
-
 use arrow::datatypes::{DataType, Field};
-use datafusion::common::{Result, plan_datafusion_err, plan_err};
-
 use arrow_schema::{DECIMAL_DEFAULT_SCALE, DECIMAL128_MAX_PRECISION, IntervalUnit, TimeUnit};
 use arroyo_types::ArroyoExtensionType;
-use datafusion::error::DataFusionError;
+use datafusion::common::{
+    DataFusionError, Diagnostic, Result, Span, plan_datafusion_err, plan_err,
+};
 use datafusion::sql::sqlparser::ast::{
-    ArrayElemTypeDef, DataType as SQLDataType, ExactNumberInfo, TimezoneInfo,
+    ArrayElemTypeDef, DataType as SQLDataType, ExactNumberInfo, Spanned, TimezoneInfo,
 };
 use itertools::Itertools;
+use std::{sync::Arc, time::SystemTime};
 // Pulled from DataFusion
 
 pub(crate) fn convert_data_type(
@@ -117,7 +116,13 @@ fn convert_simple_data_type(
 
             Ok(DataType::Struct(fields.into()))
         }
-        _ => return plan_err!("Unsupported SQL type {sql_type:?}"),
+        SQLDataType::Custom(name, _) => {
+            let message = format!("Unsupported SQL type '{name}'");
+            let diagnostic =
+                Diagnostic::new_error(message.clone(), Span::try_from_sqlparser_span(name.span()));
+            return plan_err!("{message}"; diagnostic = diagnostic);
+        }
+        _ => return plan_err!("Unsupported SQL type '{sql_type}'"),
     };
 
     Ok((dt?, None))

crates/arroyo-rpc/src/api_types/pipelines.rs

@@ -1,9 +1,8 @@
-use std::collections::HashMap;
-
 use crate::api_types::udfs::Udf;
 use crate::errors::ErrorDomain;
 use crate::grpc as grpc_proto;
 use serde::{Deserialize, Serialize};
+use std::collections::HashMap;
 use utoipa::ToSchema;
 
 #[derive(Serialize, Deserialize, Clone, Debug, ToSchema)]
@@ -17,7 +16,34 @@ pub struct ValidateQueryPost {
 #[serde(rename_all = "snake_case")]
 pub struct QueryValidationResult {
     pub graph: Option<PipelineGraph>,
-    pub errors: Vec<String>,
+    pub errors: Vec<SqlDiagnostic>,
+}
+
+#[derive(Serialize, Deserialize, Clone, Debug, PartialEq, Eq, ToSchema)]
+pub struct SqlLocation {
+    pub line: u64,
+    pub column: u64,
+}
+
+#[derive(Serialize, Deserialize, Clone, Debug, PartialEq, Eq, ToSchema)]
+pub struct SqlSpan {
+    pub start: SqlLocation,
+    pub end: SqlLocation,
+}
+
+#[derive(Serialize, Deserialize, Clone, Debug, PartialEq, Eq, ToSchema)]
+pub struct SqlDiagnostic {
+    pub message: String,
+    #[serde(skip_serializing_if = "Option::is_none")]
+    pub span: Option<SqlSpan>,
+}
+impl SqlDiagnostic {
+    pub fn message(message: impl Into<String>) -> Self {
+        Self {
+            message: message.into(),
+            span: None,
+        }
+    }
 }
 
 #[derive(Serialize, Deserialize, Clone, Debug, ToSchema)]

crates/arroyo/src/run.rs

@@ -167,7 +167,7 @@ async fn run_pipeline(
     if !errors.errors.is_empty() {
         eprintln!("There were some issues with the provided query");
         for error in errors.errors {
-            eprintln!("  * {error}");
+            eprintln!("  * {}", error.message);
         }
         exit(1);
     }