Added commands for working with the plugin cache. (#12576)

# Description
This pull request provides three new commands:
`polars store-ls` - moved from `polars ls`. It provides the list of all
object stored in the plugin cache
`polars store-rm` - deletes a cached object
`polars store-get` - gets an object from the cache. 

The addition of `polars store-get` required adding a reference_count to
cached entries. `polars get` is the only command that will increment
this value. `polars rm` will remove the value despite it's count. Calls
to PolarsPlugin::custom_value_dropped will decrement the value.

The prefix store- was chosen due to there already being a `polars cache`
command. These commands were not made sub-commands as there isn't a way
to display help for sub commands in plugins (e.g. `polars store`
displaying help) and I felt the store- seemed fine anyways.

The output of `polars store-ls` now shows the reference count for each
object.

# User-Facing Changes
polars ls has now moved to polars store-ls

---------

Co-authored-by: Jack Wright <jack.wright@disqo.com>

crates/nu_plugin_polars/src/cache/get.rs

@@ -0,0 +1,96 @@
+use nu_plugin::{EngineInterface, EvaluatedCall, PluginCommand};
+use nu_protocol::{
+    Category, Example, LabeledError, PipelineData, ShellError, Signature, Span, SyntaxShape, Type,
+    Value,
+};
+use polars::{prelude::NamedFrom, series::Series};
+use uuid::Uuid;
+
+use crate::{
+    values::{CustomValueSupport, NuDataFrame},
+    PolarsPlugin,
+};
+
+#[derive(Clone)]
+pub struct CacheGet;
+
+impl PluginCommand for CacheGet {
+    type Plugin = PolarsPlugin;
+
+    fn name(&self) -> &str {
+        "polars store-get"
+    }
+
+    fn usage(&self) -> &str {
+        "Gets a Dataframe or other object from the plugin cache."
+    }
+
+    fn signature(&self) -> Signature {
+        Signature::build(self.name())
+            .required("key", SyntaxShape::String, "Key of objects to get")
+            .input_output_types(vec![
+                (Type::Any, Type::Custom("dataframe".into())),
+                (Type::Any, Type::Custom("expression".into())),
+            ])
+            .category(Category::Custom("dataframe".into()))
+    }
+
+    fn examples(&self) -> Vec<Example> {
+        vec![Example {
+            description: "Get a stored object",
+            example: r#"let df = ([[a b];[1 2] [3 4]] | polars into-df);
+    polars store-ls | get key | first | polars store-get $in"#,
+            result: Some(
+                NuDataFrame::try_from_series_vec(
+                    vec![Series::new("a", &[1_i64, 3]), Series::new("b", &[2_i64, 4])],
+                    Span::test_data(),
+                )
+                .expect("could not create dataframe")
+                .into_value(Span::test_data()),
+            ),
+        }]
+    }
+
+    fn run(
+        &self,
+        plugin: &Self::Plugin,
+        _engine: &EngineInterface,
+        call: &EvaluatedCall,
+        _input: PipelineData,
+    ) -> Result<PipelineData, LabeledError> {
+        let key = call
+            .req::<String>(0)
+            .and_then(|ref k| as_uuid(k, call.head))?;
+
+        let value = if let Some(cache_value) = plugin.cache.get(&key, true)? {
+            let polars_object = cache_value.value;
+            polars_object.into_value(call.head)
+        } else {
+            Value::nothing(call.head)
+        };
+
+        Ok(PipelineData::Value(value, None))
+    }
+}
+
+fn as_uuid(s: &str, span: Span) -> Result<Uuid, ShellError> {
+    Uuid::parse_str(s).map_err(|e| ShellError::GenericError {
+        error: format!("Failed to convert key string to UUID: {e}"),
+        msg: "".into(),
+        span: Some(span),
+        help: None,
+        inner: vec![],
+    })
+}
+
+#[cfg(test)]
+mod test {
+    use super::*;
+    use crate::test::test_polars_plugin_command_with_decls;
+    use nu_command::{First, Get};
+
+    #[test]
+    fn test_examples() -> Result<(), ShellError> {
+        test_polars_plugin_command_with_decls(&CacheGet, vec![Box::new(Get), Box::new(First)])
+    }
+}

crates/nu_plugin_polars/src/cache/list.rs

@@ -12,7 +12,7 @@ impl PluginCommand for ListDF {
     type Plugin = PolarsPlugin;
 
     fn name(&self) -> &str {
-        "polars ls"
+        "polars store-ls"
     }
 
     fn usage(&self) -> &str {
@@ -26,8 +26,8 @@ impl PluginCommand for ListDF {
     fn examples(&self) -> Vec<Example> {
         vec![Example {
             description: "Creates a new dataframe and shows it in the dataframe list",
-            example: r#"let test = ([[a b];[1 2] [3 4]] | dfr into-df);
+            example: r#"let test = ([[a b];[1 2] [3 4]] | polars into-df);
-    polars ls"#,
+    polars store-ls"#,
             result: None,
         }]
     }
@@ -49,11 +49,12 @@ impl PluginCommand for ListDF {
                         "created" => Value::date(value.created, call.head),
                         "columns" => Value::int(df.as_ref().width() as i64, call.head),
                         "rows" => Value::int(df.as_ref().height() as i64, call.head),
-                        "type" => Value::string("NuDataFrame", call.head),
+                        "type" => Value::string("DataFrame", call.head),
                         "estimated_size" => Value::filesize(df.to_polars().estimated_size() as i64, call.head),
                         "span_contents" =>  Value::string(span_contents, value.span),
                         "span_start" => Value::int(value.span.start as i64, call.head),
                         "span_end" => Value::int(value.span.end as i64, call.head),
+                        "reference_count" => Value::int(value.reference_count as i64, call.head),
                     },
                     call.head,
                 ))),
@@ -65,11 +66,12 @@ impl PluginCommand for ListDF {
                             "created" => Value::date(value.created, call.head),
                             "columns" => Value::int(lf.as_ref().width() as i64, call.head),
                             "rows" => Value::int(lf.as_ref().height() as i64, call.head),
-                            "type" => Value::string("NuLazyFrame", call.head),
+                            "type" => Value::string("LazyFrame", call.head),
                             "estimated_size" => Value::filesize(lf.to_polars().estimated_size() as i64, call.head),
                             "span_contents" =>  Value::string(span_contents, value.span),
                             "span_start" => Value::int(value.span.start as i64, call.head),
                             "span_end" => Value::int(value.span.end as i64, call.head),
+                            "reference_count" => Value::int(value.reference_count as i64, call.head),
                         },
                         call.head,
                     )))
@@ -80,11 +82,12 @@ impl PluginCommand for ListDF {
                         "created" => Value::date(value.created, call.head),
                         "columns" => Value::nothing(call.head),
                         "rows" => Value::nothing(call.head),
-                        "type" => Value::string("NuExpression", call.head),
+                        "type" => Value::string("Expression", call.head),
                         "estimated_size" => Value::nothing(call.head),
                         "span_contents" =>  Value::string(span_contents, value.span),
                         "span_start" => Value::int(value.span.start as i64, call.head),
                         "span_end" => Value::int(value.span.end as i64, call.head),
+                        "reference_count" => Value::int(value.reference_count as i64, call.head),
                     },
                     call.head,
                 ))),
@@ -93,11 +96,12 @@ impl PluginCommand for ListDF {
                         "key" => Value::string(key.to_string(), call.head),
                         "columns" => Value::nothing(call.head),
                         "rows" => Value::nothing(call.head),
-                        "type" => Value::string("NuLazyGroupBy", call.head),
+                        "type" => Value::string("LazyGroupBy", call.head),
                         "estimated_size" => Value::nothing(call.head),
                         "span_contents" =>  Value::string(span_contents, call.head),
                         "span_start" => Value::int(call.head.start as i64, call.head),
                         "span_end" => Value::int(call.head.end as i64, call.head),
+                        "reference_count" => Value::int(value.reference_count as i64, call.head),
                     },
                     call.head,
                 ))),
@@ -106,11 +110,12 @@ impl PluginCommand for ListDF {
                         "key" => Value::string(key.to_string(), call.head),
                         "columns" => Value::nothing(call.head),
                         "rows" => Value::nothing(call.head),
-                        "type" => Value::string("NuWhen", call.head),
+                        "type" => Value::string("When", call.head),
                         "estimated_size" => Value::nothing(call.head),
                         "span_contents" =>  Value::string(span_contents.to_string(), call.head),
                         "span_start" => Value::int(call.head.start as i64, call.head),
                         "span_end" => Value::int(call.head.end as i64, call.head),
+                        "reference_count" => Value::int(value.reference_count as i64, call.head),
                     },
                     call.head,
                 ))),

crates/nu_plugin_polars/src/cache/mod.rs

@@ -1,10 +1,15 @@
+mod get;
+mod list;
+mod rm;
+
 use std::{
     collections::HashMap,
     sync::{Mutex, MutexGuard},
 };
 
 use chrono::{DateTime, FixedOffset, Local};
-use nu_plugin::EngineInterface;
+pub use list::ListDF;
+use nu_plugin::{EngineInterface, PluginCommand};
 use nu_protocol::{LabeledError, ShellError, Span};
 use uuid::Uuid;
 
@@ -16,6 +21,7 @@ pub struct CacheValue {
     pub value: PolarsPluginObject,
     pub created: DateTime<FixedOffset>,
     pub span: Span,
+    pub reference_count: i16,
 }
 
 #[derive(Default)]
@@ -35,15 +41,32 @@ impl Cache {
     }
 
     /// Removes an item from the plugin cache.
-    /// The maybe_engine parameter is required outside of testing
+    ///
+    /// * `maybe_engine` - Current EngineInterface reference. Required outside of testing
+    /// * `key` - The key of the cache entry to remove.
+    /// * `force` - Delete even if there are multiple references
     pub fn remove(
         &self,
         maybe_engine: Option<&EngineInterface>,
-        uuid: &Uuid,
+        key: &Uuid,
+        force: bool,
     ) -> Result<Option<CacheValue>, ShellError> {
         let mut lock = self.lock()?;
-        let removed = lock.remove(uuid);
+
-        plugin_debug!("PolarsPlugin: removing {uuid} from cache: {removed:?}");
+        let reference_count = lock.get_mut(key).map(|cache_value| {
+            cache_value.reference_count -= 1;
+            cache_value.reference_count
+        });
+
+        let removed = if force || reference_count.unwrap_or_default() < 1 {
+            let removed = lock.remove(key);
+            plugin_debug!("PolarsPlugin: removing {key} from cache: {removed:?}");
+            removed
+        } else {
+            plugin_debug!("PolarsPlugin: decrementing reference count for {key}");
+            None
+        };
+
         // Once there are no more entries in the cache
         // we can turn plugin gc back on
         match maybe_engine {
@@ -83,15 +106,21 @@ impl Cache {
             value,
             created: Local::now().into(),
             span,
+            reference_count: 1,
         };
         let result = lock.insert(uuid, cache_value);
         drop(lock);
         Ok(result)
     }
 
-    pub fn get(&self, uuid: &Uuid) -> Result<Option<CacheValue>, ShellError> {
+    pub fn get(&self, uuid: &Uuid, increment: bool) -> Result<Option<CacheValue>, ShellError> {
-        let lock = self.lock()?;
+        let mut lock = self.lock()?;
-        let result = lock.get(uuid).cloned();
+        let result = lock.get_mut(uuid).map(|cv| {
+            if increment {
+                cv.reference_count += 1;
+            }
+            cv.clone()
+        });
         drop(lock);
         Ok(result)
     }
@@ -134,10 +163,18 @@ pub trait Cacheable: Sized + Clone {
     }
 
     fn get_cached(plugin: &PolarsPlugin, id: &Uuid) -> Result<Option<Self>, ShellError> {
-        if let Some(cache_value) = plugin.cache.get(id)? {
+        if let Some(cache_value) = plugin.cache.get(id, false)? {
             Ok(Some(Self::from_cache_value(cache_value.value)?))
         } else {
             Ok(None)
         }
     }
 }
+
+pub(crate) fn cache_commands() -> Vec<Box<dyn PluginCommand<Plugin = PolarsPlugin>>> {
+    vec![
+        Box::new(ListDF),
+        Box::new(rm::CacheRemove),
+        Box::new(get::CacheGet),
+    ]
+}

crates/nu_plugin_polars/src/cache/rm.rs

@@ -0,0 +1,106 @@
+use nu_plugin::{EngineInterface, EvaluatedCall, PluginCommand};
+use nu_protocol::{
+    Category, Example, LabeledError, PipelineData, ShellError, Signature, Span, SyntaxShape, Type,
+    Value,
+};
+use uuid::Uuid;
+
+use crate::PolarsPlugin;
+
+#[derive(Clone)]
+pub struct CacheRemove;
+
+impl PluginCommand for CacheRemove {
+    type Plugin = PolarsPlugin;
+
+    fn name(&self) -> &str {
+        "polars store-rm"
+    }
+
+    fn usage(&self) -> &str {
+        "Removes a stored Dataframe or other object from the plugin cache."
+    }
+
+    fn signature(&self) -> Signature {
+        Signature::build(self.name())
+            .rest("keys", SyntaxShape::String, "Keys of objects to remove")
+            .input_output_type(Type::Any, Type::List(Box::new(Type::String)))
+            .category(Category::Custom("dataframe".into()))
+    }
+
+    fn examples(&self) -> Vec<Example> {
+        vec![Example {
+            description: "Removes a stored ",
+            example: r#"let df = ([[a b];[1 2] [3 4]] | polars into-df);
+    polars store-ls | get key | first | polars store-rm $in"#,
+            result: None,
+        }]
+    }
+
+    fn run(
+        &self,
+        plugin: &Self::Plugin,
+        engine: &EngineInterface,
+        call: &EvaluatedCall,
+        _input: PipelineData,
+    ) -> Result<PipelineData, LabeledError> {
+        let msgs: Vec<Value> = call
+            .rest::<String>(0)?
+            .into_iter()
+            .map(|ref key| remove_cache_entry(plugin, engine, key, call.head))
+            .collect::<Result<Vec<Value>, ShellError>>()?;
+
+        Ok(PipelineData::Value(Value::list(msgs, call.head), None))
+    }
+}
+
+fn remove_cache_entry(
+    plugin: &PolarsPlugin,
+    engine: &EngineInterface,
+    key: &str,
+    span: Span,
+) -> Result<Value, ShellError> {
+    let key = as_uuid(key, span)?;
+    let msg = plugin
+        .cache
+        .remove(Some(engine), &key, true)?
+        .map(|_| format!("Removed: {key}"))
+        .unwrap_or_else(|| format!("No value found for key: {key}"));
+    Ok(Value::string(msg, span))
+}
+
+fn as_uuid(s: &str, span: Span) -> Result<Uuid, ShellError> {
+    Uuid::parse_str(s).map_err(|e| ShellError::GenericError {
+        error: format!("Failed to convert key string to UUID: {e}"),
+        msg: "".into(),
+        span: Some(span),
+        help: None,
+        inner: vec![],
+    })
+}
+
+#[cfg(test)]
+mod test {
+    use nu_command::{First, Get};
+    use nu_plugin_test_support::PluginTest;
+    use nu_protocol::Span;
+
+    use super::*;
+
+    #[test]
+    fn test_remove() -> Result<(), ShellError> {
+        let plugin = PolarsPlugin::new_test_mode().into();
+        let pipeline_data = PluginTest::new("polars", plugin)?
+            .add_decl(Box::new(First))?
+            .add_decl(Box::new(Get))?
+            .eval("let df = ([[a b];[1 2] [3 4]] | polars into-df); polars store-ls | get key | first | polars store-rm $in")?;
+        let value = pipeline_data.into_value(Span::test_data());
+        let msg = value
+            .as_list()?
+            .first()
+            .expect("there should be a first entry")
+            .as_str()?;
+        assert!(msg.contains("Removed"));
+        Ok(())
+    }
+}

crates/nu_plugin_polars/src/dataframe/eager/mod.rs

@@ -9,7 +9,6 @@ mod filter_with;
 mod first;
 mod get;
 mod last;
-mod list;
 mod melt;
 mod open;
 mod query_df;
@@ -45,7 +44,6 @@ pub use filter_with::FilterWith;
 pub use first::FirstDF;
 pub use get::GetDF;
 pub use last::LastDF;
-pub use list::ListDF;
 pub use melt::MeltDF;
 use nu_plugin::PluginCommand;
 pub use query_df::QueryDf;
@@ -82,7 +80,6 @@ pub(crate) fn eager_commands() -> Vec<Box<dyn PluginCommand<Plugin = PolarsPlugi
         Box::new(Summary),
         Box::new(FirstDF),
         Box::new(LastDF),
-        Box::new(ListDF),
         Box::new(RenameDF),
         Box::new(SampleDF),
         Box::new(ShapeDF),

crates/nu_plugin_polars/src/dataframe/series/date/as_datetime.rs

@@ -186,10 +186,11 @@ fn command(
 #[cfg(test)]
 mod test {
     use super::*;
-    use crate::test::test_polars_plugin_command;
+    use crate::test::test_polars_plugin_command_with_decls;
+    use nu_command::IntoDatetime;
 
     #[test]
     fn test_examples() -> Result<(), ShellError> {
-        test_polars_plugin_command(&AsDateTime)
+        test_polars_plugin_command_with_decls(&AsDateTime, vec![Box::new(IntoDatetime)])
     }
 }

crates/nu_plugin_polars/src/dataframe/series/date/get_day.rs

@@ -93,10 +93,11 @@ fn command(
 #[cfg(test)]
 mod test {
     use super::*;
-    use crate::test::test_polars_plugin_command;
+    use crate::test::test_polars_plugin_command_with_decls;
+    use nu_command::IntoDatetime;
 
     #[test]
     fn test_examples() -> Result<(), ShellError> {
-        test_polars_plugin_command(&GetDay)
+        test_polars_plugin_command_with_decls(&GetDay, vec![Box::new(IntoDatetime)])
     }
 }

crates/nu_plugin_polars/src/dataframe/series/date/get_hour.rs

@@ -85,10 +85,11 @@ fn command(
 #[cfg(test)]
 mod test {
     use super::*;
-    use crate::test::test_polars_plugin_command;
+    use crate::test::test_polars_plugin_command_with_decls;
+    use nu_command::IntoDatetime;
 
     #[test]
     fn test_examples() -> Result<(), ShellError> {
-        test_polars_plugin_command(&GetHour)
+        test_polars_plugin_command_with_decls(&GetHour, vec![Box::new(IntoDatetime)])
     }
 }

crates/nu_plugin_polars/src/dataframe/series/date/get_minute.rs

@@ -83,10 +83,11 @@ fn command(
 #[cfg(test)]
 mod test {
     use super::*;
-    use crate::test::test_polars_plugin_command;
+    use crate::test::test_polars_plugin_command_with_decls;
+    use nu_command::IntoDatetime;
 
     #[test]
     fn test_examples() -> Result<(), ShellError> {
-        test_polars_plugin_command(&GetMinute)
+        test_polars_plugin_command_with_decls(&GetMinute, vec![Box::new(IntoDatetime)])
     }
 }

crates/nu_plugin_polars/src/dataframe/series/date/get_month.rs

@@ -85,10 +85,11 @@ fn command(
 #[cfg(test)]
 mod test {
     use super::*;
-    use crate::test::test_polars_plugin_command;
+    use crate::test::test_polars_plugin_command_with_decls;
+    use nu_command::IntoDatetime;
 
     #[test]
     fn test_examples() -> Result<(), ShellError> {
-        test_polars_plugin_command(&GetMonth)
+        test_polars_plugin_command_with_decls(&GetMonth, vec![Box::new(IntoDatetime)])
     }
 }

crates/nu_plugin_polars/src/dataframe/series/date/get_nanosecond.rs

@@ -85,10 +85,11 @@ fn command(
 #[cfg(test)]
 mod test {
     use super::*;
-    use crate::test::test_polars_plugin_command;
+    use crate::test::test_polars_plugin_command_with_decls;
+    use nu_command::IntoDatetime;
 
     #[test]
     fn test_examples() -> Result<(), ShellError> {
-        test_polars_plugin_command(&GetNanosecond)
+        test_polars_plugin_command_with_decls(&GetNanosecond, vec![Box::new(IntoDatetime)])
     }
 }

crates/nu_plugin_polars/src/dataframe/series/date/get_ordinal.rs

@@ -85,10 +85,11 @@ fn command(
 #[cfg(test)]
 mod test {
     use super::*;
-    use crate::test::test_polars_plugin_command;
+    use crate::test::test_polars_plugin_command_with_decls;
+    use nu_command::IntoDatetime;
 
     #[test]
     fn test_examples() -> Result<(), ShellError> {
-        test_polars_plugin_command(&GetOrdinal)
+        test_polars_plugin_command_with_decls(&GetOrdinal, vec![Box::new(IntoDatetime)])
     }
 }

crates/nu_plugin_polars/src/dataframe/series/date/get_second.rs

@@ -85,10 +85,11 @@ fn command(
 #[cfg(test)]
 mod test {
     use super::*;
-    use crate::test::test_polars_plugin_command;
+    use crate::test::test_polars_plugin_command_with_decls;
+    use nu_command::IntoDatetime;
 
     #[test]
     fn test_examples() -> Result<(), ShellError> {
-        test_polars_plugin_command(&GetSecond)
+        test_polars_plugin_command_with_decls(&GetSecond, vec![Box::new(IntoDatetime)])
     }
 }

crates/nu_plugin_polars/src/dataframe/series/date/get_week.rs

@@ -85,10 +85,11 @@ fn command(
 #[cfg(test)]
 mod test {
     use super::*;
-    use crate::test::test_polars_plugin_command;
+    use crate::test::test_polars_plugin_command_with_decls;
+    use nu_command::IntoDatetime;
 
     #[test]
     fn test_examples() -> Result<(), ShellError> {
-        test_polars_plugin_command(&GetWeek)
+        test_polars_plugin_command_with_decls(&GetWeek, vec![Box::new(IntoDatetime)])
     }
 }

crates/nu_plugin_polars/src/dataframe/series/date/get_weekday.rs

@@ -85,10 +85,11 @@ fn command(
 #[cfg(test)]
 mod test {
     use super::*;
-    use crate::test::test_polars_plugin_command;
+    use crate::test::test_polars_plugin_command_with_decls;
+    use nu_command::IntoDatetime;
 
     #[test]
     fn test_examples() -> Result<(), ShellError> {
-        test_polars_plugin_command(&GetWeekDay)
+        test_polars_plugin_command_with_decls(&GetWeekDay, vec![Box::new(IntoDatetime)])
     }
 }

crates/nu_plugin_polars/src/dataframe/series/date/get_year.rs

@@ -85,10 +85,11 @@ fn command(
 #[cfg(test)]
 mod test {
     use super::*;
-    use crate::test::test_polars_plugin_command;
+    use crate::test::test_polars_plugin_command_with_decls;
+    use nu_command::IntoDatetime;
 
     #[test]
     fn test_examples() -> Result<(), ShellError> {
-        test_polars_plugin_command(&GetYear)
+        test_polars_plugin_command_with_decls(&GetYear, vec![Box::new(IntoDatetime)])
     }
 }

crates/nu_plugin_polars/src/dataframe/series/string/strftime.rs

@@ -104,10 +104,11 @@ fn command(
 #[cfg(test)]
 mod test {
     use super::*;
-    use crate::test::test_polars_plugin_command;
+    use crate::test::test_polars_plugin_command_with_decls;
+    use nu_command::IntoDatetime;
 
     #[test]
     fn test_examples() -> Result<(), ShellError> {
-        test_polars_plugin_command(&StrFTime)
+        test_polars_plugin_command_with_decls(&StrFTime, vec![Box::new(IntoDatetime)])
     }
 }

crates/nu_plugin_polars/src/dataframe/values/mod.rs

@@ -107,6 +107,16 @@ impl PolarsPluginObject {
             PolarsPluginObject::NuWhen(w) => w.id,
         }
     }
+
+    pub fn into_value(self, span: Span) -> Value {
+        match self {
+            PolarsPluginObject::NuDataFrame(df) => df.into_value(span),
+            PolarsPluginObject::NuLazyFrame(lf) => lf.into_value(span),
+            PolarsPluginObject::NuExpression(e) => e.into_value(span),
+            PolarsPluginObject::NuLazyGroupBy(lg) => lg.into_value(span),
+            PolarsPluginObject::NuWhen(w) => w.into_value(span),
+        }
+    }
 }
 
 #[derive(Debug, Clone)]

crates/nu_plugin_polars/src/lib.rs

@@ -1,5 +1,6 @@
 use std::cmp::Ordering;
 
+use cache::cache_commands;
 pub use cache::{Cache, Cacheable};
 use dataframe::{stub::PolarsCmd, values::CustomValueType};
 use nu_plugin::{EngineInterface, Plugin, PluginCommand};
@@ -40,6 +41,7 @@ impl Plugin for PolarsPlugin {
         commands.append(&mut lazy_commands());
         commands.append(&mut expr_commands());
         commands.append(&mut series_commands());
+        commands.append(&mut cache_commands());
         commands
     }
 
@@ -50,7 +52,7 @@ impl Plugin for PolarsPlugin {
     ) -> Result<(), LabeledError> {
         if !self.disable_cache_drop {
             let id = CustomValueType::try_from_custom_value(custom_value)?.id();
-            let _ = self.cache.remove(Some(engine), &id);
+            let _ = self.cache.remove(Some(engine), &id, false);
         }
         Ok(())
     }
@@ -178,9 +180,8 @@ impl Plugin for PolarsPlugin {
 pub mod test {
     use super::*;
     use crate::values::PolarsPluginObject;
-    use nu_command::IntoDatetime;
     use nu_plugin_test_support::PluginTest;
-    use nu_protocol::{ShellError, Span};
+    use nu_protocol::{engine::Command, ShellError, Span};
 
     impl PolarsPlugin {
         /// Creates a new polars plugin in test mode
@@ -193,6 +194,13 @@ pub mod test {
     }
 
     pub fn test_polars_plugin_command(command: &impl PluginCommand) -> Result<(), ShellError> {
+        test_polars_plugin_command_with_decls(command, vec![])
+    }
+
+    pub fn test_polars_plugin_command_with_decls(
+        command: &impl PluginCommand,
+        decls: Vec<Box<dyn Command>>,
+    ) -> Result<(), ShellError> {
         let plugin = PolarsPlugin::new_test_mode();
         let examples = command.examples();
 
@@ -210,9 +218,12 @@ pub mod test {
             }
         }
 
-        PluginTest::new("polars", plugin.into())?
+        let mut plugin_test = PluginTest::new("polars", plugin.into())?;
-            .add_decl(Box::new(IntoDatetime))?
+
-            .test_examples(&examples)?;
+        for decl in decls {
+            let _ = plugin_test.add_decl(decl)?;
+        }
+        plugin_test.test_examples(&examples)?;
 
         Ok(())
     }