use std::collections::BTreeMap;
use std::num::NonZeroU64;
use onetaskgraph_plugin_api::{SecretResolver, SourceError, SourceName, SourcePlugin, TaskSource};
use schemars::{JsonSchema, Schema, schema_for};
use secrecy::ExposeSecret;
use serde::{Deserialize, Serialize};
use serde_json::Value;
use super::source::{RequestDeadline, SubprocessSource};
use crate::secrets::CredentialName;
pub(crate) const KIND: &str = "subprocess";
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize, JsonSchema)]
#[serde(deny_unknown_fields)]
pub struct SubprocessConfig {
pub command: Program,
#[serde(default)]
pub args: Vec<String>,
#[serde(default)]
pub secrets: Vec<CredentialName>,
#[serde(default)]
pub settings: Value,
#[serde(default = "default_deadline_ms")]
pub deadline_ms: NonZeroU64,
}
fn default_deadline_ms() -> NonZeroU64 {
RequestDeadline::DEFAULT.milliseconds()
}
#[derive(Debug, Clone, PartialEq, Eq, Hash, Serialize, Deserialize, JsonSchema)]
#[serde(into = "String", try_from = "String")]
#[schemars(with = "String")]
pub struct Program(String);
impl Program {
#[must_use]
pub fn new(command: &str) -> Option<Self> {
(!command.trim().is_empty()).then(|| Self(command.to_owned()))
}
#[must_use]
pub fn as_str(&self) -> &str {
&self.0
}
}
impl TryFrom<String> for Program {
type Error = String;
fn try_from(value: String) -> Result<Self, Self::Error> {
Self::new(&value)
.ok_or_else(|| "`command` must name the program that serves this source".to_owned())
}
}
impl From<Program> for String {
fn from(value: Program) -> Self {
value.0
}
}
#[derive(Debug, Clone, Copy, Default)]
pub struct Plugin;
impl SourcePlugin for Plugin {
fn kind(&self) -> &'static str {
KIND
}
fn config_schema(&self) -> Schema {
schema_for!(SubprocessConfig)
}
fn build(
&self,
name: &SourceName,
config: &Value,
secrets: &dyn SecretResolver,
) -> Result<Box<dyn TaskSource>, SourceError> {
let config: SubprocessConfig =
serde_json::from_value(config.clone()).map_err(|error| SourceError::Config {
message: format!("source {name}: {error}"),
})?;
let forwarded = resolve_named(name, &config.secrets, secrets)?;
SubprocessSource::connect_with_deadline(
config.command.as_str(),
&config.args,
name,
&config.settings,
forwarded,
RequestDeadline::from_millis(config.deadline_ms),
)
.map(|source| Box::new(source) as Box<dyn TaskSource>)
}
}
fn resolve_named(
name: &SourceName,
named: &[CredentialName],
secrets: &dyn SecretResolver,
) -> Result<BTreeMap<String, String>, SourceError> {
let mut forwarded = BTreeMap::new();
for variable in named {
let value = secrets
.get(variable.as_str())
.ok_or_else(|| SourceError::Auth {
message: format!(
"source {name}: nothing defines {variable}, which this source's \
`secrets` names; export it, or add it to the credentials file"
),
})?;
forwarded.insert(
variable.as_str().to_owned(),
value.expose_secret().to_owned(),
);
}
Ok(forwarded)
}