fix: mangled errors

Author: mesejo

python/tests/test_catalog.py

@@ -81,6 +81,12 @@ def table_exist(self, name: str) -> bool:
         return name in self.tables
 
 
+class CustomErrorSchemaProvider(CustomSchemaProvider):
+    def table(self, name: str) -> Table | None:
+        message = f"{name} is not an acceptable name"
+        raise ValueError(message)
+
+
 class CustomCatalogProvider(dfn.catalog.CatalogProvider):
     def __init__(self):
         self.schemas = {"my_schema": CustomSchemaProvider()}
@@ -197,6 +203,33 @@ def test_python_table_provider(ctx: SessionContext):
     assert schema.table_names() == {"table4"}
 
 
+def test_exception_not_mangled(ctx: SessionContext):
+    """Test registering all python providers and running a query against them."""
+
+    catalog_name = "custom_catalog"
+    schema_name = "custom_schema"
+
+    ctx.register_catalog_provider(catalog_name, CustomCatalogProvider())
+
+    catalog = ctx.catalog(catalog_name)
+
+    # Clean out previous schemas if they exist so we can start clean
+    for schema_name in catalog.schema_names():
+        catalog.deregister_schema(schema_name, cascade=False)
+
+    catalog.register_schema(schema_name, CustomErrorSchemaProvider())
+
+    schema = catalog.schema(schema_name)
+
+    for table_name in schema.table_names():
+        schema.deregister_table(table_name)
+
+    schema.register_table("test_table", create_dataset())
+
+    with pytest.raises(ValueError, match="^test_table is not an acceptable name$"):
+        ctx.sql(f"select * from {catalog_name}.{schema_name}.test_table")
+
+
 def test_schema_register_table_with_pyarrow_dataset(ctx: SessionContext):
     schema = ctx.catalog().schema()
     batch = pa.RecordBatch.from_arrays(

python/tests/test_sql.py

@@ -29,7 +29,10 @@
 
 
 def test_no_table(ctx):
-    with pytest.raises(Exception, match="DataFusion error"):
+    with pytest.raises(
+        ValueError,
+        match="^Error during planning: table 'datafusion.public.b' not found$",
+    ):
         ctx.sql("SELECT a FROM b").collect()
 
 

src/catalog.rs

@@ -364,7 +364,8 @@ impl SchemaProvider for RustWrappedPySchemaProvider {
         &self,
         name: &str,
     ) -> datafusion::common::Result<Option<Arc<dyn TableProvider>>, DataFusionError> {
-        self.table_inner(name).map_err(to_datafusion_err)
+        self.table_inner(name)
+            .map_err(|e| DataFusionError::External(Box::new(e)))
     }
 
     fn register_table(

src/errors.rs

@@ -22,6 +22,7 @@ use std::fmt::Debug;
 use datafusion::arrow::error::ArrowError;
 use datafusion::error::DataFusionError as InnerDataFusionError;
 use prost::EncodeError;
+use pyo3::exceptions::PyValueError;
 use pyo3::exceptions::PyException;
 use pyo3::PyErr;
 
@@ -96,3 +97,13 @@ pub fn py_unsupported_variant_err(e: impl Debug) -> PyErr {
 pub fn to_datafusion_err(e: impl Debug) -> InnerDataFusionError {
     InnerDataFusionError::Execution(format!("{e:?}"))
 }
+
+pub fn from_datafusion_error(err: InnerDataFusionError) -> PyErr {
+    match err {
+        InnerDataFusionError::External(boxed) => match boxed.downcast::<PyErr>() {
+            Ok(py_err) => *py_err,
+            Err(original_boxed) => PyValueError::new_err(format!("{original_boxed}")),
+        },
+        _ => PyValueError::new_err(format!("{err}")),
+    }
+}