Fix session get-or-default (#5662)

The comments described this get-or-default, but instead it was a panic

---------

Signed-off-by: Nicholas Gates <[email protected]>

Author: gatesn

vortex-session/src/lib.rs

@@ -60,12 +60,20 @@ pub trait SessionExt: Sized + private::Sealed {
     fn session(&self) -> VortexSession;
 
     /// Returns the scope variable of type `V`, or inserts a default one if it does not exist.
-    fn get<V: SessionVar>(&self) -> Ref<'_, V>;
+    fn get<V: SessionVar + Default>(&self) -> Ref<'_, V>;
+
+    /// Returns the scope variable of type `V` if it exists.
+    fn get_opt<V: SessionVar>(&self) -> Option<Ref<'_, V>>;
 
     /// Returns the scope variable of type `V`, or inserts a default one if it does not exist.
     ///
     /// Note that the returned value internally holds a lock on the variable.
-    fn get_mut<V: SessionVar>(&self) -> RefMut<'_, V>;
+    fn get_mut<V: SessionVar + Default>(&self) -> RefMut<'_, V>;
+
+    /// Returns the scope variable of type `V`, if it exists.
+    ///
+    /// Note that the returned value internally holds a lock on the variable.
+    fn get_mut_opt<V: SessionVar>(&self) -> Option<RefMut<'_, V>>;
 }
 
 mod private {
@@ -79,13 +87,12 @@ impl SessionExt for VortexSession {
     }
 
     /// Returns the scope variable of type `V`, or inserts a default one if it does not exist.
-    fn get<V: SessionVar>(&self) -> Ref<'_, V> {
+    fn get<V: SessionVar + Default>(&self) -> Ref<'_, V> {
         Ref(self
             .0
-            .get(&TypeId::of::<V>())
-            .unwrap_or_else(|| {
-                vortex_panic!("Session has not been initialized with {}", type_name::<V>())
-            })
+            .entry(TypeId::of::<V>())
+            .or_insert_with(|| Box::new(V::default()))
+            .downgrade()
             .map(|v| {
                 (**v)
                     .as_any()
@@ -94,16 +101,25 @@ impl SessionExt for VortexSession {
             }))
     }
 
+    fn get_opt<V: SessionVar>(&self) -> Option<Ref<'_, V>> {
+        self.0.get(&TypeId::of::<V>()).map(|v| {
+            Ref(v.map(|v| {
+                (**v)
+                    .as_any()
+                    .downcast_ref::<V>()
+                    .vortex_expect("Type mismatch - this is a bug")
+            }))
+        })
+    }
+
     /// Returns the scope variable of type `V`, or inserts a default one if it does not exist.
     ///
     /// Note that the returned value internally holds a lock on the variable.
-    fn get_mut<V: SessionVar>(&self) -> RefMut<'_, V> {
+    fn get_mut<V: SessionVar + Default>(&self) -> RefMut<'_, V> {
         RefMut(
             self.0
-                .get_mut(&TypeId::of::<V>())
-                .unwrap_or_else(|| {
-                    vortex_panic!("Session has not been initialized with {}", type_name::<V>())
-                })
+                .entry(TypeId::of::<V>())
+                .or_insert_with(|| Box::new(V::default()))
                 .map(|v| {
                     (**v)
                         .as_any_mut()
@@ -112,6 +128,17 @@ impl SessionExt for VortexSession {
                 }),
         )
     }
+
+    fn get_mut_opt<V: SessionVar>(&self) -> Option<RefMut<'_, V>> {
+        self.0.get_mut(&TypeId::of::<V>()).map(|v| {
+            RefMut(v.map(|v| {
+                (**v)
+                    .as_any_mut()
+                    .downcast_mut::<V>()
+                    .vortex_expect("Type mismatch - this is a bug")
+            }))
+        })
+    }
 }
 
 /// A TypeMap based on `https://docs.rs/http/1.2.0/src/http/extensions.rs.html#41-266`.