diff --git a/Cargo.lock b/Cargo.lock index 7755517..d4b853d 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -10,6 +10,7 @@ dependencies = [ "clap", "clap_complete", "dotenvy", + "serde_json", ] [[package]] @@ -141,6 +142,18 @@ version = "1.70.2" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "a6cb138bb79a146c1bd460005623e142ef0181e3d0219cb493e02f7d08a35695" +[[package]] +name = "itoa" +version = "1.0.18" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "8f42a60cbdf9a97f5d2305f08a87dc4e09308d1276d28c869c684d7777685682" + +[[package]] +name = "memchr" +version = "2.8.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "cf8baf1c55e62ffcace7a9f06f4bd9cd3f0c4beb022d3b367256b91b87513d98" + [[package]] name = "once_cell_polyfill" version = "1.70.2" @@ -165,6 +178,48 @@ dependencies = [ "proc-macro2", ] +[[package]] +name = "serde" +version = "1.0.229" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "4148590afebada386688f18773da617792bf2ef03ffc1e4cbd2b1d45b023e0ba" +dependencies = [ + "serde_core", +] + +[[package]] +name = "serde_core" +version = "1.0.229" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "67dca2c9c51e58a4791a4b1ed58308b39c64224d349a935ab5039aa360942a48" +dependencies = [ + "serde_derive", +] + +[[package]] +name = "serde_derive" +version = "1.0.229" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e7a5d71263a5a7d47b41f6b3f06ba276f10cc18b0931f1799f710578e2309348" +dependencies = [ + "proc-macro2", + "quote", + "syn", +] + +[[package]] +name = "serde_json" +version = "1.0.151" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c841b55ecdae098c80dcae9cf767f6f8a0c2cdb3416bbef72181df4d0fe73f14" +dependencies = [ + "itoa", + "memchr", + "serde", + "serde_core", + "zmij", +] + [[package]] name = "strsim" version = "0.11.1" @@ -208,3 +263,9 @@ checksum = "ae137229bcbd6cdf0f7b80a31df61766145077ddf49416a728b02cb3921ff3fc" dependencies = [ "windows-link", ] + +[[package]] +name = "zmij" +version = "1.0.23" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "29666d0abbfad1e3dc4dcf6144730dd3a3ab225bbbdac83319345b1b44ccfc1b" diff --git a/Cargo.toml b/Cargo.toml index 71d867f..332025f 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -17,6 +17,7 @@ clap = { version = "4.6.6", features = ["derive"] } clap_complete = "4.6.9" anyhow = "1.0.104" dotenvy = "0.15.7" +serde_json = "1.0.145" [build-dependencies] clap = { version = "4.6.6", features = ["derive"] } diff --git a/README.md b/README.md index 59708f9..78ca230 100644 --- a/README.md +++ b/README.md @@ -62,6 +62,24 @@ ahab compose bash # shell in the django service ahab compose exec ``` +## service detection + +`ahab` finds the services it needs in `docker compose config`, so they can be +named anything: + +- **postgres**: the service whose image is a postgres flavour, matching + `postg`, `timescale`, `pgvector` or `citus` +- **django**: the service that both builds an image and sets + `DJANGO_SETTINGS_MODULE`. A worker sharing the same build and `env_file` + matches too, so the one publishing ports wins + +If nothing matches, or two candidates cannot be told apart, `ahab` says so and +lists the services it looked at rather than guessing. + +`POSTGRES_USER` and `POSTGRES_DB` are read off the detected postgres service, so +`dropdb`, `createdb`, `pg_restore` and `pg_dump` use the role and database the +project declares, falling back to `db` when it declares neither. + ## django ```bash @@ -100,6 +118,4 @@ from the git `origin` remote. Currently `ahab` respects the following environment variables. -- `AHAB_DJANGO_CONTAINER`: control which compose service is used for sending django commands - defaults to `django` -- `AHAB_POSTGRES_CONTAINER`: control which compose service is used for sending postgres commands - defaults to `db` - `AHAB_LINK_ROOT`: root of the out-of-repo store `ahab link` moves paths into - defaults to `${XDG_DATA_HOME:-$HOME/.local/share}/ahab` diff --git a/src/compose.rs b/src/compose.rs new file mode 100644 index 0000000..d470913 --- /dev/null +++ b/src/compose.rs @@ -0,0 +1,261 @@ +use anyhow::{Context, Result, anyhow, bail}; +use serde_json::Value; + +use crate::command_builder::CommandBuilder; + +const POSTGRES_IMAGES: [&str; 4] = ["postg", "timescale", "pgvector", "citus"]; +const DJANGO_SETTINGS_MODULE: &str = "DJANGO_SETTINGS_MODULE"; + +pub struct Compose { + services: Value, +} + +impl Compose { + pub fn resolve() -> Result { + let out = CommandBuilder::docker_compose() + .args("config --format json") + .build()? + .output() + .context("running docker compose config")?; + + if !out.status.success() { + bail!( + "docker compose config failed: {}", + String::from_utf8_lossy(&out.stderr).trim() + ); + } + + let config: Value = + serde_json::from_slice(&out.stdout).context("parsing docker compose config")?; + let services = config + .get("services") + .cloned() + .context("compose config has no services")?; + + Ok(Self { services }) + } + + pub fn django(&self) -> Result { + let candidates: Vec<&str> = self + .names() + .into_iter() + .filter(|name| { + let service = &self.services[name]; + service.get("build").is_some() && env_var(service, DJANGO_SETTINGS_MODULE).is_some() + }) + .collect(); + + match candidates.as_slice() { + [] => Err(anyhow!( + "no service both builds an image and sets {DJANGO_SETTINGS_MODULE}, looked at {}", + self.names().join(", ") + )), + [only] => Ok(only.to_string()), + several => { + let serving: Vec<&&str> = several + .iter() + .filter(|name| self.publishes_ports(name)) + .collect(); + + match serving.as_slice() { + [only] => Ok(only.to_string()), + _ => Err(anyhow!( + "cannot tell which service runs django, {} all build an image, set \ + {DJANGO_SETTINGS_MODULE} and publish ports", + several.join(", ") + )), + } + } + } + } + + pub fn postgres(&self) -> Result { + let candidates: Vec<&str> = self + .names() + .into_iter() + .filter(|name| { + self.services[name] + .get("image") + .and_then(Value::as_str) + .is_some_and(is_postgres_image) + }) + .collect(); + + match candidates.as_slice() { + [] => Err(anyhow!( + "no service uses a postgres image, looked at {}", + self.names().join(", ") + )), + [only] => Ok(only.to_string()), + several => Err(anyhow!( + "cannot tell which service is the database, {} all use a postgres image", + several.join(", ") + )), + } + } + + pub fn postgres_credentials(&self, service: &str) -> (String, String) { + let service = &self.services[service]; + + let user = env_var(service, "POSTGRES_USER").unwrap_or_else(|| "db".to_string()); + let database = env_var(service, "POSTGRES_DB").unwrap_or_else(|| "db".to_string()); + + (user, database) + } + + fn names(&self) -> Vec<&str> { + let mut names: Vec<&str> = self + .services + .as_object() + .map(|services| services.keys().map(String::as_str).collect()) + .unwrap_or_default(); + + names.sort_unstable(); + names + } + + fn publishes_ports(&self, service: &str) -> bool { + self.services[service] + .get("ports") + .and_then(Value::as_array) + .is_some_and(|ports| !ports.is_empty()) + } +} + +fn env_var(service: &Value, key: &str) -> Option { + service + .get("environment")? + .get(key)? + .as_str() + .map(str::to_string) +} + +fn is_postgres_image(image: &str) -> bool { + let image = image.to_lowercase(); + POSTGRES_IMAGES.iter().any(|kind| image.contains(kind)) +} + +#[cfg(test)] +mod tests { + use super::{Compose, is_postgres_image}; + use serde_json::json; + + fn compose(services: serde_json::Value) -> Compose { + Compose { services } + } + + #[test] + fn recognises_postgres_flavours() { + for image in [ + "postgres:18", + "postgis/postgis:16-3.4", + "timescale/timescaledb:latest-pg16", + "pgvector/pgvector:pg16", + "citusdata/citus:12", + "PostgreSQL:16", + ] { + assert!(is_postgres_image(image), "{image} should count as postgres"); + } + + for image in ["redis:7-alpine", "mysql:8", "nginx", "local-django"] { + assert!(!is_postgres_image(image), "{image} should not"); + } + } + + #[test] + fn picks_the_service_that_builds_and_has_the_settings_module() { + let compose = compose(json!({ + "db": {"image": "postgres:18", "environment": {"POSTGRES_DB": "n"}}, + "django": { + "build": {"context": "."}, + "environment": {"DJANGO_SETTINGS_MODULE": "config.settings.local"}, + "ports": [{"target": 8000}], + }, + })); + + assert_eq!(compose.django().unwrap(), "django"); + assert_eq!(compose.postgres().unwrap(), "db"); + } + + #[test] + fn breaks_a_worker_tie_on_published_ports() { + let compose = compose(json!({ + "django": { + "build": {"context": "."}, + "environment": {"DJANGO_SETTINGS_MODULE": "config.settings.local"}, + "ports": [{"target": 8000}], + }, + "worker": { + "build": {"context": "."}, + "environment": {"DJANGO_SETTINGS_MODULE": "config.settings.local"}, + }, + })); + + assert_eq!(compose.django().unwrap(), "django"); + } + + #[test] + fn refuses_when_two_candidates_both_serve() { + let serving = json!({ + "build": {"context": "."}, + "environment": {"DJANGO_SETTINGS_MODULE": "config.settings.local"}, + "ports": [{"target": 8000}], + }); + let compose = compose(json!({"api": serving, "web": serving})); + + let error = compose.django().unwrap_err().to_string(); + assert!(error.contains("api, web"), "{error}"); + } + + #[test] + fn refuses_when_nothing_matches() { + let compose = compose(json!({ + "db": {"image": "postgres:18"}, + "cache": {"image": "redis:7"}, + })); + + assert!( + compose + .django() + .unwrap_err() + .to_string() + .contains("cache, db") + ); + assert_eq!(compose.postgres().unwrap(), "db"); + } + + #[test] + fn refuses_two_postgres_services() { + let compose = compose(json!({ + "db": {"image": "postgres:18"}, + "replica": {"image": "postgres:18"}, + })); + + let error = compose.postgres().unwrap_err().to_string(); + assert!(error.contains("db, replica"), "{error}"); + } + + #[test] + fn reads_credentials_off_the_detected_service() { + let compose = compose(json!({ + "db": { + "image": "postgres:18", + "environment": {"POSTGRES_USER": "myproject", "POSTGRES_DB": "myproject_db"}, + }, + })); + + let (user, database) = compose.postgres_credentials("db"); + assert_eq!( + (user.as_str(), database.as_str()), + ("myproject", "myproject_db") + ); + } + + #[test] + fn falls_back_to_db_when_the_service_says_nothing() { + let compose = compose(json!({"db": {"image": "postgres:18"}})); + + let (user, database) = compose.postgres_credentials("db"); + assert_eq!((user.as_str(), database.as_str()), ("db", "db")); + } +} diff --git a/src/lib.rs b/src/lib.rs index 99c221c..041d969 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -5,6 +5,7 @@ use std::{ pub mod cli; pub mod command_builder; +pub mod compose; pub mod scripts; // NOTE: stolen from https://docs.rs/debug_print/latest/debug_print/ diff --git a/src/main.rs b/src/main.rs index c218255..f9a55a6 100644 --- a/src/main.rs +++ b/src/main.rs @@ -35,13 +35,10 @@ fn main() -> Result<()> { cli::Django::Shell => scripts::django::shell(), cli::Django::Test => scripts::django::test(), }, - cli::Commands::Postgres { command } => { - let db_container = std::env::var("AHAB_POSTGRES_CONTAINER").unwrap_or("db".to_string()); - match command { - cli::Postgres::Import { path } => scripts::postgres::import(&db_container, &path), - cli::Postgres::Dump { path } => scripts::postgres::dump(&db_container, &path), - } - } + cli::Commands::Postgres { command } => match command { + cli::Postgres::Import { path } => scripts::postgres::import(&path), + cli::Postgres::Dump { path } => scripts::postgres::dump(&path), + }, cli::Commands::Link { command } => match command { cli::Link::Add { paths, force } => scripts::link::add(&paths, force), cli::Link::Check { diff --git a/src/scripts/django.rs b/src/scripts/django.rs index e64821e..fe4f50f 100644 --- a/src/scripts/django.rs +++ b/src/scripts/django.rs @@ -1,4 +1,3 @@ -use std::env; use std::fs::create_dir; use std::io::Write; use std::path::{Path, PathBuf}; @@ -6,6 +5,7 @@ use std::path::{Path, PathBuf}; use anyhow::{Result, anyhow}; use crate::command_builder::CommandBuilder; +use crate::compose::Compose; use crate::{create_file, safe_create_file}; const DEBUG_TEMPLATE: &str = r#"from django.core.management.base import BaseCommand @@ -55,7 +55,7 @@ pub fn make_command(app: &PathBuf, name: &str) -> Result<()> { } pub fn manage(rest: &[String]) -> Result<()> { - let container = env::var("AHAB_DJANGO_CONTAINER").unwrap_or("django".to_string()); + let container = Compose::resolve()?.django()?; let joined = rest.join(" "); let command = format!("run --rm {container} python manage.py {joined}"); CommandBuilder::docker_compose().args(&command).exec() diff --git a/src/scripts/docker_compose.rs b/src/scripts/docker_compose.rs index ddff849..00e114d 100644 --- a/src/scripts/docker_compose.rs +++ b/src/scripts/docker_compose.rs @@ -1,6 +1,7 @@ use anyhow::Result; use crate::command_builder::CommandBuilder; +use crate::compose::Compose; // simple commands pub fn build() -> Result<()> { @@ -12,10 +13,10 @@ pub fn down() -> Result<()> { } pub fn exec(rest: &[String]) -> Result<()> { - let container = std::env::var("AHAB_DJANGO_CONTAINER").unwrap_or("django".to_string()); + let service = Compose::resolve()?.django()?; CommandBuilder::docker_compose() .args("exec") - .args(&container) + .args(&service) .args(rest) .exec() } diff --git a/src/scripts/postgres.rs b/src/scripts/postgres.rs index 36575d5..b817525 100644 --- a/src/scripts/postgres.rs +++ b/src/scripts/postgres.rs @@ -8,57 +8,71 @@ use std::{ }; use super::docker_compose; -use crate::{command_builder::CommandBuilder, debug_eprintln}; +use crate::{command_builder::CommandBuilder, compose::Compose, debug_eprintln}; -fn get_containers(container: &str) -> Result<[String; 2]> { - // get db container - let db_container = CommandBuilder::docker_compose() - .args("ps -q") - .args(container) - .exec_get_stdout()? - .trim() - .to_string(); - - let no_result = db_container.is_empty(); - if no_result { - return Err(anyhow!("no container")); - } - - // get all containers and filter out db container - let app_containers = CommandBuilder::docker_compose() - .args("ps -q") - .exec_get_stdout()? - .split_whitespace() - .filter(|x| x != &db_container) - .collect::>() - .join(" "); - - Ok([db_container, app_containers]) +struct Database { + service: String, + container: String, + user: String, + name: String, } -pub fn import(container: &str, file: &Path) -> Result<()> { - let [db_container, _] = get_containers(container)?; +impl Database { + fn resolve() -> Result { + let compose = Compose::resolve()?; + let service = compose.postgres()?; + let (user, name) = compose.postgres_credentials(&service); + + let container = CommandBuilder::docker_compose() + .args("ps -q") + .args(&service) + .exec_get_stdout()? + .trim() + .to_string(); + + if container.is_empty() { + return Err(anyhow!("service {service} has no running container")); + } + + Ok(Self { + service, + container, + user, + name, + }) + } +} + +pub fn import(file: &Path) -> Result<()> { + let db = Database::resolve()?; let dump_file = file.to_string_lossy(); eprintln!("stopping all containers"); docker_compose::stop()?; eprintln!("starting db container"); - docker_compose::start(Some(container))?; + docker_compose::start(Some(&db.service))?; eprintln!("restoring database"); let commands = [ - format!("cp -L {dump_file} {db_container}:/tmp/dbdump"), - format!("exec {db_container} dropdb -U db db"), - format!("exec {db_container} createdb -U db -E utf8 -T template0 db"), - format!("exec {db_container} pg_restore -U db --dbname=db /tmp/dbdump"), + format!("cp -L {dump_file} {}:/tmp/dbdump", db.container), + format!("exec {} dropdb -U {} {}", db.container, db.user, db.name), + format!( + "exec {} createdb -U {} -E utf8 -T template0 {}", + db.container, db.user, db.name + ), + format!( + "exec {} pg_restore -U {} --dbname={} /tmp/dbdump", + db.container, db.user, db.name + ), ]; for command in commands { debug_eprintln!("waiting until pg_isready"); while !CommandBuilder::docker() .args(&format!( - "exec {db_container} pg_isready -h {container} -p 5432 -U db", + "exec {} pg_isready -U {} -d {}", + db.container, db.user, db.name )) .build()? .stdout(Stdio::null()) @@ -78,15 +92,18 @@ pub fn import(container: &str, file: &Path) -> Result<()> { Ok(()) } -pub fn dump(container: &str, file: &PathBuf) -> Result<()> { - let [db_container, _] = get_containers(container)?; +pub fn dump(file: &PathBuf) -> Result<()> { + let db = Database::resolve()?; eprintln!("dumping to local file {}", file.to_string_lossy()); let file = File::create(file)?; let stdout = Stdio::from(file); - let command = format!("exec {db_container} pg_dump -U db --format=c db"); + let command = format!( + "exec {} pg_dump -U {} --format=c {}", + db.container, db.user, db.name + ); CommandBuilder::docker() .args(&command) .exec_redirect_stdout(stdout)?;