mirror of
https://github.com/nushell/nushell.git
synced 2025-05-11 05:24:27 +02:00
Fix #15571 panic on write to source parquet file
This commit is contained in:
parent
78903724f5
commit
a7ac673d3d
@ -9,8 +9,8 @@ use nu_utils::perf;
|
||||
|
||||
use nu_plugin::{EvaluatedCall, PluginCommand};
|
||||
use nu_protocol::{
|
||||
shell_error::io::IoError, Category, Example, LabeledError, PipelineData, ShellError, Signature,
|
||||
Span, Spanned, SyntaxShape, Type, Value,
|
||||
shell_error::io::IoError, Category, DataSource, Example, LabeledError, PipelineData,
|
||||
PipelineMetadata, ShellError, Signature, Span, Spanned, SyntaxShape, Type, Value,
|
||||
};
|
||||
|
||||
use std::{fs::File, io::BufReader, num::NonZeroUsize, path::PathBuf, sync::Arc};
|
||||
@ -164,6 +164,8 @@ fn command(
|
||||
}
|
||||
|
||||
let hive_options = build_hive_options(plugin, call)?;
|
||||
let metadata = PipelineMetadata::default()
|
||||
.with_data_source(DataSource::FilePath(spanned_file.item.clone().into()));
|
||||
|
||||
match type_option {
|
||||
Some((ext, blamed)) => match PolarsFileType::from(ext.as_str()) {
|
||||
@ -199,7 +201,7 @@ fn command(
|
||||
"File without extension",
|
||||
))),
|
||||
}
|
||||
.map(|value| PipelineData::Value(value, None))
|
||||
.map(|value| PipelineData::Value(value, Some(metadata)))
|
||||
}
|
||||
|
||||
fn from_parquet(
|
||||
|
@ -15,8 +15,8 @@ use crate::{
|
||||
use log::debug;
|
||||
use nu_plugin::{EngineInterface, EvaluatedCall, PluginCommand};
|
||||
use nu_protocol::{
|
||||
shell_error::io::IoError, Category, Example, LabeledError, PipelineData, ShellError, Signature,
|
||||
Span, Spanned, SyntaxShape, Type,
|
||||
shell_error::io::IoError, Category, DataSource, Example, LabeledError, PipelineData,
|
||||
PipelineMetadata, ShellError, Signature, Span, Spanned, SyntaxShape, Type,
|
||||
};
|
||||
use polars::error::PolarsError;
|
||||
|
||||
@ -112,11 +112,20 @@ impl PluginCommand for SaveDF {
|
||||
call: &EvaluatedCall,
|
||||
input: PipelineData,
|
||||
) -> Result<PipelineData, LabeledError> {
|
||||
let spanned_file: Spanned<String> = call.req(0)?;
|
||||
debug!("file: {}", spanned_file.item);
|
||||
|
||||
let metadata = input.metadata();
|
||||
let value = input.into_value(call.head)?;
|
||||
|
||||
check_writing_into_source_file(
|
||||
metadata.as_ref(),
|
||||
&spanned_file.as_ref().map(PathBuf::from),
|
||||
)?;
|
||||
|
||||
match PolarsPluginObject::try_from_value(plugin, &value)? {
|
||||
po @ PolarsPluginObject::NuDataFrame(_) | po @ PolarsPluginObject::NuLazyFrame(_) => {
|
||||
command(plugin, engine, call, po)
|
||||
command(plugin, engine, call, po, spanned_file)
|
||||
}
|
||||
_ => Err(cant_convert_err(
|
||||
&value,
|
||||
@ -132,10 +141,8 @@ fn command(
|
||||
engine: &EngineInterface,
|
||||
call: &EvaluatedCall,
|
||||
polars_object: PolarsPluginObject,
|
||||
spanned_file: Spanned<String>,
|
||||
) -> Result<PipelineData, ShellError> {
|
||||
let spanned_file: Spanned<String> = call.req(0)?;
|
||||
debug!("file: {}", spanned_file.item);
|
||||
|
||||
let resource = Resource::new(plugin, engine, &spanned_file)?;
|
||||
let type_option: Option<(String, Span)> = call
|
||||
.get_flag("type")?
|
||||
@ -223,6 +230,28 @@ fn command(
|
||||
Ok(PipelineData::empty())
|
||||
}
|
||||
|
||||
fn check_writing_into_source_file(
|
||||
metadata: Option<&PipelineMetadata>,
|
||||
dest: &Spanned<PathBuf>,
|
||||
) -> Result<(), ShellError> {
|
||||
let Some(DataSource::FilePath(source)) = metadata.map(|meta| &meta.data_source) else {
|
||||
return Ok(());
|
||||
};
|
||||
|
||||
if &dest.item == source {
|
||||
return Err(write_into_source_error(dest.span));
|
||||
}
|
||||
|
||||
Ok(())
|
||||
}
|
||||
|
||||
fn write_into_source_error(span: Span) -> ShellError {
|
||||
polars_file_save_error(
|
||||
PolarsError::InvalidOperation("attempted to save into source".into()),
|
||||
span,
|
||||
)
|
||||
}
|
||||
|
||||
pub(crate) fn polars_file_save_error(e: PolarsError, span: Span) -> ShellError {
|
||||
ShellError::GenericError {
|
||||
error: format!("Error saving file: {e}"),
|
||||
@ -247,17 +276,13 @@ pub fn unknown_file_save_error(span: Span) -> ShellError {
|
||||
pub(crate) mod test {
|
||||
use nu_plugin_test_support::PluginTest;
|
||||
use nu_protocol::{Span, Value};
|
||||
use tempfile::TempDir;
|
||||
use uuid::Uuid;
|
||||
|
||||
use crate::PolarsPlugin;
|
||||
|
||||
fn test_save(cmd: &'static str, extension: &str) -> Result<(), Box<dyn std::error::Error>> {
|
||||
fn tmp_dir_sandbox() -> Result<(TempDir, PluginTest), Box<dyn std::error::Error>> {
|
||||
let tmp_dir = tempfile::tempdir()?;
|
||||
let mut tmp_file = tmp_dir.path().to_owned();
|
||||
tmp_file.push(format!("{}.{}", Uuid::new_v4(), extension));
|
||||
let tmp_file_str = tmp_file.to_str().expect("should be able to get file path");
|
||||
|
||||
let cmd = format!("{cmd} {tmp_file_str}");
|
||||
let mut plugin_test = PluginTest::new("polars", PolarsPlugin::new()?.into())?;
|
||||
plugin_test.engine_state_mut().add_env_var(
|
||||
"PWD".to_string(),
|
||||
@ -270,6 +295,17 @@ pub(crate) mod test {
|
||||
Span::test_data(),
|
||||
),
|
||||
);
|
||||
|
||||
Ok((tmp_dir, plugin_test))
|
||||
}
|
||||
|
||||
fn test_save(cmd: &'static str, extension: &str) -> Result<(), Box<dyn std::error::Error>> {
|
||||
let (tmp_dir, mut plugin_test) = tmp_dir_sandbox()?;
|
||||
let mut tmp_file = tmp_dir.path().to_owned();
|
||||
tmp_file.push(format!("{}.{}", Uuid::new_v4(), extension));
|
||||
let tmp_file_str = tmp_file.to_str().expect("should be able to get file path");
|
||||
|
||||
let cmd = format!("{cmd} {tmp_file_str}");
|
||||
let _pipeline_data = plugin_test.eval(&cmd)?;
|
||||
|
||||
assert!(tmp_file.exists());
|
||||
@ -290,4 +326,27 @@ pub(crate) mod test {
|
||||
extension,
|
||||
)
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_write_to_source_guard() -> Result<(), Box<dyn std::error::Error>> {
|
||||
let (tmp_dir, mut plugin_test) = tmp_dir_sandbox()?;
|
||||
let mut tmp_file = tmp_dir.path().to_owned();
|
||||
dbg!(&tmp_dir);
|
||||
tmp_file.push(format!("{}.{}", Uuid::new_v4(), "parquet"));
|
||||
let tmp_file_str = tmp_file.to_str().expect("Should be able to get file path");
|
||||
|
||||
let _setup = plugin_test.eval(&format!(
|
||||
"[1 2 3] | polars into-df | polars save {tmp_file_str}",
|
||||
))?;
|
||||
|
||||
let output = plugin_test.eval(&format!(
|
||||
"polars open {tmp_file_str} | polars save {tmp_file_str}"
|
||||
));
|
||||
|
||||
assert!(output.is_err_and(|e| e
|
||||
.to_string()
|
||||
.contains("Error saving file: attempted to save into source")));
|
||||
|
||||
Ok(())
|
||||
}
|
||||
}
|
||||
|
Loading…
Reference in New Issue
Block a user