+54
-50
src/main.rs
+54
-50
src/main.rs
···
1
-
use std::{env, fs};
2
use std::path::PathBuf;
3
4
use clap::{Parser, Subcommand};
5
use dotenvy::dotenv;
6
use log::{error, info};
7
use reqwest::header::{ACCEPT, ACCEPT_ENCODING};
8
-
use anyhow;
9
use tokio_util::compat::{FuturesAsyncReadCompatExt, TokioAsyncReadCompatExt};
10
use tokio_util::io::StreamReader;
11
-
use tokio_stream::wrappers::LinesStream;
12
-
use tokio::io::AsyncWriteExt;
13
-
use std::future::Future;
14
-
use async_compression::futures::bufread::{ZstdEncoder, ZstdDecoder};
15
-
use s3::{Bucket, Region};
16
-
use s3::creds::Credentials;
17
-
use futures::io::BufReader;
18
19
#[derive(Parser, Debug)]
20
#[command(author, version, about = "PDS Whatsit Compress Test CLI", long_about = None)]
···
63
fs::create_dir_all(parent)?;
64
}
65
}
66
-
if !path.exists() {
67
-
fs::create_dir_all(path)?;
68
-
}
69
Ok(())
70
}
71
72
#[tokio::main]
73
-
async fn main() {
74
init_logging();
75
76
let cli = Cli::parse();
77
78
-
match cli.command {
79
-
Commands::Export { did, pds_url } => {
80
81
-
info!(
82
-
"Export requested: did={}, pds_url={}",
83
-
did,
84
-
pds_url,
85
-
);
86
87
-
match do_work(pds_url, did).await{
88
Ok(_) => {
89
info!("Export completed");
90
}
91
Err(err) => {
92
error!("Export failed: {}", err);
93
}
94
}
95
}
96
Commands::Import { did } => {
97
info!("Import requested: did={}", did);
98
-
match do_import(did).await {
99
-
Ok(path) => info!("Import completed, wrote {}", path.display()),
100
-
Err(err) => error!("Import failed: {}", err),
101
}
102
}
103
}
104
}
105
106
-
async fn do_work(pds_url: String, did: String) -> anyhow::Result<()>{
107
use futures::TryStreamExt;
108
let atproto_client = reqwest::Client::new();
109
let back_up_path = format!("users/{did}");
110
111
-
// Custom region requires valid region name and endpoint
112
-
let region_name = env::var("S3_REGION")?;
113
-
let endpoint = env::var("S3_ENDPOINT")?;
114
-
let region = Region::Custom { region: region_name, endpoint };
115
-
116
-
let bucket = Bucket::new(
117
-
env::var("S3_BUCKET_NAME")?.as_str(),
118
-
region,
119
-
// Credentials are collected from environment, config, profile or instance metadata
120
-
Credentials::new(Some(env::var("S3_ACCESS_KEY")?.as_str()), Some(env::var("S3_SECRET_KEY")?.as_str()), None, None, None)?,
121
-
)?;
122
-
123
let response_reader = atproto_client
124
.get(format!("{pds_url}/xrpc/com.atproto.sync.getRepo?did={did}"))
125
.header(ACCEPT, "application/vnd.ipld.car")
···
150
Ok(())
151
}
152
153
-
async fn do_import(did: String) -> anyhow::Result<PathBuf> {
154
use futures::StreamExt;
155
156
let back_up_path = format!("users/{did}");
157
158
-
// S3 setup from env
159
-
let region_name = env::var("S3_REGION")?;
160
-
let endpoint = env::var("S3_ENDPOINT")?;
161
-
let region = Region::Custom { region: region_name, endpoint };
162
-
163
-
let bucket = Bucket::new(
164
-
env::var("S3_BUCKET_NAME")?.as_str(),
165
-
region,
166
-
Credentials::new(Some(env::var("S3_ACCESS_KEY")?.as_str()), Some(env::var("S3_SECRET_KEY")?.as_str()), None, None, None)?,
167
-
)?;
168
-
169
// Stream download from S3
170
let mut s3_stream = bucket
171
.get_object_stream(format!("/{back_up_path}/{did}.car.zst"))
···
189
let mut decoded_tokio_reader = decoder.compat();
190
191
// Prepare local output path, labeled as decompressed
192
-
let out_path: PathBuf = ["export", "users", did.as_str(), &format!("{}-decompressed.car", did)].iter().collect();
193
ensure_dir(&out_path)?;
194
195
let mut out_file = tokio::fs::File::create(&out_path).await?;
···
1
use std::path::PathBuf;
2
+
use std::{env, fs};
3
4
+
use anyhow;
5
+
use async_compression::futures::bufread::{ZstdDecoder, ZstdEncoder};
6
use clap::{Parser, Subcommand};
7
use dotenvy::dotenv;
8
+
use futures::io::BufReader;
9
use log::{error, info};
10
use reqwest::header::{ACCEPT, ACCEPT_ENCODING};
11
+
use s3::creds::Credentials;
12
+
use s3::{Bucket, Region};
13
+
use std::future::Future;
14
+
use tokio::io::AsyncWriteExt;
15
+
use tokio_stream::wrappers::LinesStream;
16
use tokio_util::compat::{FuturesAsyncReadCompatExt, TokioAsyncReadCompatExt};
17
use tokio_util::io::StreamReader;
18
19
#[derive(Parser, Debug)]
20
#[command(author, version, about = "PDS Whatsit Compress Test CLI", long_about = None)]
···
63
fs::create_dir_all(parent)?;
64
}
65
}
66
Ok(())
67
}
68
69
#[tokio::main]
70
+
async fn main() -> anyhow::Result<()> {
71
init_logging();
72
73
let cli = Cli::parse();
74
+
// Custom region requires valid region name and endpoint
75
+
let region_name = env::var("S3_REGION")?;
76
+
let endpoint = env::var("S3_ENDPOINT")?;
77
+
let region = Region::Custom {
78
+
region: region_name,
79
+
endpoint,
80
+
};
81
82
+
let bucket = Bucket::new(
83
+
env::var("S3_BUCKET_NAME")?.as_str(),
84
+
region,
85
+
// Credentials are collected from environment, config, profile or instance metadata
86
+
Credentials::new(
87
+
Some(env::var("S3_ACCESS_KEY")?.as_str()),
88
+
Some(env::var("S3_SECRET_KEY")?.as_str()),
89
+
None,
90
+
None,
91
+
None,
92
+
)?,
93
+
)?;
94
95
+
match cli.command {
96
+
Commands::Export { did, pds_url } => {
97
+
info!("Export requested: did={}, pds_url={}", did, pds_url,);
98
99
+
match do_work(pds_url, did, bucket).await {
100
Ok(_) => {
101
info!("Export completed");
102
+
Ok(())
103
}
104
Err(err) => {
105
error!("Export failed: {}", err);
106
+
Err(err)
107
}
108
}
109
}
110
Commands::Import { did } => {
111
info!("Import requested: did={}", did);
112
+
match do_import(did, bucket).await {
113
+
Ok(path) => {
114
+
info!("Import completed, wrote {}", path.display());
115
+
Ok(())
116
+
}
117
+
Err(err) => {
118
+
error!("Import failed: {}", err);
119
+
Err(err)
120
+
}
121
}
122
}
123
}
124
}
125
126
+
async fn do_work(pds_url: String, did: String, bucket: Box<Bucket>) -> anyhow::Result<()> {
127
use futures::TryStreamExt;
128
let atproto_client = reqwest::Client::new();
129
let back_up_path = format!("users/{did}");
130
131
let response_reader = atproto_client
132
.get(format!("{pds_url}/xrpc/com.atproto.sync.getRepo?did={did}"))
133
.header(ACCEPT, "application/vnd.ipld.car")
···
158
Ok(())
159
}
160
161
+
async fn do_import(did: String, bucket: Box<Bucket>) -> anyhow::Result<PathBuf> {
162
use futures::StreamExt;
163
164
let back_up_path = format!("users/{did}");
165
166
// Stream download from S3
167
let mut s3_stream = bucket
168
.get_object_stream(format!("/{back_up_path}/{did}.car.zst"))
···
186
let mut decoded_tokio_reader = decoder.compat();
187
188
// Prepare local output path, labeled as decompressed
189
+
let out_path: PathBuf = [
190
+
"export",
191
+
"users",
192
+
did.as_str(),
193
+
&format!("{}-decompressed.car", did),
194
+
]
195
+
.iter()
196
+
.collect();
197
ensure_dir(&out_path)?;
198
199
let mut out_file = tokio::fs::File::create(&out_path).await?;