feat: better exception and message for table not found (apache#851)

mesejo · web-flow · commit fe0738a9c0b5 · 2024-09-02T09:28:59.000-06:00
closes apache#796
diff --git a/python/datafusion/tests/test_context.py b/python/datafusion/tests/test_context.py
@@ -466,6 +466,13 @@ def test_table_exist(ctx):
     assert ctx.table_exist("t") is True
 
 
+def test_table_not_found(ctx):
+    from uuid import uuid4
+
+    with pytest.raises(KeyError):
+        ctx.table(f"not-found-{uuid4()}")
+
+
 def test_read_json(ctx):
     path = os.path.dirname(os.path.abspath(__file__))
 
diff --git a/src/context.rs b/src/context.rs
@@ -765,7 +765,8 @@ impl PySessionContext {
     }
 
     pub fn table(&self, name: &str, py: Python) -> PyResult<PyDataFrame> {
-        let x = wait_for_future(py, self.ctx.table(name)).map_err(DataFusionError::from)?;
+        let x = wait_for_future(py, self.ctx.table(name))
+            .map_err(|e| PyKeyError::new_err(e.to_string()))?;
         Ok(PyDataFrame::new(x))
     }
 

Original file line number	Diff line number	Diff line change
`@@ -765,7 +765,8 @@ impl PySessionContext {`
`765`	`765`	`}`
`766`	`766`
`767`	`767`	`pub fn table(&self, name: &str, py: Python) -> PyResult<PyDataFrame> {`
`768`		`- let x = wait_for_future(py, self.ctx.table(name)).map_err(DataFusionError::from)?;`
	`768`	`+ let x = wait_for_future(py, self.ctx.table(name))`
	`769`	`+ .map_err(\|e\| PyKeyError::new_err(e.to_string()))?;`
`769`	`770`	`Ok(PyDataFrame::new(x))`
`770`	`771`	`}`
`771`	`772`