diff --git a/linux/changelog.d/3899.added.md b/linux/changelog.d/3899.added.md new file mode 100644 index 0000000000..2a9eb85047 --- /dev/null +++ b/linux/changelog.d/3899.added.md @@ -0,0 +1 @@ +On PostgreSQL, the SQL editor's database menu moves a tab to another of the server's databases by opening its connection there. diff --git a/linux/crates/app/src/ui/editor/editor_session_slot.rs b/linux/crates/app/src/ui/editor/editor_session_slot.rs index e0be235a77..941cc6ef90 100644 --- a/linux/crates/app/src/ui/editor/editor_session_slot.rs +++ b/linux/crates/app/src/ui/editor/editor_session_slot.rs @@ -88,9 +88,12 @@ impl EditorSessionSlot { close_timeout: Duration, cancel: &CancellationToken, ) -> Result, AcquireError> { - let deadline = statement_timeout.map(|timeout| tokio::time::Instant::now() + timeout); let stale = match self.lookup(session) { - Lookup::Ready(editor) => return self.reach_database(session, editor, deadline, cancel).await, + Lookup::Ready(editor) => { + return self + .reach_database(session, editor, (statement_timeout, close_timeout), cancel) + .await; + } Lookup::Closed => return Err(AcquireError::Cancelled), Lookup::RolledBack { stale } => { if let Some(stale) = stale { @@ -103,10 +106,16 @@ impl EditorSessionSlot { if let Some(stale) = stale { close_quietly(stale, close_timeout).await; } - let options = EditorSessionOptions { statement_timeout }; + let options = EditorSessionOptions { + statement_timeout, + database: self.database(), + }; match session.open_editor_session(options, cancel).await { Ok(editor) => match self.adopt(session, editor.clone()) { - Ok(()) => self.reach_database(session, editor, deadline, cancel).await, + Ok(()) => { + self.reach_database(session, editor, (statement_timeout, close_timeout), cancel) + .await + } Err(orphan) => { close_quietly(orphan, close_timeout).await; Err(AcquireError::Cancelled) @@ -145,9 +154,10 @@ impl EditorSessionSlot { &self, session: &Arc, editor: Arc, - deadline: Option, + (statement_timeout, close_timeout): (Option, Duration), cancel: &CancellationToken, ) -> Result, AcquireError> { + let deadline = statement_timeout.map(|timeout| tokio::time::Instant::now() + timeout); let Some(wanted) = self.database() else { return Ok(editor); }; @@ -161,6 +171,17 @@ impl EditorSessionSlot { { Ok(()) => Ok(editor), Err(tablepro_core::DriverError::Cancelled) => Err(AcquireError::Cancelled), + // The engine binds a session's database when it connects. + Err(tablepro_core::DriverError::Unsupported { .. }) => { + self.reopen_in( + session, + editor, + (wanted, here), + (statement_timeout, close_timeout), + cancel, + ) + .await + } Err(error) => { lock(&self.state).database = here; Err(AcquireError::Moved { @@ -171,6 +192,58 @@ impl EditorSessionSlot { } } + /// Give up `editor` for a session opened in `wanted`. One holding a + /// transaction stays: closing it would roll the transaction back. + async fn reopen_in( + &self, + session: &Arc, + editor: Arc, + (wanted, here): (String, Option), + (statement_timeout, close_timeout): (Option, Duration), + cancel: &CancellationToken, + ) -> Result, AcquireError> { + let holds_transaction = lock(&self.state) + .open + .as_ref() + .is_some_and(|open| open.editor.id() == editor.id() && open.held_transaction()); + if holds_transaction { + lock(&self.state).database = here; + return Err(AcquireError::Moved { + database: wanted, + reason: crate::i18n::gettext("Commit or roll back the transaction first."), + }); + } + // The session in the old database stays until one opens in the + // new, so a move that fails keeps what the tab built there. + let options = EditorSessionOptions { + statement_timeout, + database: Some(wanted.clone()), + }; + let opened = match session.open_editor_session(options, cancel).await { + Ok(opened) => opened, + Err(error) => { + lock(&self.state).database = here; + return Err(match error { + OperationError::Cancelled => AcquireError::Cancelled, + error => AcquireError::Moved { + database: wanted, + reason: crate::ui::error_text::operation_error_message(&error), + }, + }); + } + }; + match self.adopt(session, opened.clone()) { + Ok(()) => { + close_quietly(editor, close_timeout).await; + Ok(opened) + } + Err(orphan) => { + close_quietly(orphan, close_timeout).await; + Err(AcquireError::Cancelled) + } + } + } + fn lookup(&self, owner: &Arc) -> Lookup { let mut state = lock(&self.state); if state.closed { @@ -329,6 +402,83 @@ pub(super) mod tests { const CLOSE: Duration = Duration::from_secs(10); + /// A session in `database` that cannot move while it lives, as on + /// PostgreSQL. + fn bound_to(database: &str) -> Arc { + let editor = answering(); + editor.set_database(Some(database)); + editor.refuse_moves(tablepro_core::DriverError::Unsupported { + feature: tablepro_core::UnsupportedFeature::DatabaseSwitching, + }); + editor + } + + /// A session that cannot open in the new database leaves the tab in its + /// own, on the session it had and what it built there. + #[tokio::test] + async fn a_move_that_cannot_open_keeps_the_session_it_had() { + let shop = bound_to("shop"); + let (session, _) = session_opening(vec![shop.clone()]); + let slot = EditorSessionSlot::new(); + let first = slot + .acquire(&session, None, CLOSE, &CancellationToken::new()) + .await + .expect("a session in shop"); + + slot.set_database("analytics".to_owned()); + let refused = slot.acquire(&session, None, CLOSE, &CancellationToken::new()).await; + assert!( + matches!(refused, Err(AcquireError::Moved { .. })), + "{:?}", + refused.err() + ); + assert_eq!(shop.close_count(), 0, "the session in shop closed"); + assert_eq!(slot.database().as_deref(), Some("shop")); + + let again = slot + .acquire(&session, None, CLOSE, &CancellationToken::new()) + .await + .expect("the session in shop"); + assert_eq!(again.id(), first.id()); + } + + /// Where a live session cannot move, the tab gets a session opened in + /// its database, unless the one it has holds a transaction. + #[tokio::test] + async fn a_session_that_cannot_move_gives_way_unless_it_holds_a_transaction() { + let (shop, analytics) = (bound_to("shop"), bound_to("analytics")); + let (session, connection) = session_opening(vec![shop.clone(), analytics.clone()]); + let slot = EditorSessionSlot::new(); + let first = slot + .acquire(&session, None, CLOSE, &CancellationToken::new()) + .await + .expect("a session in shop"); + slot.note_transaction(&first, Some(TransactionState::InTransaction)); + + slot.set_database("analytics".to_owned()); + let refused = slot.acquire(&session, None, CLOSE, &CancellationToken::new()).await; + assert!( + matches!(refused, Err(AcquireError::Moved { .. })), + "{:?}", + refused.err() + ); + assert_eq!(slot.database().as_deref(), Some("shop"), "the tab left its transaction"); + assert_eq!(shop.close_count(), 0); + + slot.note_transaction(&first, Some(TransactionState::Idle)); + slot.set_database("analytics".to_owned()); + let moved = slot + .acquire(&session, None, CLOSE, &CancellationToken::new()) + .await + .expect("a session in analytics"); + assert_eq!(moved.database().as_deref(), Some("analytics")); + assert_eq!(shop.close_count(), 1, "the session in shop stayed open"); + assert_eq!( + tablepro_test_support::CallCounters::get(&connection.counters().open_editor_session), + 2 + ); + } + #[tokio::test] async fn a_tab_opens_its_editor_session_once_and_reuses_it() { let editor = answering(); diff --git a/linux/crates/core/src/editor_session.rs b/linux/crates/core/src/editor_session.rs index ea2d6d1662..670d05dabb 100644 --- a/linux/crates/core/src/editor_session.rs +++ b/linux/crates/core/src/editor_session.rs @@ -43,6 +43,10 @@ pub enum EditorSessionState { #[derive(Debug, Clone, Default)] pub struct EditorSessionOptions { pub statement_timeout: Option, + /// The database to open the session in, where the engine binds a + /// session's database when it connects. Elsewhere the session opens in + /// the connection's and `EditorSession::use_database` moves it. + pub database: Option, } /// One editor tab's own connection to the server. diff --git a/linux/crates/drivers/postgres/src/connection.rs b/linux/crates/drivers/postgres/src/connection.rs index 44045d4b30..015d920bf8 100644 --- a/linux/crates/drivers/postgres/src/connection.rs +++ b/linux/crates/drivers/postgres/src/connection.rs @@ -254,6 +254,11 @@ impl Connection for PgConnection { outcome } + async fn list_databases(&self, options: CallOptions) -> Result, DriverError> { + self.read(options, read_op(|client| Box::pin(metadata::list_databases(client)))) + .await + } + async fn list_schemas(&self, options: CallOptions) -> Result, DriverError> { self.read(options, read_op(|client| Box::pin(metadata::list_schemas(client)))) .await @@ -462,7 +467,13 @@ impl Connection for PgConnection { cancel: CancellationToken, ) -> Result, DriverError> { self.ensure_open()?; - let session = PgEditorSession::open(self.endpoint.clone(), self.settings.read_only, options, cancel).await?; + // A connection stays in the database it opened in, so a tab moved + // to another one gets a session opened there. + let endpoint = match options.database.as_deref() { + Some(database) => Arc::new(self.endpoint.in_database(database)), + None => self.endpoint.clone(), + }; + let session = PgEditorSession::open(endpoint, self.settings.read_only, options, cancel).await?; Ok(session) } diff --git a/linux/crates/drivers/postgres/src/editor/editor_session.rs b/linux/crates/drivers/postgres/src/editor/editor_session.rs index 44afb469f1..7df7283131 100644 --- a/linux/crates/drivers/postgres/src/editor/editor_session.rs +++ b/linux/crates/drivers/postgres/src/editor/editor_session.rs @@ -24,6 +24,8 @@ use crate::physical::PgEndpoint; /// are enforced with cancel requests. pub(crate) struct PgEditorSession { id: EditorSessionId, + /// A connection stays in the database it opened in. + database: Option, commands: mpsc::Sender, status: Arc, shutdown: CancellationToken, @@ -39,6 +41,7 @@ impl PgEditorSession { cancel: CancellationToken, ) -> Result, DriverError> { let (client, connection) = endpoint.open(&cancel).await?; + let database = endpoint.database().map(str::to_owned); let close_timeout = endpoint.liveness().close_timeout; let (notice_sender, notices) = mpsc::unbounded_channel(); // PostgreSQL has defaulted to on since 9.1; the server reports @@ -68,6 +71,7 @@ impl PgEditorSession { )); Ok(Arc::new(Self { id: EditorSessionId::new(), + database, commands, status, shutdown, @@ -87,6 +91,10 @@ impl EditorSession for PgEditorSession { self.status.state() } + fn database(&self) -> Option { + self.database.clone() + } + async fn run( &self, sql: UserSql, diff --git a/linux/crates/drivers/postgres/src/metadata/mod.rs b/linux/crates/drivers/postgres/src/metadata/mod.rs index 767638bd5c..b7ad20caac 100644 --- a/linux/crates/drivers/postgres/src/metadata/mod.rs +++ b/linux/crates/drivers/postgres/src/metadata/mod.rs @@ -12,7 +12,7 @@ pub(crate) use estimate::table_estimate; pub(crate) use foreign_keys::{fetch_foreign_keys, fetch_outbound_foreign_keys}; pub(crate) use identity::row_identity; pub(crate) use indexes::fetch_indexes; -pub(crate) use objects::{list_objects, list_schemas}; +pub(crate) use objects::{list_databases, list_objects, list_schemas}; use futures_util::TryStreamExt; use postgres_types::{ToSql, Type}; diff --git a/linux/crates/drivers/postgres/src/metadata/objects.rs b/linux/crates/drivers/postgres/src/metadata/objects.rs index 4282d7982f..bd08c91b1b 100644 --- a/linux/crates/drivers/postgres/src/metadata/objects.rs +++ b/linux/crates/drivers/postgres/src/metadata/objects.rs @@ -50,6 +50,32 @@ pub(crate) async fn list_schemas(client: &Client) -> Result, Dri .collect() } +/// The databases the login may connect to, the connection's own the +/// default. `template0` takes no connections; `template1` and `postgres` +/// are the server's own. +const DATABASES_SQL: &str = "\ +SELECT d.datname::text, + d.datname = current_database(), + d.datistemplate OR d.datname = 'postgres' +FROM pg_catalog.pg_database d +WHERE d.datallowconn AND pg_catalog.has_database_privilege(d.datname, 'CONNECT') +ORDER BY d.datname"; + +pub(crate) async fn list_databases(client: &Client) -> Result, DriverError> { + let field = decode_error("pg_database"); + text_rows(client, DATABASES_SQL, &[]) + .await? + .iter() + .map(|row| { + Ok(SchemaInfo { + name: row.try_get(0).map_err(&field)?, + is_default: row.try_get(1).map_err(&field)?, + is_system: row.try_get(2).map_err(&field)?, + }) + }) + .collect() +} + pub(crate) async fn list_objects(client: &Client, scope: &ObjectScope) -> Result, DriverError> { let (mode, schema) = match scope { ObjectScope::DefaultSchema => ("default", ""), diff --git a/linux/crates/drivers/postgres/src/physical.rs b/linux/crates/drivers/postgres/src/physical.rs index 0a03010a40..216443a778 100644 --- a/linux/crates/drivers/postgres/src/physical.rs +++ b/linux/crates/drivers/postgres/src/physical.rs @@ -56,6 +56,22 @@ impl PgEndpoint { }) } + /// The same server, login and TLS, opening connections in `database`. + pub(crate) fn in_database(&self, database: &str) -> Self { + let mut config = self.config.clone(); + config.dbname(database); + Self { + transport: self.transport.clone(), + config, + tls: self.tls.clone(), + liveness: self.liveness, + } + } + + pub(crate) fn database(&self) -> Option<&str> { + self.config.get_dbname() + } + pub(crate) fn liveness(&self) -> &LivenessPolicy { &self.liveness } diff --git a/linux/crates/drivers/postgres/tests/integration/editor_session.rs b/linux/crates/drivers/postgres/tests/integration/editor_session.rs index 058eca484d..63eb936270 100644 --- a/linux/crates/drivers/postgres/tests/integration/editor_session.rs +++ b/linux/crates/drivers/postgres/tests/integration/editor_session.rs @@ -526,3 +526,47 @@ async fn editor_deadline_times_out_with_the_server_acknowledging() { assert_eq!(session.state(), EditorSessionState::Ready); run_ok(session.as_ref(), "SELECT 1").await; } + +/// A tab moved to another database: the server lists it, and a session +/// opened there runs in it, since a connection never leaves the database +/// it opened in. +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +#[ignore = "requires docker"] +async fn a_session_opened_in_another_database_runs_there() { + let fixture = start().await; + let connection = connect(&fixture, false).await; + + let databases = connection.list_databases(within(CALL)).await.expect("the databases"); + let own = databases + .iter() + .find(|database| database.is_default) + .expect("the connection's database"); + assert!(!own.is_system, "{own:?}"); + let postgres = databases + .iter() + .find(|database| database.name == "postgres") + .expect("postgres"); + assert!(postgres.is_system, "{postgres:?}"); + assert!( + databases.iter().all(|database| database.name != "template0"), + "{databases:?}" + ); + + let options = tablepro_core::EditorSessionOptions { + statement_timeout: None, + database: Some("postgres".to_owned()), + }; + let elsewhere = connection + .open_editor_session(options, CancellationToken::new()) + .await + .expect("a session in postgres"); + let (_, events) = run(elsewhere.as_ref(), "SELECT current_database()").await; + assert_eq!(texts(&events), ["postgres"]); + assert_eq!(elsewhere.database().as_deref(), Some("postgres")); + let moved = elsewhere.use_database("tablepro", within(CALL)).await; + assert!( + matches!(moved, Err(tablepro_core::DriverError::Unsupported { .. })), + "a live session moved: {moved:?}" + ); + elsewhere.close().await; +} diff --git a/linux/crates/mcp/src/query_tools.rs b/linux/crates/mcp/src/query_tools.rs index 912f253774..bbdb656970 100644 --- a/linux/crates/mcp/src/query_tools.rs +++ b/linux/crates/mcp/src/query_tools.rs @@ -184,6 +184,7 @@ async fn run_read(session: &Session, text: &str, limit: usize, timeout: Duration async fn run_write(session: &Session, text: &str, limit: usize, timeout: Duration, cancel: &CancellationToken) -> Ran { let options = EditorSessionOptions { statement_timeout: Some(timeout), + database: None, }; let editor = session .open_editor_session(options, cancel) diff --git a/linux/crates/test-support/src/conformance/cases/document.rs b/linux/crates/test-support/src/conformance/cases/document.rs index eec2476dc2..a38fab090b 100644 --- a/linux/crates/test-support/src/conformance/cases/document.rs +++ b/linux/crates/test-support/src/conformance/cases/document.rs @@ -605,6 +605,7 @@ pub async fn a_timed_out_statement_is_stopped_by_the_server( let table = filled(TIMEOUT_CASE, &engine, connection.as_ref(), 50).await?; let options = EditorSessionOptions { statement_timeout: Some(SETTLE), + database: None, }; let session = connection .open_editor_session(options, CancellationToken::new()) diff --git a/linux/crates/test-support/src/conformance/cases/keyspace.rs b/linux/crates/test-support/src/conformance/cases/keyspace.rs index ea14571853..f0edded84d 100644 --- a/linux/crates/test-support/src/conformance/cases/keyspace.rs +++ b/linux/crates/test-support/src/conformance/cases/keyspace.rs @@ -340,6 +340,7 @@ pub async fn a_timed_out_command_is_stopped_by_the_server( let connection = connect(TIMEOUT_CASE, &engine, writable()).await?; let options = EditorSessionOptions { statement_timeout: Some(SETTLE), + database: None, }; let session = editor(connection.as_ref(), options).await?;