Skip to content

Commit

Permalink
Added support for customizing storage backend
Browse files Browse the repository at this point in the history
  • Loading branch information
Lol3rrr committed Oct 31, 2023
1 parent f299554 commit 0929a3f
Show file tree
Hide file tree
Showing 6 changed files with 214 additions and 43 deletions.
113 changes: 113 additions & 0 deletions Cargo.lock

Some generated files are not rendered by default. Learn more about how customized files appear on GitHub.

1 change: 1 addition & 0 deletions Cargo.toml
Original file line number Diff line number Diff line change
Expand Up @@ -19,3 +19,4 @@ rust_xlsxwriter = { version = "0.44" }
prometheus = { version = "0.13.3" }
aws-creds = { version = "0.36.0", default_features = false, features = ["http-credentials", "rustls-tls"]}
rust-s3 = { version = "0.33.0", default_features = false, features = [ "tokio", "tokio-rustls-tls", "no-verify-ssl"] }
clap = { version = "4.4.7", features = ["derive"] }
62 changes: 62 additions & 0 deletions src/lib.rs
Original file line number Diff line number Diff line change
Expand Up @@ -12,3 +12,65 @@ pub use storage::*;

mod excelstats;
pub use excelstats::ExcelStats;

pub fn parse_storage(args: &str) -> Result<Box<dyn StorageBackend>, &'static str> {
args.split("->")
.filter_map(|arg| match arg {
"file" => {
let store_path =
std::env::var("STORE_PATH").unwrap_or_else(|_| "data.json".to_string());

Some(Box::new(storage::FileStorage::new(store_path)) as Box<dyn StorageBackend>)
}
"s3" => {
let s3_bucket = std::env::var("S3_BUCKET").expect("Missing `S3_BUCKET`");
let s3_access_key =
std::env::var("S3_ACCESS_KEY").expect("Missing `S3_ACCESS_KEY`");
let s3_secret_key =
std::env::var("S3_SECRET_KEY").expect("Missing `S3_SECRET_KEY`");
let s3_endpoint = std::env::var("S3_ENDPOINT").expect("Missing `S3_ENDPOINT`");

Some(Box::new(S3Storage::new(
s3::Bucket::new(
&s3_bucket,
s3::Region::Custom {
region: "default".to_string(),
endpoint: s3_endpoint.to_string(),
},
s3::creds::Credentials::new(
Some(&s3_access_key),
Some(&s3_secret_key),
None,
None,
None,
)
.unwrap(),
)
.unwrap()
.with_path_style(),
)))
}
other => {
tracing::error!("Unknown Storage {:?}", other);

None
}
})
.reduce(
|acc: Box<dyn StorageBackend>, elem: Box<dyn StorageBackend>| {
Box::new(storage::Replicated::new(acc, elem)) as Box<dyn StorageBackend>
},
)
.ok_or("")
}

#[cfg(test)]
mod tests {
use super::*;

#[test]
fn parse_single() {
std::env::set_var("S3", "");
parse_storage("s3").unwrap();
}
}
49 changes: 16 additions & 33 deletions src/main.rs
Original file line number Diff line number Diff line change
Expand Up @@ -6,10 +6,7 @@ use std::env;
use std::sync::Arc;
use std::time::{Duration, SystemTime, UNIX_EPOCH};

use gold_pass_bot::{
ClanTag, ExcelStats, FileStorage, RaidMember, RaidWeekendStats, Replicated, S3Storage, Season,
Storage, StorageBackend,
};
use gold_pass_bot::{ClanTag, ExcelStats, RaidMember, RaidWeekendStats, Season, Storage};
use serenity::async_trait;
use serenity::framework::standard::macros::{command, group};
use serenity::framework::standard::{CommandResult, StandardFramework};
Expand Down Expand Up @@ -41,13 +38,22 @@ async fn main() {
.with(tracing_subscriber::fmt::layer());
tracing::subscriber::set_global_default(layers).unwrap();

let store_path = std::env::var("STORE_PATH").unwrap_or_else(|_| "data.json".to_string());
let api_path = std::env::var("API_PATH").unwrap_or_else(|_| "api.key".to_string());
let args = clap::Command::new("Gold-Pass-Bot")
.arg(
clap::Arg::new("storage")
.long("storage")
.value_names(["storage-target"]),
)
.get_matches();
tracing::debug!("Args: {:#?}", args);

let mut storage_backend = args
.get_one::<String>("storage")
.map(|arg: &String| gold_pass_bot::parse_storage(&arg))
.unwrap()
.unwrap();

let s3_bucket = std::env::var("S3_BUCKET").unwrap();
let s3_access_key = std::env::var("S3_ACCESS_KEY").unwrap();
let s3_secret_key = std::env::var("S3_SECRET_KEY").unwrap();
let s3_endpoint = std::env::var("S3_ENDPOINT").unwrap();
let api_path = std::env::var("API_PATH").unwrap_or_else(|_| "api.key".to_string());

#[cfg(not(debug_assertions))]
let prefix = "!";
Expand All @@ -67,29 +73,6 @@ async fn main() {
.await
.expect("Error creating client");

let mut storage_backend: Box<dyn StorageBackend> = Box::new(Replicated::new(
FileStorage::new(store_path.clone()),
S3Storage::new(
s3::Bucket::new(
&s3_bucket,
s3::Region::Custom {
region: "default".to_string(),
endpoint: s3_endpoint.to_string(),
},
s3::creds::Credentials::new(
Some(&s3_access_key),
Some(&s3_secret_key),
None,
None,
None,
)
.unwrap(),
)
.unwrap()
.with_path_style(),
),
));

let elapsed = SystemTime::now()
.duration_since(UNIX_EPOCH)
.unwrap()
Expand Down
16 changes: 16 additions & 0 deletions src/storage.rs
Original file line number Diff line number Diff line change
Expand Up @@ -26,6 +26,22 @@ pub trait StorageBackend: Send {
fn load(&mut self) -> Pin<Box<dyn Future<Output = Result<Vec<u8>, ()>> + Send + 'static>>;
}

impl<S> StorageBackend for Box<S>
where
S: StorageBackend,
{
fn write(
&mut self,
content: Vec<u8>,
) -> Pin<Box<dyn Future<Output = Result<(), ()>> + Send + 'static>> {
S::write(self.as_mut(), content)
}

fn load(&mut self) -> Pin<Box<dyn Future<Output = Result<Vec<u8>, ()>> + Send + 'static>> {
S::load(self.as_mut())
}
}

#[derive(Debug, Clone, Deserialize, Serialize)]
pub struct Storage {
clans: HashMap<ClanTag, HashMap<Season, ClanStorage>>,
Expand Down
16 changes: 6 additions & 10 deletions src/storage/replicated.rs
Original file line number Diff line number Diff line change
@@ -1,21 +1,17 @@
use crate::StorageBackend;

pub struct Replicated<P, S> {
primary: P,
secondary: S,
pub struct Replicated {
primary: Box<dyn StorageBackend>,
secondary: Box<dyn StorageBackend>,
}

impl<P, S> Replicated<P, S> {
pub fn new(primary: P, secondary: S) -> Self {
impl Replicated {
pub fn new(primary: Box<dyn StorageBackend>, secondary: Box<dyn StorageBackend>) -> Self {
Self { primary, secondary }
}
}

impl<P, S> StorageBackend for Replicated<P, S>
where
P: StorageBackend,
S: StorageBackend,
{
impl StorageBackend for Replicated {
#[tracing::instrument(skip(self, content))]
fn write(
&mut self,
Expand Down

0 comments on commit 0929a3f

Please sign in to comment.