diff --git a/Cargo.lock b/Cargo.lock index 627b6466..ad9bb9d7 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -1260,6 +1260,7 @@ dependencies = [ "serde", "serde_json", "strum", + "thiserror 2.0.19", "time", "tokio", "toml", diff --git a/Cargo.toml b/Cargo.toml index 48570292..f3383f22 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -18,6 +18,7 @@ reqwest = { version = "0.12.28", features = ["json", "rustls-tls", "rustls-tls-n serde = { version = "1.0.228", features = ["derive"] } serde_json = "1.0.150" strum = { version = "0.28.0", features = ["derive"] } +thiserror = "2.0.18" time = { version = "0.3.47", features = ["serde", "serde-human-readable", "parsing", "macros"] } tokio = { version = "1.51.1", features = ["rt-multi-thread", "sync", "macros", "time"] } toml = "1.1.2" diff --git a/crates/flowrs-airflow/Cargo.toml b/crates/flowrs-airflow/Cargo.toml index f9a70491..895fa17b 100644 --- a/crates/flowrs-airflow/Cargo.toml +++ b/crates/flowrs-airflow/Cargo.toml @@ -19,6 +19,7 @@ anyhow = { workspace = true } async-trait = { workspace = true } clap = { workspace = true } strum = { workspace = true } +thiserror = { workspace = true } aws-config = { version = "1.8.18", optional = true } aws-sdk-mwaa = { version = "1.108.0", optional = true } dirs = { workspace = true, optional = true } diff --git a/crates/flowrs-airflow/src/client.rs b/crates/flowrs-airflow/src/client.rs index 625c014f..e8600845 100644 --- a/crates/flowrs-airflow/src/client.rs +++ b/crates/flowrs-airflow/src/client.rs @@ -3,28 +3,22 @@ pub mod base; pub mod v1; pub mod v2; -use anyhow::Result; +use reqwest::Response; +use serde::de::DeserializeOwned; -use crate::config::{AirflowConfig, AirflowVersion}; +use crate::error::{AirflowError, Result}; pub use base::BaseClient; pub use v1::V1Client; pub use v2::V2Client; -/// Enum wrapping the versioned API clients. -/// V1 is for Airflow v2 (uses /api/v1), V2 is for Airflow v3 (uses /api/v2). -#[derive(Debug)] -pub enum AirflowApiClient { - V1(V1Client), - V2(V2Client), -} - -/// Create an Airflow API client based on the configuration version -pub fn create_api_client(config: &AirflowConfig) -> Result { - let base = BaseClient::new(config.clone())?; - - match config.version { - AirflowVersion::V2 => Ok(AirflowApiClient::V1(V1Client::new(base))), // V2 uses API v1 - AirflowVersion::V3 => Ok(AirflowApiClient::V2(V2Client::new(base))), // V3 uses API v2 - } +/// Deserialize a response body, keeping a snippet of it in the error on failure. +/// +/// The body is read as text first rather than using `Response::json`, because some +/// deployments return payloads that do not match the documented schema (older v2.x +/// versions omit `dag_display_name`, for instance) and the raw body is what makes +/// those failures diagnosable. +pub(crate) async fn read_json(response: Response, context: &str) -> Result { + let body = response.text().await?; + serde_json::from_str(&body).map_err(|e| AirflowError::decode(context, &body, e)) } diff --git a/crates/flowrs-airflow/src/client/auth/basic.rs b/crates/flowrs-airflow/src/client/auth/basic.rs index 2f3dfada..947a058a 100644 --- a/crates/flowrs-airflow/src/client/auth/basic.rs +++ b/crates/flowrs-airflow/src/client/auth/basic.rs @@ -1,4 +1,4 @@ -use anyhow::Result; +use crate::error::Result; use async_trait::async_trait; use log::info; use reqwest::RequestBuilder; diff --git a/crates/flowrs-airflow/src/client/auth/command.rs b/crates/flowrs-airflow/src/client/auth/command.rs index 97205dc7..fd30b759 100644 --- a/crates/flowrs-airflow/src/client/auth/command.rs +++ b/crates/flowrs-airflow/src/client/auth/command.rs @@ -1,12 +1,13 @@ use std::fmt; use std::time::{Duration, Instant}; -use anyhow::{Context, Result}; +use anyhow::Context; use async_trait::async_trait; use log::info; use reqwest::RequestBuilder; use super::AuthProvider; +use crate::error::{AirflowError, Result}; /// How long a fetched token is reused before the helper command is run again. /// @@ -32,7 +33,10 @@ impl CommandTokenProvider { } /// Run the helper command and return its trimmed token output. - async fn fetch_token(&self) -> Result { + /// + /// Uses `anyhow` internally for the layered context messages; the chain is + /// flattened into `AirflowError::Auth` at the trait boundary. + async fn fetch_token(&self) -> anyhow::Result { let cmd = self.cmd.clone(); let output = tokio::task::spawn_blocking(move || { std::process::Command::new("sh") @@ -81,7 +85,10 @@ impl AuthProvider for CommandTokenProvider { if !fresh { info!("🔑 Token Auth (command): refreshing via {}", self.cmd); - let token = self.fetch_token().await?; + let token = self + .fetch_token() + .await + .map_err(|e| AirflowError::auth("token command", &e))?; *cached = Some((token, Instant::now())); } diff --git a/crates/flowrs-airflow/src/client/auth/mod.rs b/crates/flowrs-airflow/src/client/auth/mod.rs index d7784312..548dd5c1 100644 --- a/crates/flowrs-airflow/src/client/auth/mod.rs +++ b/crates/flowrs-airflow/src/client/auth/mod.rs @@ -2,7 +2,7 @@ mod basic; mod command; mod static_token; -use anyhow::Result; +use crate::error::Result; pub use basic::BasicAuthProvider; pub use command::CommandTokenProvider; @@ -46,33 +46,35 @@ pub fn create_auth_provider(auth: &AirflowAuth) -> Result> #[cfg(feature = "conveyor")] AirflowAuth::Conveyor => Ok(Box::new(ConveyorAuthProvider::new())), #[cfg(not(feature = "conveyor"))] - AirflowAuth::Conveyor => { - anyhow::bail!("Conveyor support not compiled. Enable the 'conveyor' feature.") - } + AirflowAuth::Conveyor => Err(crate::error::AirflowError::FeatureNotEnabled { + service: "Conveyor", + feature: "conveyor", + }), #[cfg(feature = "mwaa")] AirflowAuth::Mwaa(mwaa_auth) => Ok(Box::new(MwaaAuthProvider::from(mwaa_auth))), #[cfg(not(feature = "mwaa"))] - AirflowAuth::Mwaa(_) => { - anyhow::bail!("MWAA support not compiled. Enable the 'mwaa' feature.") - } + AirflowAuth::Mwaa(_) => Err(crate::error::AirflowError::FeatureNotEnabled { + service: "MWAA", + feature: "mwaa", + }), #[cfg(feature = "astronomer")] AirflowAuth::Astronomer(astro_auth) => { Ok(Box::new(AstronomerAuthProvider::from(astro_auth))) } #[cfg(not(feature = "astronomer"))] - AirflowAuth::Astronomer(_) => { - anyhow::bail!("Astronomer support not compiled. Enable the 'astronomer' feature.") - } + AirflowAuth::Astronomer(_) => Err(crate::error::AirflowError::FeatureNotEnabled { + service: "Astronomer", + feature: "astronomer", + }), #[cfg(feature = "composer")] AirflowAuth::Composer(composer_auth) => { Ok(Box::new(ComposerAuthProvider::new(composer_auth)?)) } #[cfg(not(feature = "composer"))] - AirflowAuth::Composer(_) => { - anyhow::bail!( - "Google Cloud Composer support not compiled. Enable the 'composer' feature." - ) - } + AirflowAuth::Composer(_) => Err(crate::error::AirflowError::FeatureNotEnabled { + service: "Google Cloud Composer", + feature: "composer", + }), } } diff --git a/crates/flowrs-airflow/src/client/auth/static_token.rs b/crates/flowrs-airflow/src/client/auth/static_token.rs index ab3543f8..80dd4d7c 100644 --- a/crates/flowrs-airflow/src/client/auth/static_token.rs +++ b/crates/flowrs-airflow/src/client/auth/static_token.rs @@ -1,4 +1,4 @@ -use anyhow::Result; +use crate::error::Result; use async_trait::async_trait; use log::info; use reqwest::RequestBuilder; diff --git a/crates/flowrs-airflow/src/client/base.rs b/crates/flowrs-airflow/src/client/base.rs index ebbe49b8..e75fbc65 100644 --- a/crates/flowrs-airflow/src/client/base.rs +++ b/crates/flowrs-airflow/src/client/base.rs @@ -1,18 +1,20 @@ -use anyhow::Result; use log::{debug, warn}; -use reqwest::{Method, Url}; -use std::convert::TryFrom; +use reqwest::{Method, RequestBuilder, Response, Url}; use std::fmt; use std::time::Duration; use super::auth::{create_auth_provider, AuthProvider}; use crate::config::AirflowConfig; +use crate::error::{AirflowError, Result}; /// Base HTTP client for Airflow API communication. /// Handles authentication and provides base request building functionality. pub struct BaseClient { - pub client: reqwest::Client, - pub config: AirflowConfig, + client: reqwest::Client, + config: AirflowConfig, + /// The endpoint parsed once at construction, so an unusable URL is reported + /// when the client is created rather than on every request. + endpoint: Url, auth_provider: Box, } @@ -21,6 +23,7 @@ impl fmt::Debug for BaseClient { f.debug_struct("BaseClient") .field("client", &self.client) .field("config", &self.config) + .field("endpoint", &self.endpoint) .field("auth_provider", &"") .finish() } @@ -34,6 +37,16 @@ impl BaseClient { config.name ); } + let mut endpoint = Url::parse(&config.endpoint) + .map_err(|e| AirflowError::invalid_url(config.endpoint.clone(), e))?; + // Ensure the base path ends with a slash so `base_api`'s relative `join` + // extends the endpoint rather than replacing its final segment. Without + // this, a reverse-proxy prefix such as `/airflow` would be dropped. + if !endpoint.path().ends_with('/') { + let with_slash = format!("{}/", endpoint.path()); + endpoint.set_path(&with_slash); + } + let client = reqwest::Client::builder() .timeout(Duration::from_secs(config.timeout_secs)) .use_rustls_tls() @@ -45,30 +58,106 @@ impl BaseClient { Ok(Self { client, config, + endpoint, auth_provider, }) } + /// The base URL of the Airflow deployment this client talks to. + pub const fn endpoint(&self) -> &Url { + &self.endpoint + } + + pub const fn config(&self) -> &AirflowConfig { + &self.config + } + /// Build a base request with authentication for the specified API version - pub async fn base_api( + pub(crate) async fn base_api( &self, method: Method, endpoint: &str, api_version: &str, - ) -> Result { - let base_url = Url::parse(&self.config.endpoint)?; - let url = base_url.join(format!("{api_version}/{endpoint}").as_str())?; + ) -> Result { + let path = format!("{api_version}/{endpoint}"); + let url = self + .endpoint + .join(&path) + .map_err(|e| AirflowError::invalid_url(path, e))?; debug!("🔗 Request URL: {url}"); let request = self.client.request(method, url); self.auth_provider.authenticate(request).await } + + /// Send a request, turning a non-success status into [`AirflowError::Status`]. + /// + /// The request is built rather than sent through `RequestBuilder::send` so that + /// the method and URL are available for the error message, and the body is read + /// instead of calling `error_for_status`, which discards Airflow's `detail`. + pub(crate) async fn execute(&self, request: RequestBuilder) -> Result { + let request = request.build()?; + let method = request.method().clone(); + let url = request.url().clone(); + + let response = self.client.execute(request).await?; + let status = response.status(); + if status.is_success() { + return Ok(response); + } + + let body = response.text().await.unwrap_or_default(); + Err(AirflowError::status(&method, &url, status, &body)) + } } impl TryFrom<&AirflowConfig> for BaseClient { - type Error = anyhow::Error; + type Error = AirflowError; - fn try_from(config: &AirflowConfig) -> Result { + fn try_from(config: &AirflowConfig) -> Result { Self::new(config.clone()) } } + +#[cfg(test)] +mod tests { + use super::*; + use crate::auth::{AirflowAuth, TokenSource}; + + fn config(endpoint: &str) -> AirflowConfig { + AirflowConfig { + name: "test".to_string(), + endpoint: endpoint.to_string(), + auth: AirflowAuth::Token(TokenSource::Static { + token: "tok".to_string(), + }), + managed: None, + version: crate::config::AirflowVersion::V3, + timeout_secs: crate::config::default_timeout(), + insecure: false, + } + } + + #[test] + fn rejects_an_unusable_endpoint_at_construction() { + let error = BaseClient::new(config("not a url")).expect_err("should reject"); + assert!( + matches!(error, AirflowError::InvalidUrl { .. }), + "got: {error:?}" + ); + } + + #[test] + fn accepts_a_valid_endpoint() { + let client = BaseClient::new(config("http://localhost:8080")).expect("should accept"); + assert_eq!(client.endpoint().host_str(), Some("localhost")); + } + + #[test] + fn preserves_a_reverse_proxy_prefix_when_building_requests() { + let client = + BaseClient::new(config("http://localhost:8080/airflow")).expect("should accept"); + let url = client.endpoint().join("api/v2/dags").expect("should join"); + assert_eq!(url.as_str(), "http://localhost:8080/airflow/api/v2/dags"); + } +} diff --git a/crates/flowrs-airflow/src/client/v1/dag.rs b/crates/flowrs-airflow/src/client/v1/dag.rs index 56fdbe2f..50e94d4b 100644 --- a/crates/flowrs-airflow/src/client/v1/dag.rs +++ b/crates/flowrs-airflow/src/client/v1/dag.rs @@ -1,9 +1,10 @@ -use anyhow::Result; use log::{debug, info}; use reqwest::Method; use super::model::dag::DagCollectionResponse; -use super::{parse_json_response, V1Client}; +use super::V1Client; +use crate::client::read_json; +use crate::error::Result; const PAGE_SIZE: usize = 50; @@ -15,16 +16,12 @@ impl V1Client { let mut total_entries; loop { - let response = self + let request = self .base_api(Method::GET, "dags") .await? - .query(&[("limit", limit.to_string()), ("offset", offset.to_string())]) - .send() - .await? - .error_for_status()?; - - let response_text = response.text().await?; - let page: DagCollectionResponse = parse_json_response(&response_text, "DAGs response")?; + .query(&[("limit", limit.to_string()), ("offset", offset.to_string())]); + let response = self.execute(request).await?; + let page: DagCollectionResponse = read_json(response, "DAGs response").await?; total_entries = page.total_entries; let fetched_count = page.dags.len(); @@ -53,34 +50,30 @@ impl V1Client { } pub async fn patch_dag_pause(&self, dag_id: &str, is_paused: bool) -> Result<()> { - self.base_api(Method::PATCH, &format!("dags/{dag_id}")) + let request = self + .base_api(Method::PATCH, &format!("dags/{dag_id}")) .await? .query(&[("update_mask", "is_paused")]) - .json(&serde_json::json!({"is_paused": !is_paused})) - .send() - .await? - .error_for_status()?; + .json(&serde_json::json!({"is_paused": !is_paused})); + self.execute(request).await?; Ok(()) } pub async fn fetch_dag_code(&self, file_token: &str) -> Result { - let r = self + let request = self .base_api(Method::GET, &format!("dagSources/{file_token}")) - .await? - .build()?; - let response = self.base.client.execute(r).await?.error_for_status()?; + .await?; + let response = self.execute(request).await?; let code = response.text().await?; Ok(code) } pub async fn fetch_dag_params(&self, dag_id: &str) -> Result> { - let response = self + let request = self .base_api(Method::GET, &format!("dags/{dag_id}/details")) - .await? - .send() - .await? - .error_for_status()?; - let body: serde_json::Value = response.json().await?; + .await?; + let response = self.execute(request).await?; + let body: serde_json::Value = read_json(response, "DAG details response").await?; Ok(body.get("params").cloned()) } } diff --git a/crates/flowrs-airflow/src/client/v1/dagrun.rs b/crates/flowrs-airflow/src/client/v1/dagrun.rs index 86bde95c..625b484a 100644 --- a/crates/flowrs-airflow/src/client/v1/dagrun.rs +++ b/crates/flowrs-airflow/src/client/v1/dagrun.rs @@ -1,62 +1,54 @@ -use anyhow::Result; use log::debug; -use reqwest::{Method, Response}; +use reqwest::Method; use super::model; use super::V1Client; +use crate::client::read_json; +use crate::error::Result; impl V1Client { pub async fn fetch_dagruns( &self, dag_id: &str, ) -> Result { - let response: Response = self + let request = self .base_api(Method::GET, &format!("dags/{dag_id}/dagRuns")) .await? - .query(&[("order_by", "-execution_date"), ("limit", "50")]) - .send() - .await? - .error_for_status()?; - - let dagruns: model::dagrun::DAGRunCollectionResponse = response.json().await?; - Ok(dagruns) + .query(&[("order_by", "-execution_date"), ("limit", "50")]); + let response = self.execute(request).await?; + read_json(response, "DAG runs response").await } pub async fn fetch_all_dagruns(&self) -> Result { - let response: Response = self + let request = self .base_api(Method::POST, "dags/~/dagRuns/list") .await? - .json(&serde_json::json!({"page_limit": 200})) - .send() - .await? - .error_for_status()?; - let dagruns: model::dagrun::DAGRunCollectionResponse = response.json().await?; - Ok(dagruns) + .json(&serde_json::json!({"page_limit": 200})); + let response = self.execute(request).await?; + read_json(response, "DAG runs response").await } pub async fn patch_dag_run(&self, dag_id: &str, dag_run_id: &str, status: &str) -> Result<()> { - self.base_api( - Method::PATCH, - &format!("dags/{dag_id}/dagRuns/{dag_run_id}"), - ) - .await? - .json(&serde_json::json!({"state": status})) - .send() - .await? - .error_for_status()?; + let request = self + .base_api( + Method::PATCH, + &format!("dags/{dag_id}/dagRuns/{dag_run_id}"), + ) + .await? + .json(&serde_json::json!({"state": status})); + self.execute(request).await?; Ok(()) } pub async fn post_clear_dagrun(&self, dag_id: &str, dag_run_id: &str) -> Result<()> { - self.base_api( - Method::POST, - &format!("dags/{dag_id}/dagRuns/{dag_run_id}/clear"), - ) - .await? - .json(&serde_json::json!({"dry_run": false})) - .send() - .await? - .error_for_status()?; + let request = self + .base_api( + Method::POST, + &format!("dags/{dag_id}/dagRuns/{dag_run_id}/clear"), + ) + .await? + .json(&serde_json::json!({"dry_run": false})); + self.execute(request).await?; Ok(()) } @@ -75,13 +67,11 @@ impl V1Client { body["conf"] = conf; } - let resp: Response = self + let request = self .base_api(Method::POST, &format!("dags/{dag_id}/dagRuns")) .await? - .json(&body) - .send() - .await? - .error_for_status()?; + .json(&body); + let resp = self.execute(request).await?; debug!("{resp:?}"); Ok(()) } diff --git a/crates/flowrs-airflow/src/client/v1/dagstats.rs b/crates/flowrs-airflow/src/client/v1/dagstats.rs index 5101ab31..a525d158 100644 --- a/crates/flowrs-airflow/src/client/v1/dagstats.rs +++ b/crates/flowrs-airflow/src/client/v1/dagstats.rs @@ -1,23 +1,20 @@ -use anyhow::Result; use reqwest::Method; use super::model; use super::V1Client; +use crate::client::read_json; +use crate::error::Result; impl V1Client { pub async fn fetch_dag_stats( &self, dag_ids: Vec<&str>, ) -> Result { - let response = self + let request = self .base_api(Method::GET, "dagStats") .await? - .query(&[("dag_ids", dag_ids.join(","))]) - .send() - .await? - .error_for_status()?; - - let dag_stats = response.json::().await?; - Ok(dag_stats) + .query(&[("dag_ids", dag_ids.join(","))]); + let response = self.execute(request).await?; + read_json(response, "DAG stats response").await } } diff --git a/crates/flowrs-airflow/src/client/v1/log.rs b/crates/flowrs-airflow/src/client/v1/log.rs index f8185bb1..3da74df9 100644 --- a/crates/flowrs-airflow/src/client/v1/log.rs +++ b/crates/flowrs-airflow/src/client/v1/log.rs @@ -1,11 +1,12 @@ use std::sync::LazyLock; -use anyhow::Result; use regex::Regex; use reqwest::Method; use super::model; use super::V1Client; +use crate::client::read_json; +use crate::error::Result; /// Compiled regex for parsing V1 log content (Python tuple format). /// @@ -59,7 +60,7 @@ impl V1Client { task_id: &str, task_try: u32, ) -> Result { - let response = self + let request = self .base_api( Method::GET, &format!( @@ -68,12 +69,9 @@ impl V1Client { ) .await? .query(&[("full_content", "true")]) - .header("Accept", "application/json") - .send() - .await? - .error_for_status()?; - let log = response.json::().await?; - Ok(log) + .header("Accept", "application/json"); + let response = self.execute(request).await?; + read_json(response, "task logs response").await } } diff --git a/crates/flowrs-airflow/src/client/v1/mod.rs b/crates/flowrs-airflow/src/client/v1/mod.rs index c731c4b1..8ff70d56 100644 --- a/crates/flowrs-airflow/src/client/v1/mod.rs +++ b/crates/flowrs-airflow/src/client/v1/mod.rs @@ -7,32 +7,15 @@ pub mod log; mod task; mod taskinstance; -use anyhow::Result; -use reqwest::Method; +use reqwest::{Method, RequestBuilder, Response, Url}; use super::base::BaseClient; - -/// Parse a JSON response body, including a snippet of the body in the error on failure. -/// -/// Some Airflow deployments return responses that don't match the documented schema -/// (e.g. older v2.x versions omit `dag_display_name` / `task_display_name`). Surfacing -/// the response body makes those parse failures actionable. -pub(crate) fn parse_json_response( - body: &str, - context: &str, -) -> Result { - serde_json::from_str(body).map_err(|e| { - anyhow::anyhow!( - "Failed to parse {context}: {e}. Response body (first 1000 chars): {}", - body.chars().take(1000).collect::() - ) - }) -} +use crate::error::Result; /// API v1 client implementation (for Airflow v2, uses /api/v1 endpoint) #[derive(Debug)] pub struct V1Client { - pub base: BaseClient, + base: BaseClient, } impl V1Client { @@ -42,18 +25,18 @@ impl V1Client { Self { base } } - pub(crate) async fn base_api( - &self, - method: Method, - endpoint: &str, - ) -> Result { + pub(crate) async fn base_api(&self, method: Method, endpoint: &str) -> Result { self.base .base_api(method, endpoint, Self::API_VERSION) .await } + pub(crate) async fn execute(&self, request: RequestBuilder) -> Result { + self.base.execute(request).await + } + /// Returns the base endpoint URL for this client - pub fn endpoint(&self) -> &str { - &self.base.config.endpoint + pub const fn endpoint(&self) -> &Url { + self.base.endpoint() } } diff --git a/crates/flowrs-airflow/src/client/v1/task.rs b/crates/flowrs-airflow/src/client/v1/task.rs index d15f3718..57f19a8e 100644 --- a/crates/flowrs-airflow/src/client/v1/task.rs +++ b/crates/flowrs-airflow/src/client/v1/task.rs @@ -1,19 +1,16 @@ -use anyhow::Result; use reqwest::Method; use super::model; use super::V1Client; +use crate::client::read_json; +use crate::error::Result; impl V1Client { pub async fn fetch_tasks(&self, dag_id: &str) -> Result { - let response = self + let request = self .base_api(Method::GET, &format!("dags/{dag_id}/tasks")) - .await? - .send() - .await? - .error_for_status()?; - - let task_collection: model::task::TaskCollectionResponse = response.json().await?; - Ok(task_collection) + .await?; + let response = self.execute(request).await?; + read_json(response, "tasks response").await } } diff --git a/crates/flowrs-airflow/src/client/v1/taskinstance.rs b/crates/flowrs-airflow/src/client/v1/taskinstance.rs index 9efcbc0d..79733c65 100644 --- a/crates/flowrs-airflow/src/client/v1/taskinstance.rs +++ b/crates/flowrs-airflow/src/client/v1/taskinstance.rs @@ -1,9 +1,10 @@ -use anyhow::Result; use log::{debug, info}; -use reqwest::{Method, Response}; +use reqwest::Method; use super::model; -use super::{parse_json_response, V1Client}; +use super::V1Client; +use crate::client::read_json; +use crate::error::Result; const PAGE_SIZE: usize = 100; @@ -19,20 +20,16 @@ impl V1Client { let mut total_entries; loop { - let response: Response = self + let request = self .base_api( Method::GET, &format!("dags/{dag_id}/dagRuns/{dag_run_id}/taskInstances"), ) .await? - .query(&[("limit", limit.to_string()), ("offset", offset.to_string())]) - .send() - .await? - .error_for_status()?; - - let response_text = response.text().await?; + .query(&[("limit", limit.to_string()), ("offset", offset.to_string())]); + let response = self.execute(request).await?; let page: model::taskinstance::TaskInstanceCollectionResponse = - parse_json_response(&response_text, "task instances response")?; + read_json(response, "task instances response").await?; total_entries = page.total_entries; let fetched_count = page.task_instances.len(); @@ -75,17 +72,13 @@ impl V1Client { let mut total_entries; loop { - let response: Response = self + let request = self .base_api(Method::GET, "dags/~/dagRuns/~/taskInstances") .await? - .query(&[("limit", limit.to_string()), ("offset", offset.to_string())]) - .send() - .await? - .error_for_status()?; - - let response_text = response.text().await?; + .query(&[("limit", limit.to_string()), ("offset", offset.to_string())]); + let response = self.execute(request).await?; let page: model::taskinstance::TaskInstanceCollectionResponse = - parse_json_response(&response_text, "all task instances response")?; + read_json(response, "all task instances response").await?; total_entries = page.total_entries; let fetched_count = page.task_instances.len(); @@ -119,19 +112,15 @@ impl V1Client { dag_run_id: &str, task_id: &str, ) -> Result { - let response: Response = self + let request = self .base_api( Method::GET, &format!("dags/{dag_id}/dagRuns/{dag_run_id}/taskInstances/{task_id}/tries"), ) - .await? - .send() - .await? - .error_for_status()?; - - let response_text = response.text().await?; + .await?; + let response = self.execute(request).await?; let tries: model::taskinstance::TaskInstanceTriesResponse = - parse_json_response(&response_text, "task instance tries response")?; + read_json(response, "task instance tries response").await?; debug!( "Fetched {} tries for task {task_id}", tries.task_instances.len() @@ -147,16 +136,14 @@ impl V1Client { task_id: &str, status: &str, ) -> Result<()> { - let resp: Response = self + let request = self .base_api( Method::PATCH, &format!("dags/{dag_id}/dagRuns/{dag_run_id}/taskInstances/{task_id}"), ) .await? - .json(&serde_json::json!({"new_state": status, "dry_run": false})) - .send() - .await? - .error_for_status()?; + .json(&serde_json::json!({"new_state": status, "dry_run": false})); + let resp = self.execute(request).await?; debug!("{resp:?}"); Ok(()) } @@ -167,7 +154,7 @@ impl V1Client { dag_run_id: &str, task_id: &str, ) -> Result<()> { - let resp: Response = self + let request = self .base_api(Method::POST, &format!("dags/{dag_id}/clearTaskInstances")) .await? .json(&serde_json::json!( @@ -179,10 +166,8 @@ impl V1Client { "only_failed": false, "reset_dag_runs": true, } - )) - .send() - .await? - .error_for_status()?; + )); + let resp = self.execute(request).await?; debug!("{resp:?}"); Ok(()) } diff --git a/crates/flowrs-airflow/src/client/v2/dag.rs b/crates/flowrs-airflow/src/client/v2/dag.rs index cfd3211b..dd1de6f4 100644 --- a/crates/flowrs-airflow/src/client/v2/dag.rs +++ b/crates/flowrs-airflow/src/client/v2/dag.rs @@ -1,9 +1,10 @@ -use anyhow::Result; use log::{debug, info}; use reqwest::Method; use super::model; use super::V2Client; +use crate::client::read_json; +use crate::error::Result; const PAGE_SIZE: usize = 50; @@ -15,15 +16,12 @@ impl V2Client { let mut total_entries; loop { - let response = self + let request = self .base_api(Method::GET, "dags") .await? - .query(&[("limit", limit.to_string()), ("offset", offset.to_string())]) - .send() - .await? - .error_for_status()?; - - let page: model::dag::DagList = response.json().await?; + .query(&[("limit", limit.to_string()), ("offset", offset.to_string())]); + let response = self.execute(request).await?; + let page: model::dag::DagList = read_json(response, "DAGs response").await?; total_entries = page.total_entries; let fetched_count = page.dags.len(); @@ -52,35 +50,31 @@ impl V2Client { } pub async fn patch_dag_pause(&self, dag_id: &str, is_paused: bool) -> Result<()> { - self.base_api(Method::PATCH, &format!("dags/{dag_id}")) - .await? - .json(&serde_json::json!({"is_paused": !is_paused})) - .send() + let request = self + .base_api(Method::PATCH, &format!("dags/{dag_id}")) .await? - .error_for_status()?; + .json(&serde_json::json!({"is_paused": !is_paused})); + self.execute(request).await?; Ok(()) } pub async fn fetch_dag_code(&self, dag_id: &str) -> Result { - let r = self + let request = self .base_api(Method::GET, &format!("dagSources/{dag_id}")) - .await? - .build()?; - let response = self.base.client.execute(r).await?.error_for_status()?; - let dag_source: model::dag::DagSource = response.json().await?; + .await?; + let response = self.execute(request).await?; + let dag_source: model::dag::DagSource = read_json(response, "DAG source response").await?; Ok(dag_source.content) } pub async fn fetch_dag_params(&self, dag_id: &str) -> Result> { // `params` lives on the details endpoint; the plain `dags/{dag_id}` // (DAGResponse) schema does not include it. - let response = self + let request = self .base_api(Method::GET, &format!("dags/{dag_id}/details")) - .await? - .send() - .await? - .error_for_status()?; - let body: serde_json::Value = response.json().await?; + .await?; + let response = self.execute(request).await?; + let body: serde_json::Value = read_json(response, "DAG details response").await?; Ok(body.get("params").cloned()) } } diff --git a/crates/flowrs-airflow/src/client/v2/dagrun.rs b/crates/flowrs-airflow/src/client/v2/dagrun.rs index 461aebb5..21fd6498 100644 --- a/crates/flowrs-airflow/src/client/v2/dagrun.rs +++ b/crates/flowrs-airflow/src/client/v2/dagrun.rs @@ -1,58 +1,51 @@ -use anyhow::Result; use log::debug; -use reqwest::{Method, Response}; +use reqwest::Method; use super::model; use super::V2Client; +use crate::client::read_json; +use crate::error::Result; impl V2Client { pub async fn fetch_dagruns(&self, dag_id: &str) -> Result { - let response: Response = self + let request = self .base_api(Method::GET, &format!("dags/{dag_id}/dagRuns")) .await? - .query(&[("order_by", "-run_after"), ("limit", "50")]) - .send() - .await? - .error_for_status()?; - let dagruns: model::dagrun::DagRunList = response.json().await?; - Ok(dagruns) + .query(&[("order_by", "-run_after"), ("limit", "50")]); + let response = self.execute(request).await?; + read_json(response, "DAG runs response").await } pub async fn fetch_all_dagruns(&self) -> Result { - let response: Response = self + let request = self .base_api(Method::POST, "dags/~/dagRuns/list") .await? - .json(&serde_json::json!({"page_limit": 200})) - .send() - .await? - .error_for_status()?; - let dagruns: model::dagrun::DagRunList = response.json().await?; - Ok(dagruns) + .json(&serde_json::json!({"page_limit": 200})); + let response = self.execute(request).await?; + read_json(response, "DAG runs response").await } pub async fn patch_dag_run(&self, dag_id: &str, dag_run_id: &str, status: &str) -> Result<()> { - self.base_api( - Method::PATCH, - &format!("dags/{dag_id}/dagRuns/{dag_run_id}"), - ) - .await? - .json(&serde_json::json!({"state": status})) - .send() - .await? - .error_for_status()?; + let request = self + .base_api( + Method::PATCH, + &format!("dags/{dag_id}/dagRuns/{dag_run_id}"), + ) + .await? + .json(&serde_json::json!({"state": status})); + self.execute(request).await?; Ok(()) } pub async fn post_clear_dagrun(&self, dag_id: &str, dag_run_id: &str) -> Result<()> { - self.base_api( - Method::POST, - &format!("dags/{dag_id}/dagRuns/{dag_run_id}/clear"), - ) - .await? - .json(&serde_json::json!({"dry_run": false})) - .send() - .await? - .error_for_status()?; + let request = self + .base_api( + Method::POST, + &format!("dags/{dag_id}/dagRuns/{dag_run_id}/clear"), + ) + .await? + .json(&serde_json::json!({"dry_run": false})); + self.execute(request).await?; Ok(()) } @@ -67,13 +60,11 @@ impl V2Client { body["conf"] = conf; } - let resp: Response = self + let request = self .base_api(Method::POST, &format!("dags/{dag_id}/dagRuns")) .await? - .json(&body) - .send() - .await? - .error_for_status()?; + .json(&body); + let resp = self.execute(request).await?; debug!("{resp:?}"); Ok(()) } diff --git a/crates/flowrs-airflow/src/client/v2/dagstats.rs b/crates/flowrs-airflow/src/client/v2/dagstats.rs index 1a28ea00..3cf30223 100644 --- a/crates/flowrs-airflow/src/client/v2/dagstats.rs +++ b/crates/flowrs-airflow/src/client/v2/dagstats.rs @@ -1,27 +1,22 @@ -use anyhow::Result; use reqwest::Method; use super::model; use super::V2Client; +use crate::client::read_json; +use crate::error::Result; impl V2Client { pub async fn fetch_dag_stats( &self, dag_ids: Vec<&str>, ) -> Result { - let response = self - .base_api(Method::GET, "dagStats") - .await? - .query( - &dag_ids - .into_iter() - .map(|id| ("dag_ids", id)) - .collect::>(), - ) - .send() - .await? - .error_for_status()?; - let dag_stats = response.json::().await?; - Ok(dag_stats) + let request = self.base_api(Method::GET, "dagStats").await?.query( + &dag_ids + .into_iter() + .map(|id| ("dag_ids", id)) + .collect::>(), + ); + let response = self.execute(request).await?; + read_json(response, "DAG stats response").await } } diff --git a/crates/flowrs-airflow/src/client/v2/log.rs b/crates/flowrs-airflow/src/client/v2/log.rs index be40d5c6..e019ae8b 100644 --- a/crates/flowrs-airflow/src/client/v2/log.rs +++ b/crates/flowrs-airflow/src/client/v2/log.rs @@ -1,9 +1,10 @@ -use anyhow::Result; use log::debug; use reqwest::Method; use super::model; use super::V2Client; +use crate::client::read_json; +use crate::error::Result; impl V2Client { pub async fn fetch_task_logs( @@ -13,7 +14,7 @@ impl V2Client { task_id: &str, task_try: u32, ) -> Result { - let response = self + let request = self .base_api( Method::GET, &format!( @@ -22,13 +23,11 @@ impl V2Client { ) .await? .query(&[("full_content", "true")]) - .header("Accept", "application/json") - .send() - .await? - .error_for_status()?; + .header("Accept", "application/json"); + let response = self.execute(request).await?; debug!("Response: {response:?}"); - let log = response.json::().await?; + let log: model::log::Log = read_json(response, "task logs response").await?; debug!("Parsed Log: {log:?}"); Ok(log) } diff --git a/crates/flowrs-airflow/src/client/v2/mod.rs b/crates/flowrs-airflow/src/client/v2/mod.rs index c690ac65..8a04fab8 100644 --- a/crates/flowrs-airflow/src/client/v2/mod.rs +++ b/crates/flowrs-airflow/src/client/v2/mod.rs @@ -7,15 +7,15 @@ mod log; mod task; mod taskinstance; -use anyhow::Result; -use reqwest::Method; +use reqwest::{Method, RequestBuilder, Response, Url}; use super::base::BaseClient; +use crate::error::Result; /// API v2 client implementation (for Airflow v3, uses /api/v2 endpoint) #[derive(Debug)] pub struct V2Client { - pub base: BaseClient, + base: BaseClient, } impl V2Client { @@ -25,18 +25,18 @@ impl V2Client { Self { base } } - pub(crate) async fn base_api( - &self, - method: Method, - endpoint: &str, - ) -> Result { + pub(crate) async fn base_api(&self, method: Method, endpoint: &str) -> Result { self.base .base_api(method, endpoint, Self::API_VERSION) .await } + pub(crate) async fn execute(&self, request: RequestBuilder) -> Result { + self.base.execute(request).await + } + /// Returns the base endpoint URL for this client - pub fn endpoint(&self) -> &str { - &self.base.config.endpoint + pub const fn endpoint(&self) -> &Url { + self.base.endpoint() } } diff --git a/crates/flowrs-airflow/src/client/v2/task.rs b/crates/flowrs-airflow/src/client/v2/task.rs index 4196f45f..de985cd0 100644 --- a/crates/flowrs-airflow/src/client/v2/task.rs +++ b/crates/flowrs-airflow/src/client/v2/task.rs @@ -1,19 +1,16 @@ -use anyhow::Result; use reqwest::Method; use super::model; use super::V2Client; +use crate::client::read_json; +use crate::error::Result; impl V2Client { pub async fn fetch_tasks(&self, dag_id: &str) -> Result { - let response = self + let request = self .base_api(Method::GET, &format!("dags/{dag_id}/tasks")) - .await? - .send() - .await? - .error_for_status()?; - - let task_collection: model::task::TaskCollectionResponse = response.json().await?; - Ok(task_collection) + .await?; + let response = self.execute(request).await?; + read_json(response, "tasks response").await } } diff --git a/crates/flowrs-airflow/src/client/v2/taskinstance.rs b/crates/flowrs-airflow/src/client/v2/taskinstance.rs index 49583004..5d2f7a67 100644 --- a/crates/flowrs-airflow/src/client/v2/taskinstance.rs +++ b/crates/flowrs-airflow/src/client/v2/taskinstance.rs @@ -1,9 +1,10 @@ -use anyhow::Result; use log::{debug, info}; -use reqwest::{Method, Response}; +use reqwest::Method; use super::model; use super::V2Client; +use crate::client::read_json; +use crate::error::Result; const PAGE_SIZE: usize = 100; @@ -19,18 +20,16 @@ impl V2Client { let mut total_entries; loop { - let response: Response = self + let request = self .base_api( Method::GET, &format!("dags/{dag_id}/dagRuns/{dag_run_id}/taskInstances"), ) .await? - .query(&[("limit", limit.to_string()), ("offset", offset.to_string())]) - .send() - .await? - .error_for_status()?; - - let page: model::taskinstance::TaskInstanceList = response.json().await?; + .query(&[("limit", limit.to_string()), ("offset", offset.to_string())]); + let response = self.execute(request).await?; + let page: model::taskinstance::TaskInstanceList = + read_json(response, "task instances response").await?; total_entries = page.total_entries; let fetched_count = page.task_instances.len(); @@ -67,15 +66,13 @@ impl V2Client { let mut total_entries; loop { - let response: Response = self + let request = self .base_api(Method::GET, "dags/~/dagRuns/~/taskInstances") .await? - .query(&[("limit", limit.to_string()), ("offset", offset.to_string())]) - .send() - .await? - .error_for_status()?; - - let page: model::taskinstance::TaskInstanceList = response.json().await?; + .query(&[("limit", limit.to_string()), ("offset", offset.to_string())]); + let response = self.execute(request).await?; + let page: model::taskinstance::TaskInstanceList = + read_json(response, "all task instances response").await?; total_entries = page.total_entries; let fetched_count = page.task_instances.len(); @@ -115,17 +112,15 @@ impl V2Client { dag_run_id: &str, task_id: &str, ) -> Result { - let response: Response = self + let request = self .base_api( Method::GET, &format!("dags/{dag_id}/dagRuns/{dag_run_id}/taskInstances/{task_id}/tries"), ) - .await? - .send() - .await? - .error_for_status()?; - - let tries: model::taskinstance::TaskInstanceTriesResponse = response.json().await?; + .await?; + let response = self.execute(request).await?; + let tries: model::taskinstance::TaskInstanceTriesResponse = + read_json(response, "task instance tries response").await?; debug!( "Fetched {} tries for task {task_id}", tries.task_instances.len() @@ -141,7 +136,7 @@ impl V2Client { task_id: &str, status: &str, ) -> Result<()> { - let resp: Response = self + let request = self .base_api( Method::PATCH, &format!("dags/{dag_id}/dagRuns/{dag_run_id}/taskInstances/{task_id}"), @@ -150,10 +145,8 @@ impl V2Client { // Airflow 3's `/api/v2` PatchTaskInstanceBody uses a strict schema that // forbids unknown fields; unlike the `/api/v1` endpoint it does not accept // `dry_run` (dry-run is a separate endpoint), so sending it yields a 422. - .json(&serde_json::json!({"new_state": status})) - .send() - .await? - .error_for_status()?; + .json(&serde_json::json!({"new_state": status})); + let resp = self.execute(request).await?; debug!("{resp:?}"); Ok(()) } @@ -164,7 +157,7 @@ impl V2Client { dag_run_id: &str, task_id: &str, ) -> Result<()> { - let resp: Response = self + let request = self .base_api(Method::POST, &format!("dags/{dag_id}/clearTaskInstances")) .await? .json(&serde_json::json!( @@ -176,10 +169,8 @@ impl V2Client { "only_failed": false, "reset_dag_runs": true, } - )) - .send() - .await? - .error_for_status()?; + )); + let resp = self.execute(request).await?; debug!("{resp:?}"); Ok(()) } diff --git a/crates/flowrs-airflow/src/error.rs b/crates/flowrs-airflow/src/error.rs new file mode 100644 index 00000000..b4c0d225 --- /dev/null +++ b/crates/flowrs-airflow/src/error.rs @@ -0,0 +1,201 @@ +//! The error type returned by every fallible operation in this crate. + +/// Result alias for this crate's operations. +pub type Result = std::result::Result; + +/// Everything that can go wrong while talking to an Airflow deployment. +/// +/// The variants are the categories a caller can meaningfully act on: a bad +/// endpoint is a config problem, a `Status` is the server rejecting the call, +/// and a `Decode` means the deployment speaks a schema we do not understand. +#[derive(Debug, thiserror::Error)] +#[non_exhaustive] +pub enum AirflowError { + /// The configured endpoint is not a usable base URL. + #[error("invalid endpoint URL '{url}': {source}")] + InvalidUrl { + url: String, + #[source] + source: url::ParseError, + }, + + /// The request could not be completed: DNS, TLS, connect, or timeout. + #[error("HTTP request failed: {source}")] + Http { + #[from] + source: reqwest::Error, + }, + + /// The server answered with a non-success status. + /// + /// `detail` holds the response body, which is where Airflow puts the reason a + /// call was rejected. `Response::error_for_status` throws that body away, so we + /// read it ourselves before turning the status into an error. + #[error("{method} {path} failed with HTTP {status}: {detail}")] + Status { + method: String, + path: String, + status: u16, + detail: String, + }, + + /// The response body did not match the schema this client expects. + #[error("failed to parse {context}: {source} (body starts with: {snippet})")] + Decode { + context: String, + snippet: String, + #[source] + source: serde_json::Error, + }, + + /// Obtaining credentials failed. + #[error("{provider} authentication failed: {message}")] + Auth { + provider: &'static str, + message: String, + }, + + /// Discovering servers from a managed service failed. + #[error("{service} discovery failed: {message}")] + Discovery { + service: &'static str, + message: String, + }, + + /// The configured auth or service needs a Cargo feature that was not compiled in. + #[error("{service} support is not compiled in; rebuild with the '{feature}' feature")] + FeatureNotEnabled { + service: &'static str, + feature: &'static str, + }, +} + +/// How much of a response body to keep in an error message. +const SNIPPET_LEN: usize = 1000; + +impl AirflowError { + pub(crate) fn invalid_url(url: impl Into, source: url::ParseError) -> Self { + Self::InvalidUrl { + url: url.into(), + source, + } + } + + pub(crate) fn status( + method: &reqwest::Method, + url: &url::Url, + status: reqwest::StatusCode, + body: &str, + ) -> Self { + Self::Status { + method: method.to_string(), + // Keep the query string so pagination and filter parameters stay + // visible in the error message. + path: match url.query() { + Some(query) => format!("{}?{query}", url.path()), + None => url.path().to_string(), + }, + status: status.as_u16(), + // An empty body would render as a message ending in a bare colon, so fall + // back to the status' canonical reason. + detail: if body.trim().is_empty() { + status + .canonical_reason() + .unwrap_or("no response body") + .to_string() + } else { + truncate(body) + }, + } + } + + pub(crate) fn decode( + context: impl Into, + body: &str, + source: serde_json::Error, + ) -> Self { + Self::Decode { + context: context.into(), + snippet: truncate(body), + source, + } + } + + /// Flattens an `anyhow` chain into `message`. + /// + /// The managed-service integrations use `anyhow` internally for their layered + /// `.context()` messages. Alternate formatting joins the whole chain into one + /// line, so no context is lost when only `Display` is rendered. + pub(crate) fn auth(provider: &'static str, source: &anyhow::Error) -> Self { + Self::Auth { + provider, + message: format!("{source:#}"), + } + } + + /// See [`AirflowError::auth`] for why the chain is flattened into a string. + #[cfg(any(feature = "conveyor", feature = "mwaa"))] + pub(crate) fn discovery(service: &'static str, source: &anyhow::Error) -> Self { + Self::Discovery { + service, + message: format!("{source:#}"), + } + } +} + +fn truncate(body: &str) -> String { + body.chars().take(SNIPPET_LEN).collect() +} + +#[cfg(test)] +mod tests { + use super::*; + + fn url() -> url::Url { + "http://localhost:8080/api/v2/dags/x/dagRuns" + .parse() + .expect("test URL") + } + + #[test] + fn status_error_preserves_the_query_string() { + let url: url::Url = "http://localhost:8080/api/v2/dags?limit=100&offset=200" + .parse() + .expect("test URL"); + let error = AirflowError::status( + &reqwest::Method::GET, + &url, + reqwest::StatusCode::BAD_REQUEST, + "bad request", + ); + assert!( + error.to_string().contains("?limit=100&offset=200"), + "got: {error}" + ); + } + + #[test] + fn status_error_surfaces_the_response_body() { + let error = AirflowError::status( + &reqwest::Method::POST, + &url(), + reqwest::StatusCode::UNPROCESSABLE_ENTITY, + r#"{"detail":"logical_date is in the past"}"#, + ); + assert!( + error.to_string().contains("logical_date is in the past"), + "got: {error}" + ); + } + + #[test] + fn status_error_without_a_body_falls_back_to_the_status_reason() { + let error = AirflowError::status( + &reqwest::Method::GET, + &url(), + reqwest::StatusCode::NOT_FOUND, + " ", + ); + assert!(error.to_string().ends_with("Not Found"), "got: {error}"); + } +} diff --git a/crates/flowrs-airflow/src/lib.rs b/crates/flowrs-airflow/src/lib.rs index 2456a317..676848a0 100644 --- a/crates/flowrs-airflow/src/lib.rs +++ b/crates/flowrs-airflow/src/lib.rs @@ -1,10 +1,12 @@ pub mod auth; pub mod client; pub mod config; +pub mod error; pub mod managed_services; pub use auth::{ AirflowAuth, AstronomerAuth, BasicAuth, ComposerAuth, MwaaAuth, MwaaTokenType, TokenSource, }; -pub use client::{create_api_client, AirflowApiClient, BaseClient, V1Client, V2Client}; +pub use client::{BaseClient, V1Client, V2Client}; pub use config::{AirflowConfig, AirflowVersion, GccConfig, ManagedService}; +pub use error::{AirflowError, Result}; diff --git a/crates/flowrs-airflow/src/managed_services/astronomer.rs b/crates/flowrs-airflow/src/managed_services/astronomer.rs index 8b6785ca..8505ca91 100644 --- a/crates/flowrs-airflow/src/managed_services/astronomer.rs +++ b/crates/flowrs-airflow/src/managed_services/astronomer.rs @@ -212,7 +212,7 @@ impl From<&AstronomerAuth> for AstronomerAuthProvider { #[async_trait] impl AuthProvider for AstronomerAuthProvider { - async fn authenticate(&self, request: RequestBuilder) -> Result { + async fn authenticate(&self, request: RequestBuilder) -> crate::error::Result { info!("🔑 Astronomer Auth"); Ok(request.bearer_auth(&self.api_token)) } diff --git a/crates/flowrs-airflow/src/managed_services/composer/provider.rs b/crates/flowrs-airflow/src/managed_services/composer/provider.rs index 12dde8d9..02c16798 100644 --- a/crates/flowrs-airflow/src/managed_services/composer/provider.rs +++ b/crates/flowrs-airflow/src/managed_services/composer/provider.rs @@ -1,5 +1,6 @@ use crate::client::auth::AuthProvider; -use anyhow::{Context, Result}; +use crate::error::{AirflowError, Result}; +use anyhow::Context; use async_trait::async_trait; use google_cloud_auth::credentials::{AccessTokenCredentials, Builder}; use log::info; @@ -14,10 +15,11 @@ pub struct ComposerAuthProvider { } impl ComposerAuthProvider { - pub fn new(auth: &ComposerAuth) -> Result { + pub fn new(auth: &ComposerAuth) -> crate::error::Result { let credentials = Builder::default() .build_access_token_credentials() - .context("Failed to build GCP credentials")?; + .context("Failed to build GCP credentials") + .map_err(|e| AirflowError::auth("Composer", &e))?; Ok(Self { project_id: auth.project_id.clone(), environment_name: auth.environment_name.clone(), @@ -47,7 +49,8 @@ impl AuthProvider for ComposerAuthProvider { .credentials .access_token() .await - .context("Failed to get GCP access token")?; + .context("Failed to get GCP access token") + .map_err(|e| AirflowError::auth("Composer", &e))?; Ok(request.bearer_auth(token.token)) } } diff --git a/crates/flowrs-airflow/src/managed_services/conveyor.rs b/crates/flowrs-airflow/src/managed_services/conveyor.rs index de639d26..2a8b32ee 100644 --- a/crates/flowrs-airflow/src/managed_services/conveyor.rs +++ b/crates/flowrs-airflow/src/managed_services/conveyor.rs @@ -1,6 +1,7 @@ use crate::auth::AirflowAuth; use crate::client::auth::AuthProvider; use crate::config::{AirflowConfig, ManagedService}; +use crate::error::AirflowError; use anyhow::{Context, Result}; use async_trait::async_trait; use dirs::home_dir; @@ -130,7 +131,7 @@ impl fmt::Debug for ConveyorAuthProvider { #[async_trait] impl AuthProvider for ConveyorAuthProvider { - async fn authenticate(&self, request: RequestBuilder) -> Result { + async fn authenticate(&self, request: RequestBuilder) -> crate::error::Result { let mut cached = self.cached.lock().await; let fresh = cached @@ -142,7 +143,9 @@ impl AuthProvider for ConveyorAuthProvider { // Run the blocking CLI off the async runtime. let token = tokio::task::spawn_blocking(ConveyorClient::get_token) .await - .context("Conveyor token task panicked")??; + .context("Conveyor token task panicked") + .and_then(|result| result) + .map_err(|e| AirflowError::auth("Conveyor", &e))?; *cached = Some((token, Instant::now())); } diff --git a/crates/flowrs-airflow/src/managed_services/expand.rs b/crates/flowrs-airflow/src/managed_services/expand.rs index a7f29529..e74af758 100644 --- a/crates/flowrs-airflow/src/managed_services/expand.rs +++ b/crates/flowrs-airflow/src/managed_services/expand.rs @@ -1,7 +1,7 @@ -use anyhow::Result; use log::info; use crate::config::{AirflowConfig, ManagedService}; +use crate::error::{AirflowError, Result}; #[cfg(feature = "astronomer")] use super::astronomer::get_astronomer_environment_servers; @@ -57,26 +57,38 @@ pub async fn expand_managed_services( ManagedService::Conveyor => { #[cfg(feature = "conveyor")] { - let conveyor_servers = get_conveyor_environment_servers()?; + let conveyor_servers = get_conveyor_environment_servers() + .map_err(|e| AirflowError::discovery("Conveyor", &e))?; all_servers.extend(conveyor_servers); } #[cfg(not(feature = "conveyor"))] { all_errors.push( - "Conveyor support not compiled. Enable the 'conveyor' feature.".to_string(), + AirflowError::FeatureNotEnabled { + service: "Conveyor", + feature: "conveyor", + } + .to_string(), ); } } ManagedService::Mwaa => { #[cfg(feature = "mwaa")] { - let mwaa_servers = get_mwaa_environment_servers().await?; + let mwaa_servers = get_mwaa_environment_servers() + .await + .map_err(|e| AirflowError::discovery("MWAA", &e))?; all_servers.extend(mwaa_servers); } #[cfg(not(feature = "mwaa"))] { - all_errors - .push("MWAA support not compiled. Enable the 'mwaa' feature.".to_string()); + all_errors.push( + AirflowError::FeatureNotEnabled { + service: "MWAA", + feature: "mwaa", + } + .to_string(), + ); } } ManagedService::Astronomer => { @@ -89,8 +101,11 @@ pub async fn expand_managed_services( #[cfg(not(feature = "astronomer"))] { all_errors.push( - "Astronomer support not compiled. Enable the 'astronomer' feature." - .to_string(), + AirflowError::FeatureNotEnabled { + service: "Astronomer", + feature: "astronomer", + } + .to_string(), ); } } @@ -139,8 +154,11 @@ pub async fn expand_managed_services( #[cfg(not(feature = "composer"))] { all_errors.push( - "Google Cloud Composer support not compiled. Enable the 'composer' feature." - .to_string(), + AirflowError::FeatureNotEnabled { + service: "Google Cloud Composer", + feature: "composer", + } + .to_string(), ); } } diff --git a/crates/flowrs-airflow/src/managed_services/mwaa.rs b/crates/flowrs-airflow/src/managed_services/mwaa.rs index bd001c43..5ed74d9f 100644 --- a/crates/flowrs-airflow/src/managed_services/mwaa.rs +++ b/crates/flowrs-airflow/src/managed_services/mwaa.rs @@ -187,7 +187,7 @@ impl From<&MwaaAuth> for MwaaAuthProvider { #[async_trait] impl AuthProvider for MwaaAuthProvider { - async fn authenticate(&self, request: RequestBuilder) -> Result { + async fn authenticate(&self, request: RequestBuilder) -> crate::error::Result { info!("🔑 MWAA Auth: {}", self.environment_name); match &self.token { MwaaTokenType::SessionCookie(cookie) => { diff --git a/src/airflow/client/impls/dag_ops.rs b/src/airflow/client/impls/dag_ops.rs index 60fe2af7..f638e764 100644 --- a/src/airflow/client/impls/dag_ops.rs +++ b/src/airflow/client/impls/dag_ops.rs @@ -24,22 +24,23 @@ impl DagOperations for FlowrsClient { async fn toggle_dag(&self, dag_id: &str, is_paused: bool) -> Result<()> { match self { - Self::V1(client) => client.patch_dag_pause(dag_id, is_paused).await, - Self::V2(client) => client.patch_dag_pause(dag_id, is_paused).await, + Self::V1(client) => client.patch_dag_pause(dag_id, is_paused).await?, + Self::V2(client) => client.patch_dag_pause(dag_id, is_paused).await?, } + Ok(()) } async fn get_dag_code(&self, dag: &Dag) -> Result { - match self { - Self::V1(client) => client.fetch_dag_code(&dag.file_token).await, - Self::V2(client) => client.fetch_dag_code(&dag.dag_id).await, - } + Ok(match self { + Self::V1(client) => client.fetch_dag_code(&dag.file_token).await?, + Self::V2(client) => client.fetch_dag_code(&dag.dag_id).await?, + }) } async fn get_dag_params(&self, dag_id: &str) -> Result> { - match self { - Self::V1(client) => client.fetch_dag_params(dag_id).await, - Self::V2(client) => client.fetch_dag_params(dag_id).await, - } + Ok(match self { + Self::V1(client) => client.fetch_dag_params(dag_id).await?, + Self::V2(client) => client.fetch_dag_params(dag_id).await?, + }) } } diff --git a/src/airflow/client/impls/dagrun_ops.rs b/src/airflow/client/impls/dagrun_ops.rs index fec7706d..fc57f34d 100644 --- a/src/airflow/client/impls/dagrun_ops.rs +++ b/src/airflow/client/impls/dagrun_ops.rs @@ -37,16 +37,18 @@ impl DagRunOperations for FlowrsClient { async fn mark_dag_run(&self, dag_id: &str, dag_run_id: &str, status: &str) -> Result<()> { match self { - Self::V1(client) => client.patch_dag_run(dag_id, dag_run_id, status).await, - Self::V2(client) => client.patch_dag_run(dag_id, dag_run_id, status).await, + Self::V1(client) => client.patch_dag_run(dag_id, dag_run_id, status).await?, + Self::V2(client) => client.patch_dag_run(dag_id, dag_run_id, status).await?, } + Ok(()) } async fn clear_dagrun(&self, dag_id: &str, dag_run_id: &str) -> Result<()> { match self { - Self::V1(client) => client.post_clear_dagrun(dag_id, dag_run_id).await, - Self::V2(client) => client.post_clear_dagrun(dag_id, dag_run_id).await, + Self::V1(client) => client.post_clear_dagrun(dag_id, dag_run_id).await?, + Self::V2(client) => client.post_clear_dagrun(dag_id, dag_run_id).await?, } + Ok(()) } async fn trigger_dag_run( @@ -59,13 +61,14 @@ impl DagRunOperations for FlowrsClient { Self::V1(client) => { client .post_trigger_dag_run(dag_id, logical_date, conf) - .await + .await?; } Self::V2(client) => { client .post_trigger_dag_run(dag_id, logical_date, conf) - .await + .await?; } } + Ok(()) } } diff --git a/src/airflow/client/impls/taskinstance_ops.rs b/src/airflow/client/impls/taskinstance_ops.rs index 18d74185..bbd823f4 100644 --- a/src/airflow/client/impls/taskinstance_ops.rs +++ b/src/airflow/client/impls/taskinstance_ops.rs @@ -84,14 +84,15 @@ impl TaskInstanceOperations for FlowrsClient { Self::V1(client) => { client .patch_task_instance(dag_id, dag_run_id, task_id, status) - .await + .await?; } Self::V2(client) => { client .patch_task_instance(dag_id, dag_run_id, task_id, status) - .await + .await?; } } + Ok(()) } async fn clear_task_instance( @@ -104,13 +105,14 @@ impl TaskInstanceOperations for FlowrsClient { Self::V1(client) => { client .post_clear_task_instance(dag_id, dag_run_id, task_id) - .await + .await?; } Self::V2(client) => { client .post_clear_task_instance(dag_id, dag_run_id, task_id) - .await + .await?; } } + Ok(()) } } diff --git a/src/airflow/client/open_url.rs b/src/airflow/client/open_url.rs index 89036aa7..fee41fd2 100644 --- a/src/airflow/client/open_url.rs +++ b/src/airflow/client/open_url.rs @@ -3,8 +3,8 @@ use url::{form_urlencoded, Url}; use crate::airflow::model::common::OpenItem; -pub(super) fn build_v1_open_url(endpoint: &str, item: &OpenItem) -> Result { - let mut base_url = Url::parse(endpoint)?; +pub(super) fn build_v1_open_url(endpoint: &Url, item: &OpenItem) -> Result { + let mut base_url = endpoint.clone(); match item { OpenItem::Config(config_endpoint) => { @@ -53,8 +53,8 @@ pub(super) fn build_v1_open_url(endpoint: &str, item: &OpenItem) -> Result Result { - let mut base_url = Url::parse(endpoint)?; +pub(super) fn build_v2_open_url(endpoint: &Url, item: &OpenItem) -> Result { + let mut base_url = endpoint.clone(); match item { OpenItem::Config(config_endpoint) => {