feat(linux-editor): move a PostgreSQL query tab to another database by opening its session there

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.

linux/crates/app/src/ui/editor/editor_session_slot.rs

@@ -88,9 +88,12 @@ impl EditorSessionSlot {
         close_timeout: Duration,
         cancel: &CancellationToken,
     ) -> Result<Arc<dyn EditorSession>, 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<Session>,
         editor: Arc<dyn EditorSession>,
-        deadline: Option<tokio::time::Instant>,
+        (statement_timeout, close_timeout): (Option<Duration>, Duration),
         cancel: &CancellationToken,
     ) -> Result<Arc<dyn EditorSession>, 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<Session>,
+        editor: Arc<dyn EditorSession>,
+        (wanted, here): (String, Option<String>),
+        (statement_timeout, close_timeout): (Option<Duration>, Duration),
+        cancel: &CancellationToken,
+    ) -> Result<Arc<dyn EditorSession>, 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<Session>) -> 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<FakeEditorSession> {
+        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();

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<std::time::Duration>,
+    /// 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<String>,
 }
 
 /// One editor tab's own connection to the server.

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

@@ -254,6 +254,11 @@ impl Connection for PgConnection {
         outcome
     }
 
+    async fn list_databases(&self, options: CallOptions) -> Result<Vec<SchemaInfo>, DriverError> {
+        self.read(options, read_op(|client| Box::pin(metadata::list_databases(client))))
+            .await
+    }
+
     async fn list_schemas(&self, options: CallOptions) -> Result<Vec<SchemaInfo>, 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<Arc<dyn EditorSession>, 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)
     }
 

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<String>,
     commands: mpsc::Sender<Command>,
     status: Arc<SessionStatus>,
     shutdown: CancellationToken,
@@ -39,6 +41,7 @@ impl PgEditorSession {
         cancel: CancellationToken,
     ) -> Result<Arc<Self>, 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<String> {
+        self.database.clone()
+    }
+
     async fn run(
         &self,
         sql: UserSql,

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};

linux/crates/drivers/postgres/src/metadata/objects.rs

@@ -50,6 +50,32 @@ pub(crate) async fn list_schemas(client: &Client) -> Result<Vec<SchemaInfo>, 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<Vec<SchemaInfo>, 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<Vec<ObjectInfo>, DriverError> {
     let (mode, schema) = match scope {
         ObjectScope::DefaultSchema => ("default", ""),

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
     }

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;
+}