Skip to main content

tatara_engine/drivers/
kube.rs

1//! Kubernetes execution driver.
2//!
3//! Implements the Driver trait for Kubernetes workloads via the kube-rs
4//! client. Leverages tatara-kube's Server-Side Apply reconciler for
5//! resource management.
6
7use 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
20/// Kubernetes execution driver.
21///
22/// Manages workloads on K8s clusters via Server-Side Apply. Requires
23/// a valid kubeconfig. Uses tatara-kube's reconciler for resource
24/// lifecycle management.
25pub struct KubeDriver {
26    /// Kubeconfig path (None = default location).
27    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    /// Check if kubectl/kubeconfig is available.
42    async fn check_kubeconfig(&self) -> bool {
43        // 1. Explicit kubeconfig path from constructor
44        if let Some(ref path) = self.kubeconfig {
45            return tokio::fs::metadata(path).await.is_ok();
46        }
47
48        // 2. KUBECONFIG environment variable
49        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        // 3. Default location
56        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        // In the full implementation, this would:
83        // 1. Parse the task's flake_ref as a K8s manifest source
84        // 2. Use tatara-kube's reconciler to apply via SSA
85        // 3. Wait for the resource to be ready
86        // 4. Return a handle with the resource reference
87
88        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        // Would delete the K8s resource via tatara-kube
108        Ok(())
109    }
110
111    async fn status(&self, handle: &TaskHandle) -> Result<TaskRunState> {
112        // Would query pod status via kube-rs
113        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        // Would stream pod logs via kube-rs
120        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}