From 7ec9d4f2b8fcceb46f4daf64db6eae4cfe4a0356 Mon Sep 17 00:00:00 2001 From: Alek Westover Date: Thu, 13 Jul 2023 13:41:53 -0400 Subject: [PATCH] use tar + archive crates instead --- Cargo.lock | 22 ++++++++++++++- compute_tools/Cargo.toml | 1 + compute_tools/src/extension_server.rs | 28 ++++++------------- .../regress/test_download_extensions.py | 14 +++++----- 4 files changed, 38 insertions(+), 27 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index 0d345c30b2..e44f6deaf3 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -593,7 +593,7 @@ dependencies = [ "cc", "cfg-if", "libc", - "miniz_oxide", + "miniz_oxide 0.6.2", "object", "rustc-demangle", ] @@ -885,6 +885,7 @@ dependencies = [ "chrono", "clap", "compute_api", + "flate2", "futures", "hyper", "notify", @@ -1370,6 +1371,16 @@ version = "0.4.2" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "0ce7134b9999ecaf8bcd65542e436736ef32ddca1b3e06094cb6ec5755203b80" +[[package]] +name = "flate2" +version = "1.0.26" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "3b9429470923de8e8cbd4d2dc513535400b4b3fef0319fb5c4e1f520a7bef743" +dependencies = [ + "crc32fast", + "miniz_oxide 0.7.1", +] + [[package]] name = "fnv" version = "1.0.7" @@ -2154,6 +2165,15 @@ dependencies = [ "adler", ] +[[package]] +name = "miniz_oxide" +version = "0.7.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e7810e0be55b428ada41041c41f32c9f1a42817901b4ccf45fa3d4b6561e74c7" +dependencies = [ + "adler", +] + [[package]] name = "mio" version = "0.8.6" diff --git a/compute_tools/Cargo.toml b/compute_tools/Cargo.toml index 43d122c90d..ad5cda7fc8 100644 --- a/compute_tools/Cargo.toml +++ b/compute_tools/Cargo.toml @@ -32,3 +32,4 @@ utils.workspace = true workspace_hack.workspace = true toml_edit.workspace = true remote_storage = { version = "0.1", path = "../libs/remote_storage/" } +flate2 = "1.0.26" diff --git a/compute_tools/src/extension_server.rs b/compute_tools/src/extension_server.rs index 2c79e92614..f37eee2518 100644 --- a/compute_tools/src/extension_server.rs +++ b/compute_tools/src/extension_server.rs @@ -18,24 +18,19 @@ extensions enabled for specific tenant-ids. use crate::compute::ComputeNode; use anyhow::Context; use anyhow::{self, Result}; +use flate2::read::GzDecoder; use remote_storage::*; use serde_json::{self, Value}; use std::collections::HashSet; -use std::fs::File; -use std::io::BufWriter; -use std::io::Write; use std::num::{NonZeroU32, NonZeroUsize}; use std::path::Path; use std::str; use std::sync::Arc; use std::thread; +use tar::Archive; use tokio::io::AsyncReadExt; use tracing::info; -// TODO: use these crates for untarring, it's better -// use tar::Archive; -// use flate2::read::GzDecoder; - fn get_pg_config(argument: &str, pgbin: &str) -> String { // gives the result of `pg_config [argument]` // where argument is a flag like `--version` or `--sharedir` @@ -169,19 +164,14 @@ pub async fn download_extension( .download_stream .read_to_end(&mut write_data_buffer) .await?; - let zip_name = ext_path.object_name().context("invalid extension path")?; - let mut output_file = BufWriter::new(File::create(zip_name)?); - output_file.write_all(&write_data_buffer)?; - info!("Download {:?} completed successfully", &ext_path); - info!("Unzipping extension to {:?}", zip_name); - std::process::Command::new("tar") - .arg("xzvf") - .arg(zip_name) - .spawn()? - .wait()?; + let unzip_dest = pgbin.strip_suffix("/bin/postgres").expect("bad pgbin"); + let tar = GzDecoder::new(std::io::Cursor::new(write_data_buffer)); + let mut archive = Archive::new(tar); + archive.unpack(unzip_dest)?; + info!("Download + unzip {:?} completed successfully", &ext_path); let local_sharedir = Path::new(&get_pg_config("--sharedir", pgbin)).join("extension"); - let zip_sharedir = format!("extensions/{ext_name}/share/extension"); + let zip_sharedir = format!("{unzip_dest}/extensions/{ext_name}/share/extension"); info!("mv {zip_sharedir:?}/* {local_sharedir:?}"); for file in std::fs::read_dir(zip_sharedir)? { let old_file = file?.path(); @@ -190,7 +180,7 @@ pub async fn download_extension( std::fs::rename(old_file, new_file)?; } let local_libdir = Path::new(&get_pg_config("--libdir", pgbin)).join("postgresql"); - let zip_libdir = format!("extensions/{ext_name}/lib"); + let zip_libdir = format!("{unzip_dest}/extensions/{ext_name}/lib"); info!("mv {zip_libdir:?}/* {local_libdir:?}"); for file in std::fs::read_dir(zip_libdir)? { let old_file = file?.path(); diff --git a/test_runner/regress/test_download_extensions.py b/test_runner/regress/test_download_extensions.py index 0e09a2ba3f..2ce96eb7c0 100644 --- a/test_runner/regress/test_download_extensions.py +++ b/test_runner/regress/test_download_extensions.py @@ -96,17 +96,17 @@ def test_remote_extensions( # Cleaning up downloaded files is important for local tests # or else one test could reuse the files from another test or another test run cleanup_files = [ - "embedding.tar.gz", - "anon.tar.gz", - f"pg_install/v{pg_version}/lib/postgresql/anon.so", - f"pg_install/v{pg_version}/lib/postgresql/embedding.so", - f"pg_install/v{pg_version}/share/postgresql/extension/anon.control", - f"pg_install/v{pg_version}/share/postgresql/extension/embedding--0.1.0.sql", - f"pg_install/v{pg_version}/share/postgresql/extension/embedding.control", + "lib/postgresql/anon.so", + "lib/postgresql/embedding.so", + "share/postgresql/extension/anon.control", + "share/postgresql/extension/embedding--0.1.0.sql", + "share/postgresql/extension/embedding.control", ] + cleanup_files = [f"pg_install/v{pg_version}/" + x for x in cleanup_files] cleanup_folders = [ "extensions", f"pg_install/v{pg_version}/share/postgresql/extension/anon", + f"pg_install/v{pg_version}/extensions", ] for file in cleanup_files: try: