tatara_engine/drivers/
kube.rs1use anyhow::Result;
8use async_trait::async_trait;
9use chrono::Utc;
10use std::path::Path;
11use std::time::Duration;
12use tokio::sync::mpsc;
13
14use tatara_core::domain::allocation::TaskRunState;
15use tatara_core::domain::job::Task;
16
17use super::{Driver, LogEntry, TaskHandle};
18use tatara_core::domain::job::DriverType;
19
20pub struct KubeDriver {
26 kubeconfig: Option<String>,
28}
29
30impl KubeDriver {
31 pub fn new() -> Self {
32 Self { kubeconfig: None }
33 }
34
35 pub fn with_kubeconfig(kubeconfig: impl Into<String>) -> Self {
36 Self {
37 kubeconfig: Some(kubeconfig.into()),
38 }
39 }
40
41 async fn check_kubeconfig(&self) -> bool {
43 if let Some(ref path) = self.kubeconfig {
45 return tokio::fs::metadata(path).await.is_ok();
46 }
47
48 if let Ok(env_path) = std::env::var("KUBECONFIG") {
50 if !env_path.is_empty() {
51 return tokio::fs::metadata(&env_path).await.is_ok();
52 }
53 }
54
55 if let Some(home) = dirs::home_dir() {
57 let default = home.join(".kube").join("config");
58 return default.exists();
59 }
60
61 false
62 }
63}
64
65impl Default for KubeDriver {
66 fn default() -> Self {
67 Self::new()
68 }
69}
70
71#[async_trait]
72impl Driver for KubeDriver {
73 fn name(&self) -> &str {
74 "kube"
75 }
76
77 async fn available(&self) -> bool {
78 self.check_kubeconfig().await
79 }
80
81 async fn start(&self, task: &Task, _alloc_dir: &Path) -> Result<TaskHandle> {
82 tracing::info!(
89 task = %task.name,
90 driver = "kube",
91 "starting K8s workload"
92 );
93
94 Ok(TaskHandle {
95 driver: DriverType::Kube,
96 pid: None,
97 container_id: Some(format!("kube:{}", task.name)),
98 started_at: Utc::now(),
99 })
100 }
101
102 async fn stop(&self, handle: &TaskHandle, _timeout: Duration) -> Result<()> {
103 tracing::info!(
104 container_id = ?handle.container_id,
105 "stopping K8s workload"
106 );
107 Ok(())
109 }
110
111 async fn status(&self, handle: &TaskHandle) -> Result<TaskRunState> {
112 let _ = handle;
114 Ok(TaskRunState::Running)
115 }
116
117 async fn logs(&self, _handle: &TaskHandle) -> Result<mpsc::Receiver<LogEntry>> {
118 let (tx, rx) = mpsc::channel(100);
119 drop(tx);
121 Ok(rx)
122 }
123}
124
125#[cfg(test)]
126mod tests {
127 use super::*;
128
129 #[test]
130 fn test_kube_driver_name() {
131 let driver = KubeDriver::new();
132 assert_eq!(driver.name(), "kube");
133 }
134
135 #[test]
136 fn test_kube_driver_with_kubeconfig() {
137 let driver = KubeDriver::with_kubeconfig("/path/to/config");
138 assert_eq!(driver.kubeconfig.as_deref(), Some("/path/to/config"));
139 }
140}