use crate::Result; use crate::cloud::CloudOptions; use crate::config::HazeConfig; use crate::database::DatabaseFamily; use crate::service::ServiceTrait; use bollard::Docker; use futures_util::future::try_join_all; use maplit::hashmap; use miette::Report; use serde_json::Value; use serde_json::json; use std::collections::HashMap; use std::convert::identity; #[derive(Debug, Clone, Eq, PartialEq)] pub struct Sharding; const SHARDS: &[&str] = &["-1", "-2", "-3", "-4"]; #[async_trait::async_trait] impl ServiceTrait for Sharding { fn name(&self) -> &str { "sharding" } async fn spawn( &self, docker: &Docker, cloud_id: &str, network: &str, _config: &HazeConfig, options: &CloudOptions, ) -> Result> { if options.db.family() == DatabaseFamily::Sqlite { return Err(Report::msg("Sharding is not supported with sqlite")); } let containers = try_join_all( SHARDS .iter() .copied() .map(|postfix| options.db.spawn(docker, cloud_id, network, postfix)), ) .await?; Ok(containers.into_iter().flatten().collect()) } async fn is_healthy( &self, docker: &Docker, cloud_id: &str, options: &CloudOptions, ) -> Result { let running = try_join_all( SHARDS .iter() .copied() .map(|postfix| options.db.is_healthy(docker, cloud_id, postfix)), ) .await?; Ok(running.iter().copied().all(identity)) } fn config( &self, _docker: &Docker, _cloud_id: &str, _config: &HazeConfig, ) -> Result> { let shard_config = json!({ "filecache": { "shards": [ { "dbname": "haze", "host": "db-1", "user": "haze", "password": "haze", }, { "dbname": "haze", "host": "db-2", "user": "haze", "password": "haze", }, { "dbname": "haze", "host": "db-3", "user": "haze", "password": "haze", }, { "dbname": "haze", "host": "db-4", "user": "haze", "password": "haze", } ], } }); Ok(hashmap! {String::from("dbsharding") => shard_config}) } } #[derive(Debug, Clone, Eq, PartialEq)] pub struct SingleShard; #[async_trait::async_trait] impl ServiceTrait for SingleShard { fn name(&self) -> &str { "single-shard" } async fn spawn( &self, docker: &Docker, cloud_id: &str, network: &str, _config: &HazeConfig, options: &CloudOptions, ) -> Result> { if options.db.family() == DatabaseFamily::Sqlite { return Err(Report::msg("Sharding is not supported with sqlite")); } let container = options .db .spawn(docker, cloud_id, network, "-shard") .await?; Ok(container.into_iter().collect()) } async fn is_healthy( &self, docker: &Docker, cloud_id: &str, options: &CloudOptions, ) -> Result { options.db.is_healthy(docker, cloud_id, "-shard").await } fn config( &self, _docker: &Docker, _cloud_id: &str, _config: &HazeConfig, ) -> Result> { let shard_config = json!({ "filecache": { "shards": [ { "dbname": "haze", "host": "db-shard", "user": "haze", "password": "haze", } ], } }); Ok(hashmap! {String::from("dbsharding") => shard_config}) } } #[derive(Debug, Clone, Eq, PartialEq)] pub struct ShardingMigrate; #[async_trait::async_trait] impl ServiceTrait for ShardingMigrate { fn name(&self) -> &str { "sharding-migrate" } async fn spawn( &self, docker: &Docker, cloud_id: &str, network: &str, config: &HazeConfig, options: &CloudOptions, ) -> Result> { Sharding::spawn(&Sharding, docker, cloud_id, network, config, options).await } async fn is_healthy( &self, docker: &Docker, cloud_id: &str, options: &CloudOptions, ) -> Result { Sharding::is_healthy(&Sharding, docker, cloud_id, options).await } fn config( &self, _docker: &Docker, _cloud_id: &str, _config: &HazeConfig, ) -> Result> { let shard_config = json!({ "filecache": { "from_shard_key": 99, "from_primary_key": 9999, "shards": [ { "dbname": "haze", "host": "db-1", "user": "haze", "password": "haze", }, { "dbname": "haze", "host": "db-2", "user": "haze", "password": "haze", }, { "dbname": "haze", "host": "db-3", "user": "haze", "password": "haze", }, { "dbname": "haze", "host": "db-4", "user": "haze", "password": "haze", } ], } }); Ok(hashmap! {String::from("dbsharding") => shard_config}) } } #[derive(Debug, Clone, Eq, PartialEq)] pub struct ShardingMigrateUnset; #[async_trait::async_trait] impl ServiceTrait for ShardingMigrateUnset { fn name(&self) -> &str { "sharding-migrate-unset" } async fn spawn( &self, docker: &Docker, cloud_id: &str, network: &str, config: &HazeConfig, options: &CloudOptions, ) -> Result> { Sharding::spawn(&Sharding, docker, cloud_id, network, config, options).await } async fn is_healthy( &self, docker: &Docker, cloud_id: &str, options: &CloudOptions, ) -> Result { Sharding::is_healthy(&Sharding, docker, cloud_id, options).await } fn config( &self, _docker: &Docker, _cloud_id: &str, _config: &HazeConfig, ) -> Result> { let shard_config = json!({ "filecache": { "from_shard_key": 99, "from_primary_key": 9999, "shards": [ { "dbname": "haze", "host": "db-1", "user": "haze", "password": "haze", }, { "dbname": "haze", "host": "db-2", "user": "haze", "password": "haze", }, { "dbname": "haze", "host": "db-3", "user": "haze", "password": "haze", }, { "dbname": "haze", "host": "db-4", "user": "haze", "password": "haze", } ], } }); Ok(hashmap! {String::from("_dbsharding") => shard_config}) } }