From c535c24d03aa6c1eeee5dd4654aa9e9109b017a5 Mon Sep 17 00:00:00 2001
From: Jack Wright <56345+ayax79@users.noreply.github.com>
Date: Sun, 15 Sep 2024 05:21:02 -0700
Subject: [PATCH] catch unwrap on panics with `polars collect` (#13850)

# Description
This resurrects the work from #12866 and fixes #12732.

Polars panics for a plethora or reasons. While handling panics is
generally frowned upon, in cases like with `polars collect` a panic
cause a lot of work to be lost. Often you might have multiple dataframes
in memory and you are trying one operation and lose all state.

While it possible the panic can leave things a strange state, it is
pretty unlikely as part of a polars pipeline. Most of the time polars
objects are not manipulating dataframes in memory mutability, but rather
creating a new dataframe the operations being applied. This is always
the case with a lazy pipeline. After the collect call, the original
dataframes are intact still and I haven't observed any side effects.
---
 .../src/dataframe/values/nu_lazyframe/mod.rs  | 25 +++++++++++--------
 crates/nu_plugin_polars/src/lib.rs            | 21 +++++++++++++++-
 2 files changed, 35 insertions(+), 11 deletions(-)

diff --git a/crates/nu_plugin_polars/src/dataframe/values/nu_lazyframe/mod.rs b/crates/nu_plugin_polars/src/dataframe/values/nu_lazyframe/mod.rs
index e89f14c31680..dad72cab74f3 100644
--- a/crates/nu_plugin_polars/src/dataframe/values/nu_lazyframe/mod.rs
+++ b/crates/nu_plugin_polars/src/dataframe/values/nu_lazyframe/mod.rs
@@ -55,16 +55,21 @@ impl NuLazyFrame {
     }
 
     pub fn collect(self, span: Span) -> Result<NuDataFrame, ShellError> {
-        self.to_polars()
-            .collect()
-            .map_err(|e| ShellError::GenericError {
-                error: "Error collecting lazy frame".into(),
-                msg: e.to_string(),
-                span: Some(span),
-                help: None,
-                inner: vec![],
-            })
-            .map(|df| NuDataFrame::new(true, df))
+        crate::handle_panic(
+            || {
+                self.to_polars()
+                    .collect()
+                    .map_err(|e| ShellError::GenericError {
+                        error: "Error collecting lazy frame".into(),
+                        msg: e.to_string(),
+                        span: Some(span),
+                        help: None,
+                        inner: vec![],
+                    })
+                    .map(|df| NuDataFrame::new(true, df))
+            },
+            span,
+        )
     }
 
     pub fn apply_with_expr<F>(self, expr: NuExpression, f: F) -> Self
diff --git a/crates/nu_plugin_polars/src/lib.rs b/crates/nu_plugin_polars/src/lib.rs
index f70ab2b6efd5..342e5cd28c9f 100644
--- a/crates/nu_plugin_polars/src/lib.rs
+++ b/crates/nu_plugin_polars/src/lib.rs
@@ -1,4 +1,7 @@
-use std::cmp::Ordering;
+use std::{
+    cmp::Ordering,
+    panic::{catch_unwind, AssertUnwindSafe},
+};
 
 use cache::cache_commands;
 pub use cache::{Cache, Cacheable};
@@ -209,6 +212,22 @@ impl Plugin for PolarsPlugin {
     }
 }
 
+pub(crate) fn handle_panic<F, R>(f: F, span: Span) -> Result<R, ShellError>
+where
+    F: FnOnce() -> Result<R, ShellError>,
+{
+    match catch_unwind(AssertUnwindSafe(f)) {
+        Ok(inner_result) => inner_result,
+        Err(_) => Err(ShellError::GenericError {
+            error: "Panic occurred".into(),
+            msg: "".into(),
+            span: Some(span),
+            help: None,
+            inner: vec![],
+        }),
+    }
+}
+
 #[cfg(test)]
 pub mod test {
     use super::*;