feat!(rust/ffi): catch panics at FFI boundary

.github/workflows/native-unix.yml

@@ -598,12 +598,18 @@ jobs:
           else
             make -C ./go/adbc/pkg libadbc_driver_panicdummy.so
           fi
+      - name: Build Panic Dummy (Rust)
+        working-directory: rust
+        run: |
+          cargo build -padbc_dummy
       - name: Test Python Driver Manager
         run: |
           if [[ $(uname) = "Darwin" ]]; then
             export PANICDUMMY_LIBRARY_PATH=$(pwd)/go/adbc/pkg/libadbc_driver_panicdummy.dylib
+            export PANICDUMMY_RUST_LIBRARY_PATH=$(pwd)/rust/target/debug/libadbc_dummy.dylib
           else
             export PANICDUMMY_LIBRARY_PATH=$(pwd)/go/adbc/pkg/libadbc_driver_panicdummy.so
+            export PANICDUMMY_RUST_LIBRARY_PATH=$(pwd)/rust/target/debug/libadbc_dummy.so
           fi
           export PATH=$RUNNER_TOOL_CACHE/go/${GO_VERSION}/x64/bin:$PATH
           env BUILD_ALL=0 BUILD_DRIVER_MANAGER=1 ./ci/scripts/python_test.sh "$(pwd)" "$(pwd)/build" "$HOME/local"

python/adbc_driver_manager/tests/test_panic.py

@@ -15,7 +15,7 @@
 # specific language governing permissions and limitations
 # under the License.
 
-"""Tests for the panic behavior of the Go driver FFI wrapper."""
+"""Tests for the panic behavior of the Go/Rust FFI wrappers."""
 
 import os
 import subprocess
@@ -26,25 +26,33 @@
 
 pytestmark = pytest.mark.panicdummy
 
-_LIB_ENV_VAR = "PANICDUMMY_LIBRARY_PATH"
+_GO_LIB_ENV_VAR = "PANICDUMMY_LIBRARY_PATH"
+_RUST_LIB_ENV_VAR = "PANICDUMMY_RUST_LIBRARY_PATH"
 
 
 @pytest.fixture(scope="module")
-def libpath() -> str:
-    if _LIB_ENV_VAR not in os.environ:
-        pytest.skip(f"{_LIB_ENV_VAR} not specified", allow_module_level=True)
-    return os.environ[_LIB_ENV_VAR]
+def go_driver() -> str:
+    if _GO_LIB_ENV_VAR not in os.environ:
+        pytest.skip(f"{_GO_LIB_ENV_VAR} not specified", allow_module_level=True)
+    return os.environ[_GO_LIB_ENV_VAR]
 
 
-def test_panic_close(libpath) -> None:
+@pytest.fixture(scope="module")
+def rust_driver() -> str:
+    if _RUST_LIB_ENV_VAR not in os.environ:
+        pytest.skip(f"{_RUST_LIB_ENV_VAR} not specified", allow_module_level=True)
+    return os.environ[_RUST_LIB_ENV_VAR]
+
+
+def test_panic_close_go(go_driver) -> None:
     env = os.environ.copy()
     env["PANICDUMMY_FUNC"] = "StatementClose"
     env["PANICDUMMY_MESSAGE"] = "Boo!"
     output = subprocess.run(
         [
             sys.executable,
             Path(__file__).parent / "panictest.py",
-            libpath,
+            go_driver,
         ],
         stdout=subprocess.PIPE,
         stderr=subprocess.PIPE,
@@ -57,15 +65,15 @@ def test_panic_close(libpath) -> None:
     assert "Go panicked, driver is in unknown state" in output.stderr
 
 
-def test_panic_execute(libpath) -> None:
+def test_panic_execute_go(go_driver) -> None:
     env = os.environ.copy()
     env["PANICDUMMY_FUNC"] = "StatementExecuteQuery"
     env["PANICDUMMY_MESSAGE"] = "Boo!"
     output = subprocess.run(
         [
             sys.executable,
             Path(__file__).parent / "panictest.py",
-            libpath,
+            go_driver,
         ],
         stdout=subprocess.PIPE,
         stderr=subprocess.PIPE,
@@ -76,3 +84,45 @@ def test_panic_execute(libpath) -> None:
     assert "Go panic in PanicDummy driver" in output.stderr
     assert "Boo!" in output.stderr
     assert "Go panicked, driver is in unknown state" in output.stderr
+
+
+def test_panic_close_rust(rust_driver) -> None:
+    env = os.environ.copy()
+    env["PANICDUMMY_FUNC"] = "StatementClose"
+    env["PANICDUMMY_MESSAGE"] = "Boo!"
+    output = subprocess.run(
+        [
+            sys.executable,
+            Path(__file__).parent / "panictest.py",
+            rust_driver,
+        ],
+        stdout=subprocess.PIPE,
+        stderr=subprocess.PIPE,
+        env=env,
+        encoding="utf-8",
+    )
+    assert output.returncode != 0
+    assert "Uncaught panic in driver" in output.stderr
+    assert "Boo!" in output.stderr
+    assert "Driver panicked and is in unknown state" in output.stderr
+
+
+def test_panic_execute_rust(rust_driver) -> None:
+    env = os.environ.copy()
+    env["PANICDUMMY_FUNC"] = "StatementExecuteQuery"
+    env["PANICDUMMY_MESSAGE"] = "Boo!"
+    output = subprocess.run(
+        [
+            sys.executable,
+            Path(__file__).parent / "panictest.py",
+            rust_driver,
+        ],
+        stdout=subprocess.PIPE,
+        stderr=subprocess.PIPE,
+        env=env,
+        encoding="utf-8",
+    )
+    assert output.returncode != 0
+    assert "Uncaught panic in driver" in output.stderr
+    assert "Boo!" in output.stderr
+    assert "Driver panicked and is in unknown state" in output.stderr

rust/driver/dummy/src/lib.rs

@@ -175,6 +175,18 @@ where
     }
 }
 
+fn maybe_panic(fnname: impl AsRef<str>) {
+    if let Some(func) = std::env::var_os("PANICDUMMY_FUNC").map(|x| x.to_string_lossy().to_string())
+    {
+        if fnname.as_ref() == func {
+            let message = std::env::var_os("PANICDUMMY_MESSAGE")
+                .map(|x| x.to_string_lossy().to_string())
+                .unwrap_or_else(|| format!("We panicked in {}!", fnname.as_ref()));
+            panic!("{}", message);
+        }
+    }
+}
+
 /// A dummy driver used for testing purposes.
 #[derive(Default)]
 pub struct DummyDriver {}
@@ -837,6 +849,7 @@ impl Statement for DummyStatement {
     }
 
     fn execute(&mut self) -> Result<impl RecordBatchReader> {
+        maybe_panic("StatementExecuteQuery");
         let batch = get_table_data();
         let reader = SingleBatchReader::new(batch);
         Ok(reader)
@@ -875,4 +888,10 @@ impl Statement for DummyStatement {
     }
 }
 
-adbc_ffi::export_driver!(DummyDriverInit, DummyDriver);
+impl Drop for DummyStatement {
+    fn drop(&mut self) {
+        maybe_panic("StatementClose");
+    }
+}
+
+adbc_ffi::export_driver!(AdbcDummyInit, DummyDriver);

rust/driver/dummy/tests/driver_exporter_dummy.rs

@@ -67,7 +67,7 @@ fn get_exported() -> (
 ) {
     let mut driver = ManagedDriver::load_dynamic_from_name(
         "adbc_dummy",
-        Some(b"DummyDriverInit"),
+        Some(b"AdbcDummyInit"),
         AdbcVersion::V110,
     )
     .unwrap();