use async_trait::async_trait;
use chrono::{DateTime, Utc};
use crate::error::Result;
use crate::models::{
Job, JobRevision, PartitionAssignment, Run, RunStatus, SchedulerLeader, Script, Worker,
};
#[async_trait]
pub trait SchedulerStore: Send + Sync {
async fn upsert_job(&self, job: &Job) -> Result<()>;
async fn get_job(&self, job_id: &str) -> Result<Option<Job>>;
async fn get_job_by_name(&self, job_name: &str) -> Result<Option<Job>>;
async fn list_jobs(&self) -> Result<Vec<Job>>;
async fn list_due_jobs(&self, before: DateTime<Utc>) -> Result<Vec<Job>>;
async fn pause_job(&self, job_id: &str) -> Result<()>;
async fn resume_job(&self, job_id: &str) -> Result<()>;
async fn create_run(&self, run: &Run) -> Result<()>;
async fn update_run(&self, run: &Run) -> Result<()>;
async fn get_run(&self, run_id: &str) -> Result<Option<Run>>;
async fn list_runs_for_job(&self, job_id: &str, limit: usize) -> Result<Vec<Run>>;
async fn list_runs_filtered(
&self,
job_id: Option<&str>,
status: Option<RunStatus>,
offset: usize,
limit: usize,
) -> Result<Vec<Run>>;
async fn claim_next_queued(
&self,
pool_id: &str,
worker_id: &str,
now: DateTime<Utc>,
lease_ttl_secs: i64,
) -> Result<Option<Run>>;
async fn claim_run_by_id(
&self,
run_id: &str,
pool_id: &str,
worker_id: &str,
now: DateTime<Utc>,
lease_ttl_secs: i64,
) -> Result<Option<Run>>;
async fn claim_runs_by_ids(
&self,
run_ids: &[&str],
pool_id: &str,
worker_id: &str,
now: DateTime<Utc>,
lease_ttl_secs: i64,
) -> Result<Vec<Run>> {
let mut claimed = Vec::new();
for run_id in run_ids {
if let Some(run) = self
.claim_run_by_id(run_id, pool_id, worker_id, now, lease_ttl_secs)
.await?
{
claimed.push(run);
}
}
Ok(claimed)
}
async fn renew_run_lease(
&self,
run_id: &str,
worker_id: &str,
now: DateTime<Utc>,
lease_ttl_secs: i64,
) -> Result<bool>;
async fn append_revision(&self, revision: &JobRevision) -> Result<()>;
async fn list_revisions(&self, job_id: &str) -> Result<Vec<JobRevision>>;
async fn upsert_script(&self, script: &Script) -> Result<()>;
async fn get_script(&self, script_name: &str) -> Result<Option<Script>>;
async fn try_claim_run_once(
&self,
job_id: &str,
claimed_by: &str,
now: DateTime<Utc>,
claim_ttl_secs: i64,
) -> Result<bool>;
async fn mark_run_once_completed(&self, job_id: &str, completed_at: DateTime<Utc>)
-> Result<()>;
async fn release_run_once_claim(
&self,
job_id: &str,
claimed_by: &str,
now: DateTime<Utc>,
) -> Result<()>;
async fn find_due_job_ids_in_partitions(
&self,
owned_partitions: &[u32],
due_until: DateTime<Utc>,
limit: u32,
) -> Result<Vec<String>>;
async fn min_next_run_at_in_partitions(
&self,
owned_partitions: &[u32],
) -> Result<Option<DateTime<Utc>>>;
async fn claim_job_for_tick(
&self,
job_id: &str,
claim_id: &str,
now: DateTime<Utc>,
lease_ttl_secs: i64,
) -> Result<bool>;
async fn release_job_tick_claim(&self, job_id: &str) -> Result<()>;
async fn persist_post_tick_job_state(
&self,
job_id: &str,
next_run_at: Option<DateTime<Utc>>,
) -> Result<()>;
async fn try_acquire_leader(&self, instance_id: &str, ttl_secs: i64) -> Result<bool>;
async fn renew_leader_lease(&self, instance_id: &str, ttl_secs: i64) -> Result<()>;
async fn get_leader(&self) -> Result<Option<SchedulerLeader>>;
async fn upsert_partition_assignment(&self, assignment: &PartitionAssignment) -> Result<()>;
async fn list_partition_assignments(&self) -> Result<Vec<PartitionAssignment>>;
async fn register_worker(&self, worker: &Worker) -> Result<()>;
async fn heartbeat_worker(&self, worker_id: &str, at: DateTime<Utc>) -> Result<()>;
}