feat(linux-editor): switch a query tab's database on SQL Server, ClickHouse, Cassandra and MongoDB

linux/changelog.d/3898.added.md

@@ -0,0 +1 @@
+The SQL editor's database menu also moves a tab on SQL Server, ClickHouse, Cassandra, ScyllaDB and MongoDB.

linux/changelog.d/3898.fixed.md

@@ -0,0 +1 @@
+On ClickHouse, a USE in the SQL editor now applies to the statements that follow it.

linux/crates/core/src/sql_syntax/catalog_change.rs

@@ -39,6 +39,22 @@ pub fn may_change_database(text: &str, grammar: SqlGrammar) -> bool {
         .any(|token| token.is_word("USE"))
 }
 
+/// The database the last `USE <name>` in `text` names, for an engine
+/// whose client, not its server, keeps the database a request runs in. A
+/// name that would need unescaping is left out.
+pub fn used_database(text: &str, grammar: SqlGrammar) -> Option<String> {
+    lex_statements(text, grammar)
+        .statements
+        .iter()
+        .rev()
+        .find_map(|statement| match statement.as_slice() {
+            [first, name] if first.is_word("USE") => name.name(),
+            _ => None,
+        })
+        .filter(|name| !name.is_empty() && !name.contains(['`', '"', '\\']))
+        .map(str::to_owned)
+}
+
 fn plan_changes(text: &str, plan: &ScriptPlan) -> bool {
     let is = |word: &str, set: &[&str]| set.iter().any(|verb| word.eq_ignore_ascii_case(verb));
     plan.statements().iter().any(|statement| {
@@ -60,6 +76,16 @@ fn plan_changes(text: &str, plan: &ScriptPlan) -> bool {
 mod tests {
     use super::*;
 
+    #[test]
+    fn a_use_names_its_database() {
+        let grammar = SqlGrammar::ClickHouse;
+        assert_eq!(used_database("USE analytics", grammar).as_deref(), Some("analytics"));
+        assert_eq!(used_database("use `my db`;", grammar).as_deref(), Some("my db"));
+        assert_eq!(used_database("USE a; USE b", grammar).as_deref(), Some("b"));
+        assert_eq!(used_database("SELECT 1", grammar), None);
+        assert_eq!(used_database("USE `a``b`", grammar), None, "a name to unescape");
+    }
+
     #[test]
     fn a_use_can_change_the_database() {
         for grammar in SqlGrammar::SQL {

linux/crates/core/src/sql_syntax/mod.rs

@@ -17,7 +17,7 @@ mod statement_access;
 mod statement_tokens;
 
 pub use builtin_snippet::{BuiltinSnippet, builtin_snippets};
-pub use catalog_change::{changes_catalog, may_change_database};
+pub use catalog_change::{changes_catalog, may_change_database, used_database};
 pub use credential_rules::CredentialRules;
 pub use errors::{ExpressionSyntaxError, TypeSyntaxError};
 pub use grammar::SqlGrammar;

linux/crates/drivers/cassandra/src/connection.rs

@@ -243,6 +243,12 @@ impl Connection for CqlConnection {
         outcome
     }
 
+    /// The schemas this engine lists are the server's databases, which a
+    /// query tab can move between.
+    async fn list_databases(&self, options: CallOptions) -> Result<Vec<SchemaInfo>, DriverError> {
+        self.list_schemas(options).await
+    }
+
     async fn list_schemas(&self, options: CallOptions) -> Result<Vec<SchemaInfo>, DriverError> {
         let default = self.endpoint.keyspace().map(str::to_owned);
         self.read(options, move |link, _| async move {

linux/crates/drivers/cassandra/src/editor.rs

@@ -4,8 +4,8 @@
 //! only that tab, as cqlsh moves only its own connection. CQL keeps no
 //! other state on a connection: no transaction, no temporary table.
 
-use std::sync::Arc;
 use std::sync::atomic::{AtomicU8, Ordering};
+use std::sync::{Arc, Mutex, PoisonError};
 use std::time::Duration;
 
 use async_trait::async_trait;
@@ -21,7 +21,7 @@ use tablepro_core::{
 use tokio::time::Instant;
 use tokio_util::sync::CancellationToken;
 
-use crate::errors::execution_error;
+use crate::errors::{execution_error, use_keyspace_error};
 use crate::rows::{ColumnsRead, MAX_PAGE_ROWS, next_state, page_size};
 use crate::session::{CqlLink, LinkSlot};
 
@@ -40,12 +40,17 @@ pub(crate) struct CqlEditorSession {
     statement_timeout: Option<Duration>,
     state: Arc<AtomicU8>,
     shutdown: CancellationToken,
+    /// The keyspace the server confirmed. The driver records one it is
+    /// still moving to, or failed to.
+    keyspace: Arc<Mutex<Option<String>>>,
 }
 
 impl CqlEditorSession {
     pub(crate) fn new(link: CqlLink, read_only: bool, statement_timeout: Option<Duration>) -> Self {
+        let keyspace = link.session().get_keyspace().map(|keyspace| keyspace.to_string());
         Self {
             id: EditorSessionId::new(),
+            keyspace: Arc::new(Mutex::new(keyspace)),
             link: LinkSlot::new(link),
             read_only,
             statement_timeout,
@@ -74,10 +79,39 @@ impl EditorSession for CqlEditorSession {
         self.id
     }
 
-    /// The driver moves the session itself when a `USE` answers.
     fn database(&self) -> Option<String> {
-        let link = self.link.get().ok()?;
-        link.session().get_keyspace().map(|keyspace| keyspace.to_string())
+        self.keyspace.lock().unwrap_or_else(PoisonError::into_inner).clone()
+    }
+
+    /// The driver checks the name, which CQL takes as letters, digits and
+    /// `_` alone, and moves every connection of the session. A move that
+    /// does not finish leaves them in keyspaces nothing knows, so the
+    /// session ends.
+    async fn use_database(&self, database: &str, options: CallOptions) -> Result<(), DriverError> {
+        let link = self.link.get()?;
+        self.begin()?;
+        let using = link.session().use_keyspace(database, true);
+        let used = tokio::select! {
+            biased;
+            () = options.cancel().cancelled() => Err(DriverError::Cancelled),
+            () = until(options.deadline()) => Err(DriverError::Timeout {
+                phase: tablepro_core::TimeoutPhase::Statement,
+                server_cancelled: false,
+            }),
+            used = using => used.map_err(|error| use_keyspace_error(error, LossPhase::Statement)),
+        };
+        match used {
+            Ok(()) => {
+                *self.keyspace.lock().unwrap_or_else(PoisonError::into_inner) = Some(database.to_owned());
+                self.state.store(READY, Ordering::SeqCst);
+                Ok(())
+            }
+            Err(error) => {
+                self.state.store(LOST, Ordering::SeqCst);
+                self.link.release();
+                Err(error)
+            }
+        }
     }
 
     fn state(&self) -> EditorSessionState {
@@ -106,8 +140,15 @@ impl EditorSession for CqlEditorSession {
             shutdown: self.shutdown.clone(),
         };
         let state = self.state.clone();
+        let (keyspace, link) = (self.keyspace.clone(), run.link.clone());
+        let follows = tablepro_core::sql_syntax::may_change_database(sql.as_str(), SqlGrammar::Cql);
         let running = tokio::spawn(async move {
             let outcome = run.run(sql, options, sink).await;
+            // The driver records a keyspace once the server confirms it.
+            if outcome.is_ok() && follows {
+                *keyspace.lock().unwrap_or_else(PoisonError::into_inner) =
+                    link.session().get_keyspace().map(|keyspace| keyspace.to_string());
+            }
             let next = match &outcome {
                 Ok(_) => READY,
                 Err(_) => LOST,
@@ -375,3 +416,10 @@ fn lost(error: &DriverError) -> bool {
         ErrorCategory::Network | ErrorCategory::ConnectionLost | ErrorCategory::Tls
     )
 }
+
+async fn until(deadline: Option<Instant>) {
+    match deadline {
+        Some(at) => tokio::time::sleep_until(at).await,
+        None => std::future::pending().await,
+    }
+}

linux/crates/drivers/cassandra/src/errors.rs

@@ -88,7 +88,7 @@ fn pool_error(error: ConnectionPoolError, during: LossPhase) -> DriverError {
     }
 }
 
-fn use_keyspace_error(error: UseKeyspaceError, during: LossPhase) -> DriverError {
+pub(crate) fn use_keyspace_error(error: UseKeyspaceError, during: LossPhase) -> DriverError {
     match error {
         UseKeyspaceError::RequestError(attempt) => attempt_error(attempt, during),
         other => DriverError::server(other.to_string()),

linux/crates/drivers/clickhouse/src/connection.rs

@@ -261,6 +261,12 @@ impl Connection for ClickHouseConnection {
         outcome
     }
 
+    /// The schemas this engine lists are the server's databases, which a
+    /// query tab can move between.
+    async fn list_databases(&self, options: CallOptions) -> Result<Vec<SchemaInfo>, DriverError> {
+        self.list_schemas(options).await
+    }
+
     async fn list_schemas(&self, options: CallOptions) -> Result<Vec<SchemaInfo>, DriverError> {
         self.workspace_read(options, |calls| async move { metadata::list_schemas(&calls).await })
             .await

linux/crates/drivers/clickhouse/src/editor_session.rs

@@ -1,10 +1,10 @@
-use std::sync::Arc;
+use std::sync::{Arc, Mutex, PoisonError};
 use std::time::Duration;
 
 use async_trait::async_trait;
 use tablepro_core::{
     CallOptions, DriverError, EditorSession, EditorSessionId, EditorSessionOptions, EditorSessionState, LossPhase,
-    ResultSink, RunOptions, RunSummary, TimeoutPhase, TransactionState, UserSql,
+    ResultSink, RunEnd, RunOptions, RunSummary, TimeoutPhase, TransactionState, UserSql,
 };
 use tokio_util::sync::CancellationToken;
 use uuid::Uuid;
@@ -24,7 +24,9 @@ use crate::session_status::SessionStatus;
 /// takes no scripts.
 pub(crate) struct ClickHouseEditorSession {
     id: EditorSessionId,
-    session: clickhouse::Client,
+    /// Every request names its database, which overrides a `USE` the
+    /// session ran, so the tab's database lives in the client.
+    session: Arc<Mutex<clickhouse::Client>>,
     control: clickhouse::Client,
     statement_timeout: Option<Duration>,
     grace: Duration,
@@ -34,6 +36,10 @@ pub(crate) struct ClickHouseEditorSession {
 }
 
 impl ClickHouseEditorSession {
+    fn client(&self) -> clickhouse::Client {
+        self.session.lock().unwrap_or_else(PoisonError::into_inner).clone()
+    }
+
     /// Open the session with one trivial question, so a session the
     /// server will not create fails here rather than on the first run.
     pub(crate) async fn open(
@@ -59,7 +65,7 @@ impl ClickHouseEditorSession {
             id: EditorSessionId::new(),
             // Without the check, a run on a session the server let expire
             // quietly starts a new, empty one under the same id.
-            session: session.with_setting("session_check", "1"),
+            session: Arc::new(Mutex::new(session.with_setting("session_check", "1"))),
             control: endpoint.control().clone(),
             statement_timeout: options.statement_timeout,
             grace: liveness.cancel_grace,
@@ -76,10 +82,26 @@ impl EditorSession for ClickHouseEditorSession {
         self.id
     }
 
-    /// Every request names the database, which a `USE` in the session
-    /// does not override.
     fn database(&self) -> Option<String> {
-        self.session.database().map(str::to_owned)
+        self.client().database().map(str::to_owned)
+    }
+
+    /// The next request names `database`, once a request there answers.
+    async fn use_database(&self, database: &str, options: CallOptions) -> Result<(), DriverError> {
+        let moved = self.client().with_database(database);
+        let asking = drain(moved.query_raw("SELECT 1"));
+        let answered = tokio::select! {
+            biased;
+            () = options.cancel().cancelled() => return Err(DriverError::Cancelled),
+            () = until(options.deadline()) => return Err(DriverError::Timeout {
+                phase: TimeoutPhase::Statement,
+                server_cancelled: false,
+            }),
+            answered = asking => answered,
+        };
+        answered?;
+        *self.session.lock().unwrap_or_else(PoisonError::into_inner) = moved;
+        Ok(())
     }
 
     fn state(&self) -> EditorSessionState {
@@ -97,17 +119,27 @@ impl EditorSession for ClickHouseEditorSession {
     ) -> Result<RunSummary, DriverError> {
         self.status.begin()?;
         let guard = options.cancel.clone().drop_guard();
+        let used =
+            tablepro_core::sql_syntax::used_database(sql.as_str(), tablepro_core::sql_syntax::SqlGrammar::ClickHouse);
         let run = EditorRun {
-            session: self.session.clone(),
+            session: self.client(),
             control: self.control.clone(),
             statement_timeout: self.statement_timeout,
             grace: self.grace,
             shutdown: self.shutdown.clone(),
             lingering: self.lingering.clone(),
         };
-        let status = self.status.clone();
+        let (status, client) = (self.status.clone(), self.session.clone());
         let running = tokio::spawn(async move {
             let outcome = run.run(sql, options, sink).await;
+            // The server took the `USE`; the requests after it name that
+            // database, as the session would have.
+            if let (Ok(summary), Some(database)) = (&outcome, used)
+                && summary.end == RunEnd::Completed
+            {
+                let mut client = client.lock().unwrap_or_else(PoisonError::into_inner);
+                *client = client.clone().with_database(database);
+            }
             match &outcome {
                 // The server still held the session for a stopped query:
                 // the session itself is fine.
@@ -143,6 +175,13 @@ impl EditorSession for ClickHouseEditorSession {
     }
 }
 
+async fn until(deadline: Option<tokio::time::Instant>) {
+    match deadline {
+        Some(at) => tokio::time::sleep_until(at).await,
+        None => std::future::pending().await,
+    }
+}
+
 #[cfg(test)]
 mod tests {
     use async_trait::async_trait;
@@ -171,6 +210,63 @@ mod tests {
             .await
     }
 
+    struct Collect(Arc<Mutex<Vec<Vec<tablepro_core::Value>>>>);
+
+    #[async_trait]
+    impl ResultSink for Collect {
+        async fn accept(&mut self, event: ResultEvent) -> SinkFlow {
+            if let ResultEvent::Rows(rows) = event {
+                self.0.lock().unwrap_or_else(PoisonError::into_inner).extend(rows);
+            }
+            SinkFlow::Continue
+        }
+    }
+
+    async fn current_database(session: &ClickHouseEditorSession) -> Vec<Vec<tablepro_core::Value>> {
+        let rows = Arc::new(Mutex::new(Vec::new()));
+        session
+            .run(
+                UserSql::new("SELECT currentDatabase()"),
+                RunOptions::new(CancellationToken::new()),
+                Box::new(Collect(rows.clone())),
+            )
+            .await
+            .expect("a run");
+        rows.lock().unwrap_or_else(PoisonError::into_inner).clone()
+    }
+
+    /// Every request names its database, so a `USE` lasts only because
+    /// the session's next requests name the database it chose.
+    #[tokio::test]
+    #[ignore = "requires docker"]
+    async fn a_session_runs_in_the_database_it_moved_to() {
+        let fixture = ClickHouseFixture::start().await.expect("a ClickHouse container");
+        let endpoint = ClickHouseEndpoint::for_fixture(&fixture, test_liveness());
+        let session =
+            ClickHouseEditorSession::open(&endpoint, EditorSessionOptions::default(), CancellationToken::new())
+                .await
+                .expect("an editor session");
+        let text = |name: &str| vec![vec![tablepro_core::Value::Text(name.to_owned())]];
+
+        session
+            .use_database("system", CallOptions::new(None, CancellationToken::new()))
+            .await
+            .expect("a move to system");
+        assert_eq!(current_database(&session).await, text("system"));
+        assert_eq!(session.database().as_deref(), Some("system"));
+
+        let used = run(&session, "USE default").await;
+        assert_eq!(used.map(|summary| summary.end), Ok(RunEnd::Completed));
+        assert_eq!(current_database(&session).await, text("default"));
+        assert_eq!(session.database().as_deref(), Some("default"));
+
+        let refused = session
+            .use_database("no_such_database", CallOptions::new(None, CancellationToken::new()))
+            .await;
+        assert!(refused.is_err(), "{refused:?}");
+        assert_eq!(session.database().as_deref(), Some("default"), "a refused move moved");
+    }
+
     #[tokio::test]
     #[ignore = "requires docker"]
     async fn a_session_the_server_let_expire_is_lost_not_replaced() {

linux/crates/drivers/mongodb/src/connection.rs

@@ -282,6 +282,12 @@ impl Connection for MongoConnection {
         outcome
     }
 
+    /// The schemas this engine lists are the server's databases, which a
+    /// query tab can move between.
+    async fn list_databases(&self, options: CallOptions) -> Result<Vec<SchemaInfo>, DriverError> {
+        self.list_schemas(options).await
+    }
+
     async fn list_schemas(&self, options: CallOptions) -> Result<Vec<SchemaInfo>, DriverError> {
         let default = self.database.clone();
         self.read(options, move |context, session| {