From 79459e8c0a04457f83e9217ecf2f5d09479a71a1 Mon Sep 17 00:00:00 2001 From: Alek Westover Date: Tue, 13 Jun 2023 14:51:00 -0400 Subject: [PATCH] realized that we actually need to add a new function to remote_storage for listing files. working on it --- alek_ext/Cargo.lock | 40 ++---------- alek_ext/Cargo.toml | 9 ++- .../alek1.out | 0 alek_ext/alek2.out | 7 ++ ..._api.rs => download_with_remote_api_v2.rs} | 0 alek_ext/src/main.rs | 60 +++++++++-------- alek_ext/src/real_main.rs | 65 +++++++++++++++++++ .../fuzzystrmatch.control} | 0 fakes3/public_extensions/pg_cron.control | 7 ++ fakes3/tenant123/test.control | 3 + libs/remote_storage/src/lib.rs | 23 +++++++ libs/remote_storage/src/local_fs.rs | 30 +++++++++ libs/remote_storage/src/s3_bucket.rs | 33 ++++++++++ libs/remote_storage/src/simulate_failures.rs | 8 +++ 14 files changed, 221 insertions(+), 64 deletions(-) rename fakes3/neon-dev-extensions/fuzzystrmatch.control => alek_ext/alek1.out (100%) create mode 100644 alek_ext/alek2.out rename alek_ext/src/{download_with_remote_api.rs => download_with_remote_api_v2.rs} (100%) create mode 100644 alek_ext/src/real_main.rs rename fakes3/{neon-dev-extensions/pg_cron.control => public_extensions/fuzzystrmatch.control} (100%) create mode 100644 fakes3/public_extensions/pg_cron.control create mode 100644 fakes3/tenant123/test.control diff --git a/alek_ext/Cargo.lock b/alek_ext/Cargo.lock index 2a011e82c7..87ccb5f3d8 100644 --- a/alek_ext/Cargo.lock +++ b/alek_ext/Cargo.lock @@ -31,13 +31,16 @@ name = "alek_ext" version = "0.1.0" dependencies = [ "anyhow", + "aws-config", "aws-credential-types", - "aws-sdk-s3 0.28.0", + "aws-sdk-s3", "aws-smithy-http", + "aws-types", "remote_storage", "tokio", "toml_edit", "tracing", + "tracing-subscriber", ] [[package]] @@ -256,39 +259,6 @@ dependencies = [ "url", ] -[[package]] -name = "aws-sdk-s3" -version = "0.28.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "fba197193cbb4bcb6aad8d99796b2291f36fa89562ded5d4501363055b0de89f" -dependencies = [ - "aws-credential-types", - "aws-endpoint", - "aws-http", - "aws-sig-auth", - "aws-sigv4", - "aws-smithy-async", - "aws-smithy-checksums", - "aws-smithy-client", - "aws-smithy-eventstream", - "aws-smithy-http", - "aws-smithy-http-tower", - "aws-smithy-json", - "aws-smithy-types", - "aws-smithy-xml", - "aws-types", - "bytes", - "http", - "http-body", - "once_cell", - "percent-encoding", - "regex", - "tokio-stream", - "tower", - "tracing", - "url", -] - [[package]] name = "aws-sdk-sts" version = "0.28.0" @@ -1828,7 +1798,7 @@ dependencies = [ "async-trait", "aws-config", "aws-credential-types", - "aws-sdk-s3 0.27.0", + "aws-sdk-s3", "aws-smithy-http", "aws-types", "hyper", diff --git a/alek_ext/Cargo.toml b/alek_ext/Cargo.toml index 5ac9811953..f84c39f7cb 100644 --- a/alek_ext/Cargo.toml +++ b/alek_ext/Cargo.toml @@ -7,12 +7,15 @@ edition = "2021" [dependencies] anyhow = "1.0.71" -aws-credential-types = "0.55.3" -aws-sdk-s3 = "0.28.0" -aws-smithy-http = "0.55.3" +aws-config = { version = "0.55", default-features = false, features=["rustls"] } +aws-sdk-s3 = "0.27" +aws-smithy-http = "0.55" +aws-credential-types = "0.55" +aws-types = "0.55" remote_storage = { version = "0.1", path = "../libs/remote_storage/" } tokio = "1.28.2" toml_edit = "0.19.10" tracing = "0.1.37" +tracing-subscriber = "0.3.17" [workspace] diff --git a/fakes3/neon-dev-extensions/fuzzystrmatch.control b/alek_ext/alek1.out similarity index 100% rename from fakes3/neon-dev-extensions/fuzzystrmatch.control rename to alek_ext/alek1.out diff --git a/alek_ext/alek2.out b/alek_ext/alek2.out new file mode 100644 index 0000000000..6447eb8ef9 --- /dev/null +++ b/alek_ext/alek2.out @@ -0,0 +1,7 @@ +this +ax mod p < p/2 + +its gonna be big! +coming soon at 4pm +lol + diff --git a/alek_ext/src/download_with_remote_api.rs b/alek_ext/src/download_with_remote_api_v2.rs similarity index 100% rename from alek_ext/src/download_with_remote_api.rs rename to alek_ext/src/download_with_remote_api_v2.rs diff --git a/alek_ext/src/main.rs b/alek_ext/src/main.rs index cc7dca98d6..353da50dc3 100644 --- a/alek_ext/src/main.rs +++ b/alek_ext/src/main.rs @@ -1,6 +1,13 @@ /* **WIP** * This is a MWE of using our RemoteStorage API to call the aws stuff and download multiple files + * + * TODO: + * fix local_fs and s3_bucket list_files functions. specifically + * 1. local_fs mostly just clean up the error handling or whatever + * 2. s3_bucket needs a lot of work; for example to make sure the pagination thing goes fine. + * + * 3. figure out how to write dirrectly from the buffer to a local file */ use remote_storage::*; @@ -8,48 +15,49 @@ use std::path::Path; use std::fs::File; use std::io::{BufWriter, Write}; use toml_edit; -use anyhow; -use tokio::io::AsyncReadExt; +use anyhow::{self, Context}; +use tokio::io::AsyncReadExt; +use tracing::*; +use tracing_subscriber; #[tokio::main] async fn main() -> anyhow::Result<()> { - let from_path = "fuzzystrmatch.control"; - let remote_from_path = RemotePath::new(Path::new(from_path))?; - println!("{:?}", remote_from_path.clone()); + let subscriber = tracing_subscriber::FmtSubscriber::new(); + tracing::subscriber::set_global_default(subscriber)?; - // read configurations from `pageserver.toml` let cfg_file_path = Path::new("./../.neon/pageserver.toml"); let cfg_file_contents = std::fs::read_to_string(cfg_file_path).unwrap(); let toml = cfg_file_contents .parse::() - .expect("Error parsing toml"); + .context("Error parsing toml")?; let remote_storage_data = toml.get("remote_storage") .expect("field should be present"); let remote_storage_config = RemoteStorageConfig::from_toml(remote_storage_data) .expect("error parsing toml") .expect("error parsing toml"); - - // query S3 bucket + info!("{:?}", remote_storage_config); let remote_storage = GenericRemoteStorage::from_config(&remote_storage_config)?; - let from_path = "fuzzystrmatch.control"; - let remote_from_path = RemotePath::new(Path::new(from_path))?; - - println!("{:?}", remote_from_path.clone()); - // if let GenericRemoteStorage::AwsS3(mybucket) = remote_storage { - // println!("{:?}",mybucket.relative_path_to_s3_object(&remote_from_path)); - // } - let mut data = remote_storage.download(&remote_from_path).await.expect("data yay"); - let mut write_data_buffer = Vec::new(); - data.download_stream.read_to_end(&mut write_data_buffer).await?; - // write `data` to a file locally - let f = File::create("alek.out").expect("problem creating file"); - let mut f = BufWriter::new(f); - f.write_all(&mut write_data_buffer).expect("error writing data"); + let folder = RemotePath::new(Path::new("public_extensions"))?; + let from_paths = remote_storage.list_files(Some(&folder)).await?; - // let stuff = response.body; - // let data = stuff.collect().await.expect("error reading data").to_vec(); - // println!("data: {:?}", std::str::from_utf8(&data)); + let mut i = 0; + for remote_from_path in from_paths { + i += 1; + println!("number {i}{:?}", remote_from_path.object_name()); + if remote_from_path.extension() == Some("control") { + let mut data = remote_storage.download(&remote_from_path).await?; + // write `data` to a file locally + // TODO: I think that the way I'm doing this is not optimal; + // It should be possible to write the data directly to a file + // rather than first writing it to a vector... + let mut write_data_buffer = Vec::new(); + data.download_stream.read_to_end(&mut write_data_buffer).await?; + let f = File::create(format!("alek{i}.out")).expect("problem creating file"); + let mut f = BufWriter::new(f); + f.write_all(&mut write_data_buffer).expect("error writing data"); + } + } Ok(()) } diff --git a/alek_ext/src/real_main.rs b/alek_ext/src/real_main.rs new file mode 100644 index 0000000000..36bb745ebc --- /dev/null +++ b/alek_ext/src/real_main.rs @@ -0,0 +1,65 @@ +/* +**WIP** + * This is a MWE of using our RemoteStorage API to call the aws stuff and download multiple files +*/ + +#![allow(unused_imports)] +use remote_storage::*; +use std::path::Path; +use std::fs::File; +use std::io::{BufWriter, Write}; +use toml_edit; +use anyhow::{self, Context}; +use tokio::io::AsyncReadExt; + +#[tokio::main] +async fn main() -> anyhow::Result<()> { + /* me trying to hack RemotePath into submission */ + + let cfg_file_path = Path::new("./../.neon/pageserver.toml"); + let cfg_file_contents = std::fs::read_to_string(cfg_file_path) + .expect("couldn't find pageserver.toml; make sure you are in neon/alek_ext"); + let toml = cfg_file_contents + .parse::() + .expect("Error parsing toml"); + let remote_storage_data = toml.get("remote_storage") + .expect("field should be present"); + let remote_storage_config = RemoteStorageConfig::from_toml(remote_storage_data) + .expect("error parsing toml") + .expect("error parsing toml"); + let remote_storage = GenericRemoteStorage::from_config(&remote_storage_config)?; + + if let GenericRemoteStorage::AwsS3(mybucket) = remote_storage { + let resp = mybucket + .client + .list_objects_v2() + .bucket("neon-dev-extensions") + .set_prefix(Some("public_extensions".to_string())) + .delimiter("/".to_string()) + .send().await?; + + let z = resp.common_prefixes.unwrap(); + for yy in z { + println!("plzplz: {:?}",yy); + } + let mut i = 0; + for remote_from_path in from_paths { + i += 1; + println!("{:?}", &remote_from_path); + if remote_from_path.extension() == Some("control") { + let mut data = remote_storage.download(&remote_from_path).await?; + // write `data` to a file locally + // TODO: I think that the way I'm doing this is not optimal; + // It should be possible to write the data directly to a file + // rather than first writing it to a vector... + let mut write_data_buffer = Vec::new(); + data.download_stream.read_to_end(&mut write_data_buffer).await?; + let f = File::create("alek{i}.out").expect("problem creating file"); + let mut f = BufWriter::new(f); + f.write_all(&mut write_data_buffer).expect("error writing data"); + } + } + } + + Ok(()) +} diff --git a/fakes3/neon-dev-extensions/pg_cron.control b/fakes3/public_extensions/fuzzystrmatch.control similarity index 100% rename from fakes3/neon-dev-extensions/pg_cron.control rename to fakes3/public_extensions/fuzzystrmatch.control diff --git a/fakes3/public_extensions/pg_cron.control b/fakes3/public_extensions/pg_cron.control new file mode 100644 index 0000000000..6447eb8ef9 --- /dev/null +++ b/fakes3/public_extensions/pg_cron.control @@ -0,0 +1,7 @@ +this +ax mod p < p/2 + +its gonna be big! +coming soon at 4pm +lol + diff --git a/fakes3/tenant123/test.control b/fakes3/tenant123/test.control new file mode 100644 index 0000000000..3ecf3e5a09 --- /dev/null +++ b/fakes3/tenant123/test.control @@ -0,0 +1,3 @@ + +Lorem ipsum dolor sit amet, consetetur sadipscing elitr, sed diam nonumy eirmod tempor invidunt ut labore et dolore magna aliquyam erat, sed diam voluptua. At vero eos et accusam et justo duo dolores et ea rebum. Stet clita kasd gubergren, no sea takimata sanctus est Lorem ipsum dolor sit amet. + diff --git a/libs/remote_storage/src/lib.rs b/libs/remote_storage/src/lib.rs index ec52448025..a02343cda8 100644 --- a/libs/remote_storage/src/lib.rs +++ b/libs/remote_storage/src/lib.rs @@ -90,6 +90,12 @@ pub trait RemoteStorage: Send + Sync + 'static { prefix: Option<&RemotePath>, ) -> Result, DownloadError>; + /// Lists all files in a subdirectories + async fn list_files( + &self, + prefix: Option<&RemotePath>, + ) -> anyhow::Result>; + /// Streams the local file contents into remote into the remote storage entry. async fn upload( &self, @@ -165,6 +171,23 @@ pub enum GenericRemoteStorage { } impl GenericRemoteStorage { + // A function for listing all the files in a "directory" + // Example: + // list_files("foo/bar") = ["foo/bar/a.txt", "foo/bar/b.txt"] + pub async fn list_files( + &self, + folder: Option<&RemotePath> + ) -> anyhow::Result>{ + match self { + Self::LocalFs(s) => s.list_files(folder).await, + Self::AwsS3(s) => s.list_files(folder).await, + Self::Unreliable(s) => s.list_files(folder).await, + } + } + + // lists common *prefixes*, if any of files + // Example: + // list_prefixes("foo123","foo567","bar123","bar432") = ["foo", "bar"] pub async fn list_prefixes( &self, prefix: Option<&RemotePath>, diff --git a/libs/remote_storage/src/local_fs.rs b/libs/remote_storage/src/local_fs.rs index c081a6d361..acc9c4a004 100644 --- a/libs/remote_storage/src/local_fs.rs +++ b/libs/remote_storage/src/local_fs.rs @@ -116,6 +116,36 @@ impl RemoteStorage for LocalFs { .collect()) } + async fn list_files( + &self, + folder: Option<&RemotePath> + ) -> anyhow::Result> { + /* Note: if you want, you can return a DownloadError instead of an anyhow::Error + as follows: replace all ?'s with: + .map_err(|e| DownloadError::Other(anyhow::Error::from(e)))?; + */ + let full_path = match folder.clone() { + Some(folder) => folder.with_base(&self.storage_root), + None => self.storage_root.clone(), + }; + println!("{:?}", full_path); + let mut entries = fs::read_dir(full_path).await?; + let mut files = vec![]; + + while let Some(entry) = entries.next_entry().await? { + let file_name: PathBuf = entry.file_name().into(); + let file_type = entry.file_type().await?; + if file_type.is_file() { + let mut file_remote_path = RemotePath::new(&file_name)?; + if let Some(folder) = folder { + file_remote_path = folder.join(&file_name); + } + files.push(file_remote_path); + } + } + Ok(files) + } + async fn upload( &self, data: impl io::AsyncRead + Unpin + Send + Sync + 'static, diff --git a/libs/remote_storage/src/s3_bucket.rs b/libs/remote_storage/src/s3_bucket.rs index 0be8c72fe0..5dd380c0a2 100644 --- a/libs/remote_storage/src/s3_bucket.rs +++ b/libs/remote_storage/src/s3_bucket.rs @@ -5,6 +5,7 @@ //! their bucket prefixes are both specified and different. use std::sync::Arc; +use std::path::Path; use anyhow::Context; use aws_config::{ @@ -332,6 +333,38 @@ impl RemoteStorage for S3Bucket { Ok(document_keys) } + async fn list_files( + &self, + folder: Option<&RemotePath> + ) -> anyhow::Result>{ + + // THIS IS A somewhat messed up prototype + + let folder_name = folder.map(|x| + String::from( + x.object_name().expect("something went wrong computing folder name") + ) + ); + + // TODO: worry about continuation character and limits and permit and stuff + let resp = self + .client + .list_objects_v2() + .bucket("neon-dev-extensions") + .set_prefix(folder_name) + .send().await + .unwrap(); + + let mut all_files = vec![]; + for object in resp.contents().unwrap() { + let x = object.key().unwrap(); + let y = RemotePath::new(Path::new(x)).unwrap(); + all_files.push(y); + } + + Ok(all_files) + } + async fn upload( &self, from: impl io::AsyncRead + Unpin + Send + Sync + 'static, diff --git a/libs/remote_storage/src/simulate_failures.rs b/libs/remote_storage/src/simulate_failures.rs index cb40859831..bd1841ae21 100644 --- a/libs/remote_storage/src/simulate_failures.rs +++ b/libs/remote_storage/src/simulate_failures.rs @@ -82,6 +82,14 @@ impl RemoteStorage for UnreliableWrapper { self.inner.list_prefixes(prefix).await } + async fn list_files( + &self, + folder: Option<&RemotePath> + ) -> anyhow::Result>{ + self.attempt(RemoteOp::ListPrefixes(folder.cloned()))?; + self.inner.list_files(folder).await + } + async fn upload( &self, data: impl tokio::io::AsyncRead + Unpin + Send + Sync + 'static,