286 lines
7.7 KiB
Rust
286 lines
7.7 KiB
Rust
//! In-memory storage layer for Supervisor
|
|
//!
|
|
//! Provides CRUD operations for:
|
|
//! - API Keys
|
|
//! - Runners
|
|
//! - Jobs
|
|
|
|
use crate::auth::{ApiKey, ApiKeyScope};
|
|
use crate::error::{SupervisorError, SupervisorResult};
|
|
use hero_job::Job;
|
|
use std::collections::{HashMap, HashSet};
|
|
|
|
/// In-memory storage for all supervisor data
|
|
pub struct Store {
|
|
/// API keys (key_value -> ApiKey)
|
|
api_keys: HashMap<String, ApiKey>,
|
|
/// Registered runner IDs
|
|
runners: HashSet<String>,
|
|
/// In-memory job storage (job_id -> Job)
|
|
jobs: HashMap<String, Job>,
|
|
}
|
|
|
|
impl Store {
|
|
/// Create a new store
|
|
pub fn new() -> Self {
|
|
Self {
|
|
api_keys: HashMap::new(),
|
|
runners: HashSet::new(),
|
|
jobs: HashMap::new(),
|
|
}
|
|
}
|
|
|
|
// ==================== API Key Operations ====================
|
|
|
|
/// Create an API key with a specific value
|
|
pub fn key_create(&mut self, key: ApiKey) -> ApiKey {
|
|
self.api_keys.insert(key.name.clone(), key.clone());
|
|
key
|
|
}
|
|
|
|
/// Create a new API key with generated UUID
|
|
pub fn key_create_new(&mut self, name: String, scope: ApiKeyScope) -> ApiKey {
|
|
let key = ApiKey::new(name, scope);
|
|
self.api_keys.insert(key.name.clone(), key.clone());
|
|
key
|
|
}
|
|
|
|
/// Get an API key by its value
|
|
pub fn key_get(&self, key_name: &str) -> Option<&ApiKey> {
|
|
self.api_keys.get(key_name)
|
|
}
|
|
|
|
/// Delete an API key
|
|
pub fn key_delete(&mut self, key_name: &str) -> Option<ApiKey> {
|
|
self.api_keys.remove(key_name)
|
|
}
|
|
|
|
/// List all API keys
|
|
pub fn key_list(&self) -> Vec<ApiKey> {
|
|
self.api_keys.values().cloned().collect()
|
|
}
|
|
|
|
/// List API keys by scope
|
|
pub fn key_list_by_scope(&self, scope: ApiKeyScope) -> Vec<ApiKey> {
|
|
self.api_keys
|
|
.values()
|
|
.filter(|k| k.scope == scope)
|
|
.cloned()
|
|
.collect()
|
|
}
|
|
|
|
// ==================== Runner Operations ====================
|
|
|
|
/// Add a runner
|
|
pub fn runner_add(&mut self, runner_id: String) -> SupervisorResult<()> {
|
|
self.runners.insert(runner_id);
|
|
Ok(())
|
|
}
|
|
|
|
/// Remove a runner
|
|
pub fn runner_remove(&mut self, runner_id: &str) -> SupervisorResult<()> {
|
|
self.runners.remove(runner_id);
|
|
Ok(())
|
|
}
|
|
|
|
/// Check if a runner exists
|
|
pub fn runner_exists(&self, runner_id: &str) -> bool {
|
|
self.runners.contains(runner_id)
|
|
}
|
|
|
|
/// List all runner IDs
|
|
pub fn runner_list_all(&self) -> Vec<String> {
|
|
self.runners.iter().cloned().collect()
|
|
}
|
|
|
|
// ==================== Job Operations ====================
|
|
|
|
/// Store a job in memory
|
|
pub fn job_store(&mut self, job: Job) -> SupervisorResult<()> {
|
|
self.jobs.insert(job.id.clone(), job);
|
|
Ok(())
|
|
}
|
|
|
|
/// Get a job from memory
|
|
pub fn job_get(&self, job_id: &str) -> SupervisorResult<Job> {
|
|
self.jobs
|
|
.get(job_id)
|
|
.cloned()
|
|
.ok_or_else(|| SupervisorError::JobNotFound {
|
|
job_id: job_id.to_string(),
|
|
})
|
|
}
|
|
|
|
/// Delete a job from memory
|
|
pub fn job_delete(&mut self, job_id: &str) -> SupervisorResult<()> {
|
|
self.jobs
|
|
.remove(job_id)
|
|
.ok_or_else(|| SupervisorError::JobNotFound {
|
|
job_id: job_id.to_string(),
|
|
})?;
|
|
Ok(())
|
|
}
|
|
|
|
/// List all job IDs
|
|
pub fn job_list(&self) -> Vec<String> {
|
|
self.jobs.keys().cloned().collect()
|
|
}
|
|
|
|
/// Check if a job exists
|
|
pub fn job_exists(&self, job_id: &str) -> bool {
|
|
self.jobs.contains_key(job_id)
|
|
}
|
|
}
|
|
|
|
impl Clone for Store {
|
|
fn clone(&self) -> Self {
|
|
Self {
|
|
api_keys: self.api_keys.clone(),
|
|
runners: self.runners.clone(),
|
|
jobs: self.jobs.clone(),
|
|
}
|
|
}
|
|
}
|
|
|
|
#[cfg(test)]
|
|
mod tests {
|
|
use super::*;
|
|
use hero_job::JobBuilder;
|
|
|
|
fn create_test_store() -> Store {
|
|
Store::new()
|
|
}
|
|
|
|
fn create_test_job(id: &str, runner: &str) -> Job {
|
|
let job = JobBuilder::new()
|
|
.caller_id("test_caller")
|
|
.context_id("test_context")
|
|
.runner(runner)
|
|
.payload("test payload")
|
|
.build()
|
|
.unwrap();
|
|
job.id = id.to_string(); // Set ID manually
|
|
job
|
|
}
|
|
|
|
#[test]
|
|
fn test_api_key_operations() {
|
|
let mut store = create_test_store();
|
|
|
|
// Create key
|
|
let key = store.key_create_new("test_key".to_string(), ApiKeyScope::Admin);
|
|
assert_eq!(key.name, "test_key");
|
|
assert_eq!(key.scope, ApiKeyScope::Admin);
|
|
|
|
// Get key
|
|
let retrieved = store.key_get(&key.key);
|
|
assert!(retrieved.is_some());
|
|
assert_eq!(retrieved.unwrap().name, "test_key");
|
|
|
|
// List keys
|
|
let keys = store.key_list();
|
|
assert_eq!(keys.len(), 1);
|
|
|
|
// List by scope
|
|
let admin_keys = store.key_list_by_scope(ApiKeyScope::Admin);
|
|
assert_eq!(admin_keys.len(), 1);
|
|
|
|
// Delete key
|
|
let removed = store.key_delete(&key.key);
|
|
assert!(removed.is_some());
|
|
assert!(store.key_get(&key.key).is_none());
|
|
}
|
|
|
|
#[test]
|
|
fn test_runner_operations() {
|
|
let mut store = create_test_store();
|
|
|
|
// Add runner
|
|
assert!(store.runner_add("runner1".to_string()).is_ok());
|
|
assert!(store.runner_exists("runner1"));
|
|
|
|
// List runners
|
|
let runners = store.runner_list_all();
|
|
assert_eq!(runners.len(), 1);
|
|
assert!(runners.contains(&"runner1".to_string()));
|
|
|
|
// List all runners
|
|
let all_runners = store.runner_list_all();
|
|
assert_eq!(all_runners.len(), 1);
|
|
|
|
// Remove runner
|
|
assert!(store.runner_remove("runner1").is_ok());
|
|
assert!(!store.runner_exists("runner1"));
|
|
}
|
|
|
|
#[test]
|
|
fn test_job_operations() {
|
|
let mut store = create_test_store();
|
|
let job = create_test_job("job1", "runner1");
|
|
|
|
// Store job
|
|
assert!(store.job_store(job.clone()).is_ok());
|
|
assert!(store.job_exists("job1"));
|
|
|
|
// Get job
|
|
let retrieved = store.job_get("job1");
|
|
assert!(retrieved.is_ok());
|
|
assert_eq!(retrieved.unwrap().id, "job1");
|
|
|
|
// List jobs
|
|
let jobs = store.job_list();
|
|
assert_eq!(jobs.len(), 1);
|
|
assert!(jobs.contains(&"job1".to_string()));
|
|
|
|
// Delete job
|
|
assert!(store.job_delete("job1").is_ok());
|
|
assert!(!store.job_exists("job1"));
|
|
assert!(store.job_get("job1").is_err());
|
|
}
|
|
|
|
#[test]
|
|
fn test_job_not_found() {
|
|
let store = create_test_store();
|
|
let result = store.job_get("nonexistent");
|
|
assert!(result.is_err());
|
|
}
|
|
|
|
#[test]
|
|
fn test_multiple_jobs() {
|
|
let mut store = create_test_store();
|
|
|
|
// Add multiple jobs
|
|
for i in 1..=3 {
|
|
let job = create_test_job(&format!("job{}", i), "runner1");
|
|
assert!(store.job_store(job).is_ok());
|
|
}
|
|
|
|
// Verify all exist
|
|
assert_eq!(store.job_list().len(), 3);
|
|
assert!(store.job_exists("job1"));
|
|
assert!(store.job_exists("job2"));
|
|
assert!(store.job_exists("job3"));
|
|
|
|
// Delete one
|
|
assert!(store.job_delete("job2").is_ok());
|
|
assert_eq!(store.job_list().len(), 2);
|
|
assert!(!store.job_exists("job2"));
|
|
}
|
|
|
|
#[test]
|
|
fn test_store_clone() {
|
|
let mut store = create_test_store();
|
|
store.runner_add("runner1".to_string()).unwrap();
|
|
|
|
let job = create_test_job("job1", "runner1");
|
|
store.job_store(job).unwrap();
|
|
|
|
// Clone the store
|
|
let cloned = store.clone();
|
|
|
|
// Verify cloned data
|
|
assert!(cloned.runner_exists("runner1"));
|
|
assert!(cloned.job_exists("job1"));
|
|
}
|
|
}
|