fix(vine): bugs in updating vine-session

ulagbulag · Jul 8, 2024 · 1133a24 · 1133a24
1 parent 1faa154
commit 1133a24
Show file tree

Hide file tree

Showing 5 changed files with 44 additions and 23 deletions.
diff --git a/crates/vine/plugin/src/routes/desktop/single.rs b/crates/vine/plugin/src/routes/desktop/single.rs
@@ -8,7 +8,7 @@ use kube::Client;
 use tracing::{instrument, Level};
 use vine_api::user_session::{UserSessionCommand, UserSessionRef};
 use vine_rbac::auth::{AuthUserSession, AuthUserSessionRef};
-use vine_session::SessionExec;
+use vine_session::exec::SessionExecExt;
 
 #[instrument(level = Level::INFO, skip(request, kube))]
 #[post("/user/desktop/exec")]
@@ -26,6 +26,6 @@ pub async fn post_exec(
         Err(error) => return HttpResponse::from(Result::<()>::Err(error.to_string())),
     };
 
-    let result = session.exec(kube, command).await.map(|_| ());
+    let result = session.exec_without_tty(kube, command).await.map(|_| ());
     HttpResponse::from(Result::from(result))
 }
diff --git a/crates/vine/plugin/src/routes/session.rs b/crates/vine/plugin/src/routes/session.rs
@@ -5,7 +5,7 @@ use kube::Client;
 use tracing::{instrument, Level};
 use vine_api::user_session::UserSessionMetadata;
 use vine_rbac::auth::AuthUserSession;
-use vine_session::SessionExec;
+use vine_session::exec::SessionExec;
 
 #[instrument(level = Level::INFO, skip(request, kube))]
 #[get("/batch/user/session")]

diff --git a/crates/vine/session/src/batch.rs b/crates/vine/session/src/batch.rs
@@ -4,15 +4,12 @@ use anyhow::{anyhow, Error, Result};
 use ark_api::{NamespaceAny, SessionRef};
 use futures::{stream::FuturesUnordered, StreamExt};
 use k8s_openapi::api::core::v1::Node;
-use kube::{
-    api::{AttachParams, ListParams},
-    Api, Client, ResourceExt,
-};
+use kube::{api::ListParams, Api, Client, ResourceExt};
 use regex::Regex;
 use tokio::spawn;
 use tracing::{debug, error, instrument, warn, Level};
 
-use crate::exec::SessionExec;
+use crate::exec::SessionExecExt;
 
 pub struct BatchCommandArgs<C, U> {
     pub command: C,
@@ -54,15 +51,8 @@ impl<C, U> BatchCommandArgs<C, U> {
 
         let processes = sessions_filtered.into_iter().map(|session| {
             let kube = kube.clone();
-            let ap = AttachParams {
-                stdin: false,
-                stdout: true,
-                stderr: true,
-                tty: false,
-                ..Default::default()
-            };
             let command = command.clone();
-            spawn(async move { session.exec(kube, ap, command).await })
+            spawn(async move { session.exec_without_tty(kube, command).await })
         });
 
         processes

diff --git a/crates/vine/session/src/exec.rs b/crates/vine/session/src/exec.rs
@@ -13,6 +13,41 @@ use tokio::{spawn, task::yield_now};
 use tracing::{instrument, Level};
 use vine_api::user::UserCrd;
 
+#[async_trait]
+pub trait SessionExecExt
+where
+    Self: fmt::Debug + SessionExec,
+{
+    #[instrument(level = Level::INFO, skip(kube, command), err(Display))]
+    async fn exec_with_tty<I>(&self, kube: Client, command: I) -> Result<Vec<Process>>
+    where
+        I: 'static + Send + Sync + Clone + fmt::Debug + IntoIterator,
+        <I as IntoIterator>::Item: Sync + Into<String>,
+    {
+        let ap = AttachParams::interactive_tty();
+        <Self as SessionExec>::exec(self, kube, ap, command).await
+    }
+
+    #[instrument(level = Level::INFO, skip(kube, command), err(Display))]
+    async fn exec_without_tty<I>(&self, kube: Client, command: I) -> Result<Vec<Process>>
+    where
+        I: 'static + Send + Sync + Clone + fmt::Debug + IntoIterator,
+        <I as IntoIterator>::Item: Sync + Into<String>,
+    {
+        let ap = AttachParams {
+            stdin: false,
+            stdout: true,
+            stderr: true,
+            tty: false,
+            ..Default::default()
+        };
+        <Self as SessionExec>::exec(self, kube, ap, command).await
+    }
+}
+
+#[async_trait]
+impl<T> SessionExecExt for T where Self: fmt::Debug + SessionExec {}
+
 #[async_trait]
 pub trait SessionExec {
     async fn list(kube: Client) -> Result<Vec<Self>>

diff --git a/crates/vine/session/src/shell.rs b/crates/vine/session/src/shell.rs
@@ -4,10 +4,7 @@ use anyhow::{Error, Result};
 use avt::Vt;
 use chrono::Utc;
 use futures::{channel::mpsc, stream::FuturesUnordered, SinkExt, StreamExt};
-use kube::{
-    api::{AttachParams, TerminalSize},
-    Client,
-};
+use kube::{api::TerminalSize, Client};
 use ratatui::{
     backend::CrosstermBackend,
     buffer::Buffer,
@@ -31,7 +28,7 @@ use tracing::{error, info};
 
 use crate::{
     batch::{collect_user_sessions, BatchCommandUsers},
-    exec::{Process, SessionExec},
+    exec::{Process, SessionExecExt},
 };
 
 pub struct BatchShellArgs<C, U> {
@@ -59,9 +56,8 @@ impl<C, U> BatchShellArgs<C, U> {
             .into_iter()
             .map(|session| {
                 let kube = kube.clone();
-                let ap = AttachParams::interactive_tty();
                 let command = [command.clone()];
-                async move { session.exec(kube, ap, command).await }
+                async move { session.exec_with_tty(kube, command).await }
             })
             .collect::<FuturesUnordered<_>>()
             .filter_map(|result| async move {