use std::{borrow::Cow, time::SystemTime};
use schemars::JsonSchema;
use serde::{de::DeserializeOwned, Serialize};
use crate::{
job_definition::{JobDefinition, RetryPolicy, TimeoutPolicy},
schedule_definition::ScheduleDefinition,
IndexMap,
};
pub trait JobType: Serialize + DeserializeOwned + JsonSchema + Send + Sync + 'static {
type Output: Serialize + DeserializeOwned + JsonSchema + Send + Sync + 'static;
fn id() -> &'static str;
#[must_use]
fn default_timeout_policy() -> TimeoutPolicy {
TimeoutPolicy::default()
}
#[must_use]
fn default_retry_policy() -> RetryPolicy {
RetryPolicy::default()
}
#[must_use]
fn default_labels() -> IndexMap<String, String> {
IndexMap::default()
}
}
pub trait JobTypeExt: JobType {
fn job(&self) -> TypedJobDefinition<Self>;
#[must_use]
fn metadata() -> JobTypeMetadata;
}
impl<J> JobTypeExt for J
where
J: JobType,
{
fn job(&self) -> TypedJobDefinition<Self> {
JobDefinition {
job_type_id: Self::id().into(),
target_execution_time: SystemTime::now(),
input_payload_json: serde_json::to_string(self).unwrap(),
labels: Self::default_labels(),
timeout_policy: Self::default_timeout_policy(),
retry_policy: Self::default_retry_policy(),
metadata_json: None,
}
.into()
}
fn metadata() -> JobTypeMetadata {
let input_schema = schemars::schema_for!(J);
JobTypeMetadata {
id: Self::id().into(),
name: input_schema
.schema
.metadata
.as_ref()
.and_then(|m| m.title.clone()),
description: input_schema
.schema
.metadata
.as_ref()
.and_then(|m| m.description.clone()),
input_schema_json: Some(serde_json::to_string(&input_schema).unwrap()),
output_schema_json: Some(
serde_json::to_string(&schemars::schema_for!(J::Output)).unwrap(),
),
}
}
}
#[derive(Debug, Clone)]
#[must_use]
pub struct TypedJobDefinition<J = ()> {
pub(crate) inner: JobDefinition,
_job_type: std::marker::PhantomData<fn() -> J>,
}
impl<J> TypedJobDefinition<J> {
pub fn at(mut self, target_execution_time: impl Into<std::time::SystemTime>) -> Self {
self.inner.target_execution_time = target_execution_time.into();
self
}
pub fn now(mut self) -> Self {
self.inner.target_execution_time = std::time::SystemTime::now();
self
}
pub fn with_label(mut self, key: impl Into<String>, value: impl Into<String>) -> Self {
self.inner.labels.insert(key.into(), value.into());
self
}
pub fn with_timeout(mut self, timeout: impl Into<std::time::Duration>) -> Self {
self.inner.timeout_policy.timeout = Some(timeout.into());
self
}
pub fn with_retries(mut self, retries: u64) -> Self {
self.inner.retry_policy.retries = retries;
self
}
pub fn into_inner(self) -> JobDefinition {
self.inner
}
pub fn cast_type<T>(self) -> TypedJobDefinition<T> {
TypedJobDefinition {
inner: self.inner,
_job_type: std::marker::PhantomData,
}
}
}
impl<J> TypedJobDefinition<J> {
pub fn repeat_every(self, interval: std::time::Duration) -> ScheduleDefinition {
self.inner.repeat_every(interval)
}
pub fn repeat_cron(
self,
cron_expression: impl Into<String>,
) -> eyre::Result<ScheduleDefinition> {
self.inner.repeat_cron(cron_expression)
}
}
impl<J> TypedJobDefinition<J>
where
J: JobType,
{
#[must_use]
pub fn input(&self) -> J {
serde_json::from_str(&self.inner.input_payload_json).unwrap()
}
pub fn try_input(&self) -> Result<J, serde_json::Error> {
serde_json::from_str(&self.inner.input_payload_json)
}
pub fn with_input(mut self, input: impl AsRef<J>) -> Self {
self.inner.input_payload_json = serde_json::to_string(input.as_ref()).unwrap();
self
}
#[must_use]
pub fn is_valid(&self) -> bool {
self.inner.job_type_id == J::id() && self.try_input().is_ok()
}
}
impl<J> From<JobDefinition> for TypedJobDefinition<J> {
fn from(job_definition: JobDefinition) -> Self {
Self {
inner: job_definition,
_job_type: std::marker::PhantomData,
}
}
}
impl<J> From<TypedJobDefinition<J>> for JobDefinition {
fn from(typed: TypedJobDefinition<J>) -> Self {
typed.inner
}
}
impl JobDefinition {
pub fn into_typed_unknown(self) -> TypedJobDefinition {
TypedJobDefinition {
inner: self,
_job_type: std::marker::PhantomData,
}
}
pub fn into_typed<J>(self) -> TypedJobDefinition<J> {
TypedJobDefinition {
inner: self,
_job_type: std::marker::PhantomData,
}
}
}
#[derive(Debug, Clone)]
pub struct JobTypeMetadata {
pub id: Cow<'static, str>,
pub name: Option<String>,
pub description: Option<String>,
pub input_schema_json: Option<String>,
pub output_schema_json: Option<String>,
}