fix: better error handling, avoiding panics

Author: davidsmfreire

sqlshield-py/src/lib.rs

@@ -1,4 +1,4 @@
-use pyo3::prelude::*;
+use pyo3::{exceptions::PyValueError, prelude::*};
 
 extern crate sqlshield as sqlshield_rs;
 
@@ -38,7 +38,7 @@ fn validate_files(dir: String, schema_file_path: String) -> PyResult<Vec<PySqlVa
 
 #[pyfunction]
 fn validate_query(query: String, schema: String) -> PyResult<Vec<String>> {
-    Ok(sqlshield_rs::validate_query(query, schema))
+    sqlshield_rs::validate_query(query, schema).map_err(|err| PyValueError::new_err(err))
 }
 
 /// A Python module implemented in Rust.

sqlshield/src/lib.rs

@@ -9,14 +9,20 @@ use regex::Regex;
 use validation::{validate_queries_in_code, validate_statements_with_schema, SqlValidationError};
 use walkdir::WalkDir;
 
-pub fn validate_query(query: String, schema: String) -> Vec<String> {
+pub fn validate_query(query: String, schema: String) -> Result<Vec<String>, String> {
     let dialect = sqlparser::dialect::GenericDialect {};
 
-    let statements = sqlparser::parser::Parser::parse_sql(&dialect, &query).unwrap();
+    let statements = match sqlparser::parser::Parser::parse_sql(&dialect, &query) {
+        Ok(statements) => statements,
+        Err(err) => return Err(err.to_string()),
+    };
 
-    let loaded_schema = schema::load_schema(&schema.into_bytes(), "sql").unwrap();
+    let loaded_schema = match schema::load_schema(&schema.into_bytes(), "sql") {
+        Ok(loaded_schema) => loaded_schema,
+        Err(err) => return Err(err),
+    };
 
-    validate_statements_with_schema(&statements, &loaded_schema)
+    Ok(validate_statements_with_schema(&statements, &loaded_schema))
 }
 
 pub fn validate_files(dir: &PathBuf, schema_file_path: &PathBuf) -> Vec<SqlValidationError> {

sqlshield/src/schema/sql.rs

@@ -5,8 +5,10 @@ pub fn load_schema(schema: &[u8]) -> Result<super::TablesAndColumns, String> {
     let schema_str = String::from_utf8_lossy(schema);
 
     let dialect = GenericDialect {};
-    let statements =
-        Parser::parse_sql(&dialect, schema_str.as_ref()).expect("Could not parse schema file");
+    let statements = match Parser::parse_sql(&dialect, schema_str.as_ref()) {
+        Ok(statements) => statements,
+        Err(err) => return Err(format!("Could not parse schema file: {err}")),
+    };
 
     let mut tables: HashMap<String, HashSet<String>> = HashMap::new();
     for statement in statements {