Repository: pola-rs/polars
File: crates/polars-io/src/cloud/adaptors.rs (Rust)

Code before:
                Err(e)
            },
        }
    }
}

impl std::io::Write for BlockingCloudWriter {
    fn write(&mut self, buf: &[u8]) -> std::io::Result<usize> {
        // SAFETY:
        // We extend the lifetime for the duration of this function. This is safe as we block the
        // async runtime here
        let buf = unsafe { std::mem::transmute::<&[u8], &'static [u8]>(buf) };

        self.try_with_writer(|writer| {
            get_runtime()
                .block_in_place_on(async { writer.write_all(buf).await.map(|_t| buf.len()) })
        })
    }

    fn flush(&mut self) -> std::io::Result<()> {
        self.try_with_writer(|writer| get_runtime().block_in_place_on(writer.flush()))
    }
}

impl PutOwned for BlockingCloudWriter {
    fn put(&mut self, data: Bytes) -> std::io::Result<()> {
        self.try_with_writer(|writer| {
            let data = data.clone();
            get_runtime().block_in_place_on(async move {
                writer.put(data).await?;
                Ok(())
            })
        })
    }
}

impl WriteableTrait for BlockingCloudWriter {
    fn close(&mut self) -> std::io::Result<()> {
        BlockingCloudWriter::close(self)
    }

    fn sync_all(&self) -> std::io::Result<()> {
        Ok(())
    }

    fn sync_data(&self) -> std::io::Result<()> {
        Ok(())
    }
}

impl Drop for BlockingCloudWriter {

Review comment (suggestion):
Since this is being called from an async function that's already in the tokio runtime, we can just use the object-store inner writer directly (`.try_into_inner()`) instead of going through block_on.


Code after:
        }
    }
}

impl std::io::Write for BlockingCloudWriter {
    fn write(&mut self, buf: &[u8]) -> std::io::Result<usize> {
        // SAFETY:
        // We extend the lifetime for the duration of this function. This is safe as we block the
        // async runtime here
        let buf = unsafe { std::mem::transmute::<&[u8], &'static [u8]>(buf) };

        self.try_with_writer(|writer| {
            get_runtime()
                .block_in_place_on(async { writer.write_all(buf).await.map(|_t| buf.len()) })
        })
    }

    fn flush(&mut self) -> std::io::Result<()> {
        self.try_with_writer(|writer| get_runtime().block_in_place_on(writer.flush()))
    }
}

impl WriteableTrait for BlockingCloudWriter {
    fn close(&mut self) -> std::io::Result<()> {
        BlockingCloudWriter::close(self)
    }

    fn sync_all(&self) -> std::io::Result<()> {
        Ok(())
    }

    fn sync_data(&self) -> std::io::Result<()> {
        Ok(())
    }
}

impl Drop for BlockingCloudWriter {
    fn drop(&mut self) {
        if self.state.is_err() {
            return;
        }

        // Note: We should not hit here - the writer should instead be explicitly closed.
        // But we still have this here as a safety measure to prevent silently dropping errors.
        match self.close() {
            Ok(()) => {},
            e @ Err(_) => {
                if std::thread::panicking() {
                    eprintln!("ERROR: CloudWriter errored on close: {e:?}")
                } else {
                    e.unwrap()