diff --git a/crates/ide-db/src/conn.rs b/crates/ide-db/src/conn.rs index af0825e06..0132af333 100644 --- a/crates/ide-db/src/conn.rs +++ b/crates/ide-db/src/conn.rs @@ -12,14 +12,17 @@ use ide_core::{CancellationToken, IdeError, IdeResult, error::ErrorKind}; use serde::{Deserialize, Serialize}; use specta::Type; use sqlx::{ - Column, ColumnIndex, Connection, Decode, Executor, MySql, Postgres, Row, SqlSafeStr, Sqlite, TypeInfo, ValueRef, - mysql::{MySqlConnectOptions, MySqlPool, MySqlPoolOptions, MySqlSslMode}, + Column, ColumnIndex, Connection, Decode, Executor, Row, SqlSafeStr, Sqlite, TypeInfo, ValueRef, + mysql::{MySqlConnectOptions, MySqlConnection, MySqlSslMode}, pool::PoolConnection, - postgres::{PgConnectOptions, PgPool, PgPoolOptions, PgSslMode}, + postgres::{PgConnectOptions, PgConnection, PgSslMode}, sqlite::{SqliteConnectOptions, SqlitePool, SqlitePoolOptions}, }; use std::path::{Path, PathBuf}; -use std::sync::Arc; +use std::sync::{ + Arc, + atomic::{AtomicBool, Ordering}, +}; use std::time::{Duration, Instant}; #[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize, Type)] @@ -47,12 +50,135 @@ pub struct QueryResult { pub execution_ms: f64, } +/// Session context belongs to the physical connection, including across pool checkouts. +#[derive(Debug, Clone)] +pub struct Context { + default: Option, + current: Option, + known: bool, +} + +impl Context { + pub(crate) fn new(default: Option) -> Self { + Self { current: default.clone(), default, known: true } + } + + pub(crate) fn set_current(&mut self, current: Option) { + self.current = current; + } + + fn wanted(&self, driver: DbDriver, schema: Option<&str>) -> Option { + schema + .filter(|s| !s.is_empty()) + .and_then(|s| crate::sql::use_schema(driver, s)) + .or_else(|| self.default.clone()) + } +} + +macro_rules! sqlx_pool { + ($manager:ident, $handle:ident, $options:ty, $connection:ty, $driver:expr, $query:expr) => { + pub struct $manager { + options: $options, + init: Arc>, + timeout: Duration, + } + + pub struct $handle { + conn: Option<$connection>, + context: Context, + options: $options, + init: Arc>, + timeout: Duration, + in_transaction: bool, + } + + impl std::ops::Deref for $handle { + type Target = $connection; + fn deref(&self) -> &Self::Target { + self.conn.as_ref().expect("live connection") + } + } + + impl std::ops::DerefMut for $handle { + fn deref_mut(&mut self) -> &mut Self::Target { + self.conn.as_mut().expect("live connection") + } + } + + impl ManageConnection for $manager { + type Connection = $handle; + type Error = sqlx::Error; + + async fn connect(&self) -> Result { + tokio::time::timeout(self.timeout, async { + let mut conn = <$connection>::connect_with(&self.options).await?; + let context_sql = |value: Option| { + value.map(|v| { + if $driver == DbDriver::Postgres { + format!("SET search_path TO {v}") + } else { + crate::sql::use_schema($driver, &v).expect("switchable driver") + } + }) + }; + for sql in self.init.iter() { + sqlx::raw_sql(sqlx::AssertSqlSafe(sql.clone())).execute(&mut conn).await?; + } + let row = sqlx::raw_sql(sqlx::AssertSqlSafe($query.to_owned())).fetch_one(&mut conn).await?; + let context = Context::new(context_sql(row.try_get(0)?)); + Ok($handle { + conn: Some(conn), + context, + options: self.options.clone(), + init: self.init.clone(), + timeout: self.timeout, + in_transaction: false, + }) + }) + .await + .map_err(|_| sqlx::Error::PoolTimedOut)? + } + + async fn is_valid(&self, conn: &mut Self::Connection) -> Result<(), Self::Error> { + conn.ping().await + } + + fn has_broken(&self, conn: &mut Self::Connection) -> bool { + conn.conn.is_none() + } + } + }; +} + +sqlx_pool!(PgManager, PgHandle, PgConnectOptions, PgConnection, DbDriver::Postgres, "SHOW search_path"); +sqlx_pool!(MysqlManager, MysqlHandle, MySqlConnectOptions, MySqlConnection, DbDriver::Mysql, "SELECT DATABASE()"); + +pub struct ServerPool { + inner: Bb8Pool, + closed: AtomicBool, +} + +impl ServerPool { + async fn get(&self) -> IdeResult> + where + M::Error: std::fmt::Display, + { + if self.closed.load(Ordering::Relaxed) { + return Err(IdeError::new(ErrorKind::Io, "database connection is closed")); + } + self.inner.get_owned().await.map_err(|e| match e { + bb8::RunError::User(e) => IdeError::new(ErrorKind::Io, e.to_string()), + bb8::RunError::TimedOut => IdeError::new(ErrorKind::Io, "timed out connecting to the database"), + }) + } +} + // Keep an empty slot when a connection is discarded after cancellation or rollback failure. pub struct MssqlConnectionManager(tiberius::Config); type MssqlClient = tiberius::Client>; -pub struct MssqlConnection(Option); +pub struct MssqlConnection(Option, Context); impl std::ops::Deref for MssqlConnection { type Target = MssqlClient; @@ -76,7 +202,7 @@ impl ManageConnection for MssqlConnectionManager { let tcp = tokio::net::TcpStream::connect(self.0.get_addr()).await?; tcp.set_nodelay(true)?; - let client = match MssqlClient::connect(self.0.clone(), tcp.compat_write()).await { + let mut client = match MssqlClient::connect(self.0.clone(), tcp.compat_write()).await { // Preserve the single redirect supported by the previous pool adapter. Err(tiberius::error::Error::Routing { host, port }) => { let mut config = self.0.clone(); @@ -88,7 +214,11 @@ impl ManageConnection for MssqlConnectionManager { } result => result?, }; - Ok(MssqlConnection(Some(client))) + let row = client.simple_query("SELECT DB_NAME()").await?.into_row().await?; + let default = row + .and_then(|r| r.get::<&str, _>(0).map(str::to_owned)) + .and_then(|db| crate::sql::use_schema(DbDriver::SqlServer, &db)); + Ok(MssqlConnection(Some(client), Context::new(default))) } async fn is_valid(&self, conn: &mut Self::Connection) -> Result<(), Self::Error> { @@ -103,25 +233,13 @@ impl ManageConnection for MssqlConnectionManager { pub enum Pool { Sqlite(SqlitePool), - Postgres(PgPool), - Mysql(MySqlPool), + Postgres(ServerPool), + Mysql(ServerPool), SqlServer(Bb8Pool), Clickhouse(crate::clickhouse::ClickhousePool), Oracle(crate::oracle::OraclePool), } -/// Runs `$body` with `$p` bound to the concrete pool of each driver. -macro_rules! with_pool { - ($pool:expr, $p:ident => $body:expr) => { - match $pool { - Pool::Sqlite($p) => $body, - Pool::Postgres($p) => $body, - Pool::Mysql($p) => $body, - Pool::SqlServer(_) | Pool::Clickhouse(_) | Pool::Oracle(_) => unreachable!(), - } - }; -} - pub(crate) fn db_err(e: sqlx::Error) -> IdeError { match &e { sqlx::Error::Database(d) => { @@ -438,6 +556,19 @@ macro_rules! backend_id { }}; } +// Session control, scripts and PostgreSQL set_config can invalidate the tracked context. +fn context_may_change(sql: &str) -> bool { + let text = sql.trim().trim_end_matches(';'); + text.contains(';') + || text.to_ascii_lowercase().contains("set_config") + || !text.split_whitespace().next().is_some_and(|s| { + matches!( + s.to_ascii_uppercase().as_str(), + "SELECT" | "INSERT" | "UPDATE" | "DELETE" | "CREATE" | "DROP" | "TRUNCATE" | "SHOW" | "EXPLAIN" + ) + }) +} + fn looks_like_query(sql: &str) -> bool { let first = sql .lines() @@ -535,7 +666,13 @@ impl Setup { return Err(IdeError::invalid("ClickHouse HTTP does not support statement caching")); } (_, "statement_cache_capacity") => s.statement_cache = Some(number(&name)? as usize), - (_, "connect_timeout") => s.connect_timeout = Some(Duration::from_secs(number(&name)?)), + (_, "connect_timeout") => { + let seconds = number(&name)?; + if seconds == 0 { + return Err(IdeError::invalid("connect_timeout must be greater than zero")); + } + s.connect_timeout = Some(Duration::from_secs(seconds)); + } (DbDriver::Postgres, "application_name") => s.application_name = Some(value.to_owned()), (DbDriver::Mysql, "charset") => s.charset = Some(value.to_owned()), (DbDriver::Mysql, "collation") => s.collation = Some(value.to_owned()), @@ -662,14 +799,12 @@ impl Pool { if let Some(n) = setup.statement_cache { opts = opts.statement_cache_capacity(n); } - let pool = PgPoolOptions::new() - .max_connections(MAX) - .acquire_timeout(timeout) - .after_connect(after_connect!()) - .connect_with(opts) - .await - .map_err(db_err)?; - Ok(Pool::Postgres(pool)) + let manager = PgManager { options: opts, init: init.clone(), timeout }; + // Open once directly so authentication failures retain their database error. + let first = manager.connect().await.map_err(db_err)?; + let pool = Bb8Pool::builder().max_size(MAX).connection_timeout(timeout).build_unchecked(manager); + let _ = pool.add(first); + Ok(Pool::Postgres(ServerPool { inner: pool, closed: AtomicBool::new(false) })) } DbDriver::Mysql => { let mut opts = MySqlConnectOptions::new() @@ -711,14 +846,12 @@ impl Pool { if !tz.is_empty() { opts = opts.timezone(Some(tz.to_owned())); } - let pool = MySqlPoolOptions::new() - .max_connections(MAX) - .acquire_timeout(timeout) - .after_connect(after_connect!()) - .connect_with(opts) - .await - .map_err(db_err)?; - Ok(Pool::Mysql(pool)) + let manager = MysqlManager { options: opts, init: init.clone(), timeout }; + // Open once directly so authentication failures retain their database error. + let first = manager.connect().await.map_err(db_err)?; + let pool = Bb8Pool::builder().max_size(MAX).connection_timeout(timeout).build_unchecked(manager); + let _ = pool.add(first); + Ok(Pool::Mysql(ServerPool { inner: pool, closed: AtomicBool::new(false) })) } DbDriver::SqlServer => { use tiberius::Config; @@ -757,7 +890,9 @@ impl Pool { match self { Pool::SqlServer(_) | Pool::Oracle(_) => {} Pool::Clickhouse(p) => p.close(), - _ => with_pool!(self, p => p.close().await), + Pool::Sqlite(p) => p.close().await, + Pool::Postgres(p) => p.closed.store(true, Ordering::Relaxed), + Pool::Mysql(p) => p.closed.store(true, Ordering::Relaxed), } } @@ -766,7 +901,9 @@ impl Pool { Pool::SqlServer(_) => false, Pool::Oracle(p) => p.is_closed(), Pool::Clickhouse(p) => p.is_closed(), - _ => with_pool!(self, p => p.is_closed()), + Pool::Sqlite(p) => p.is_closed(), + Pool::Postgres(p) => p.closed.load(Ordering::Relaxed), + Pool::Mysql(p) => p.closed.load(Ordering::Relaxed), } } @@ -775,18 +912,9 @@ impl Pool { if let Pool::Clickhouse(p) = self { return Ok(p.execute(sql, max, &CancellationToken::new()).await?.rows); } - if let Pool::Oracle(p) = self { - let mut conn = p.get().await?; - let fetched = crate::oracle::run(&mut conn, sql, 0, max, Cells::Text, &CancellationToken::new()).await?; - return Ok(fetched.1); - } - let fetched: IdeResult = match self { - Pool::SqlServer(p) => { - let mut conn = p.get_owned().await.map_err(mssql_pool_err)?; - mssql_fetch(&mut conn, sql, 0, max).await - } - _ => with_pool!(self, p => fetch!(p, sql, max, Cells::Text)), - }; + let mut conn = Held::acquire(self).await?; + conn.apply_context(None).await?; + let fetched = conn.fetch_rows(sql, max).await; Ok(fetched?.1) } @@ -809,7 +937,10 @@ impl Pool { return p.execute_from(sql, skip, max_rows, cancel).await; } let mut conn = Held::acquire(self).await?; - let r = conn.execute_from(self, sql, skip, max_rows, cancel).await; + let r = match conn.apply_context(None).await { + Ok(()) => conn.execute_from(self, sql, skip, max_rows, cancel).await, + Err(e) => Err(e), + }; if cancel.is_cancelled() { conn.discard().await; return Err(IdeError::cancelled()); @@ -822,8 +953,10 @@ impl Pool { pub async fn transaction(&self, statements: &[(String, bool)]) -> IdeResult { match self { Pool::Clickhouse(_) => Err(crate::clickhouse::transaction_error()), - Pool::SqlServer(p) => { - let mut conn = p.get_owned().await.map_err(mssql_pool_err)?; + Pool::SqlServer(_) => { + let mut held = Held::acquire(self).await?; + held.apply_context(None).await?; + let Held::SqlServer(mut conn) = held else { unreachable!() }; let result = async { mssql_command(&mut conn, "BEGIN TRAN").await?; let mut total = 0; @@ -846,8 +979,10 @@ impl Pool { } result } - Pool::Oracle(p) => { - let mut conn = p.get().await?; + Pool::Oracle(_) => { + let mut held = Held::acquire(self).await?; + held.apply_context(None).await?; + let Held::Oracle(mut conn) = held else { unreachable!() }; let result = crate::oracle::transaction(&mut conn, statements).await; if result.is_err() { // An open transaction must never return to the pool with the connection. @@ -855,21 +990,37 @@ impl Pool { } result } - _ => with_pool!(self, p => { - let mut tx = p.begin().await.map_err(db_err)?; - let mut total = 0; - for (sql, expect_one) in statements { - let done = sqlx::raw_sql(sqlx::AssertSqlSafe(sql.clone())).execute(&mut *tx).await.map_err(db_err)?; - let n = done.rows_affected(); - // MySQL reports 0 for an UPDATE that leaves the row unchanged. - if *expect_one && (n > 1 || (n == 0 && !matches!(self, Pool::Mysql(_)))) { - return Err(IdeError::invalid(format!("{sql}: expected to change 1 row, changed {n}; nothing was saved"))); - } - total += n; + _ => { + let mut conn = Held::acquire(self).await?; + conn.apply_context(None).await?; + macro_rules! transaction { + ($c:expr) => {{ + let mut tx = (&mut **$c).begin().await.map_err(db_err)?; + let mut total = 0; + for (sql, expect_one) in statements { + let done = sqlx::raw_sql(sqlx::AssertSqlSafe(sql.clone())) + .execute(&mut *tx) + .await + .map_err(db_err)?; + let n = done.rows_affected(); + if *expect_one && (n > 1 || (n == 0 && !matches!(self, Pool::Mysql(_)))) { + return Err(IdeError::invalid(format!( + "{sql}: expected to change 1 row, changed {n}; nothing was saved" + ))); + } + total += n; + } + tx.commit().await.map_err(db_err)?; + Ok(total) + }}; } - tx.commit().await.map_err(db_err)?; - Ok(total) - }), + match &mut conn { + Held::Sqlite(c) => transaction!(c), + Held::Postgres(c) => transaction!(&mut ***c), + Held::Mysql(c) => transaction!(&mut ***c), + _ => unreachable!(), + } + } } } } @@ -878,8 +1029,8 @@ impl Pool { /// manual-commit console across statements with its transaction open. pub enum Held { Sqlite(PoolConnection), - Postgres(PoolConnection), - Mysql(PoolConnection), + Postgres(Box>), + Mysql(Box>), SqlServer(Box>), Clickhouse(crate::clickhouse::ClickhousePool), Oracle(crate::oracle::OracleConnection), @@ -889,8 +1040,8 @@ impl Held { pub async fn acquire(pool: &Pool) -> IdeResult { Ok(match pool { Pool::Sqlite(p) => Held::Sqlite(p.acquire().await.map_err(db_err)?), - Pool::Postgres(p) => Held::Postgres(p.acquire().await.map_err(db_err)?), - Pool::Mysql(p) => Held::Mysql(p.acquire().await.map_err(db_err)?), + Pool::Postgres(p) => Held::Postgres(Box::new(p.get().await?)), + Pool::Mysql(p) => Held::Mysql(Box::new(p.get().await?)), Pool::SqlServer(p) => Held::SqlServer(Box::new(p.get_owned().await.map_err(mssql_pool_err)?)), Pool::Clickhouse(p) => { if p.is_closed() { @@ -902,8 +1053,121 @@ impl Held { }) } + fn context(&mut self) -> Option<&mut Context> { + match self { + Held::Postgres(c) => Some(&mut c.context), + Held::Mysql(c) => Some(&mut c.context), + Held::SqlServer(c) => Some(&mut c.1), + Held::Oracle(c) => Some(c.context()), + _ => None, + } + } + + fn invalidate_context(&mut self) { + if let Some(context) = self.context() { + context.known = false; + } + } + + fn before_statement(&mut self, sql: &str) { + let control = sql.trim().trim_end_matches(';').trim().to_ascii_uppercase(); + if let Held::Mysql(c) = self { + if matches!(control.as_str(), "BEGIN" | "START TRANSACTION") { + c.in_transaction = true; + return; + } + if matches!(control.as_str(), "COMMIT" | "ROLLBACK") { + return; + } + } + if context_may_change(sql) { + self.invalidate_context(); + } + } + + fn after_statement(&mut self, sql: &str) { + if let Held::Mysql(c) = self + && matches!(sql.trim().trim_end_matches(';').trim().to_ascii_uppercase().as_str(), "COMMIT" | "ROLLBACK") + { + c.in_transaction = false; + } + } + + pub async fn apply_context(&mut self, schema: Option<&str>) -> IdeResult<()> { + let driver = match self { + Held::Postgres(_) => DbDriver::Postgres, + Held::Mysql(_) => DbDriver::Mysql, + Held::SqlServer(_) => DbDriver::SqlServer, + Held::Oracle(_) => DbDriver::Oracle, + _ => return Ok(()), + }; + let context = self.context().expect("switchable connection"); + let wanted = context.wanted(driver, schema); + if context.known && context.current == wanted { + return Ok(()); + } + context.known = false; + if let Some(sql) = &wanted { + self.raw_command(sql).await?; + } else { + // MySQL cannot USE an absent database. Re-open with the original options. + macro_rules! reopen { + ($c:expr, $manager:ident) => {{ + let manager = $manager { options: $c.options.clone(), init: $c.init.clone(), timeout: $c.timeout }; + if $c.in_transaction { + return Err(IdeError::invalid("commit or roll back before clearing the selected database")); + } + let fresh = manager.connect().await.map_err(db_err)?; + if let Some(raw) = $c.conn.take() { + let _ = raw.close().await; + } + ***$c = fresh; + }}; + } + match self { + Held::Mysql(c) => reopen!(c, MysqlManager), + Held::Postgres(c) => reopen!(c, PgManager), + _ => unreachable!("server context has a default"), + } + } + let context = self.context().expect("switchable connection"); + context.current = wanted; + context.known = true; + Ok(()) + } + + /// Restore before a manual session returns its physical connection to the pool. + pub async fn restore_default(&mut self) -> IdeResult<()> { + self.apply_context(None).await + } + + async fn fetch_rows(&mut self, sql: &str, max: usize) -> IdeResult { + self.before_statement(sql); + let result = match self { + Held::Sqlite(c) => fetch!(&mut **c, sql, max, Cells::Text), + Held::Postgres(c) => fetch!(&mut ****c, sql, max, Cells::Text), + Held::Mysql(c) => fetch!(&mut ****c, sql, max, Cells::Text), + Held::SqlServer(c) => mssql_fetch(c, sql, 0, max).await, + Held::Oracle(c) => crate::oracle::run(c, sql, 0, max, Cells::Text, &CancellationToken::new()).await, + Held::Clickhouse(_) => unreachable!(), + }; + if result.is_ok() { + self.after_statement(sql); + } + result + } + /// Runs a statement without results (`BEGIN`, `COMMIT`, `ROLLBACK`). pub async fn command(&mut self, sql: &str) -> IdeResult<()> { + self.before_statement(sql); + let result = self.raw_command(sql).await; + if result.is_ok() { + self.after_statement(sql); + } + result + } + + async fn raw_command(&mut self, sql: &str) -> IdeResult<()> { match self { Held::SqlServer(c) => mssql_command(c, sql).await, Held::Clickhouse(_) => Err(crate::clickhouse::transaction_error()), @@ -918,8 +1182,8 @@ impl Held { let sql = sqlx::AssertSqlSafe(sql.to_owned()); match other { Held::Sqlite(c) => sqlx::raw_sql(sql).execute(&mut **c).await.map(drop), - Held::Postgres(c) => sqlx::raw_sql(sql).execute(&mut **c).await.map(drop), - Held::Mysql(c) => sqlx::raw_sql(sql).execute(&mut **c).await.map(drop), + Held::Postgres(c) => sqlx::raw_sql(sql).execute(&mut ****c).await.map(drop), + Held::Mysql(c) => sqlx::raw_sql(sql).execute(&mut ****c).await.map(drop), Held::SqlServer(_) | Held::Clickhouse(_) | Held::Oracle(_) => unreachable!(), } .map_err(db_err) @@ -932,11 +1196,15 @@ impl Held { pub async fn discard(self) { match self { Held::Sqlite(c) => drop(c.detach()), - Held::Postgres(c) => { - let _ = c.detach().close().await; + Held::Postgres(mut c) => { + if let Some(raw) = c.conn.take() { + let _ = raw.close().await; + } } - Held::Mysql(c) => { - let _ = c.detach().close().await; + Held::Mysql(mut c) => { + if let Some(raw) = c.conn.take() { + let _ = raw.close().await; + } } Held::SqlServer(mut c) => { c.0.take(); @@ -971,6 +1239,7 @@ impl Held { if let Held::Clickhouse(c) = self { return c.execute_from(sql, skip, max_rows, cancel).await; } + self.before_statement(sql); let started = Instant::now(); // `$id`: the query naming the server process, `$kill`: how another connection stops it. macro_rules! run { @@ -1003,9 +1272,9 @@ impl Held { let fetched: Fetched = match self { Held::Sqlite(c) => run!(c, None::<&str>, |_: String| String::new()), Held::Postgres(c) => { - run!(c, Some("SELECT pg_backend_pid()"), |b: String| format!("SELECT pg_cancel_backend({b})")) + run!(&mut ***c, Some("SELECT pg_backend_pid()"), |b: String| format!("SELECT pg_cancel_backend({b})")) } - Held::Mysql(c) => run!(c, Some("SELECT CONNECTION_ID()"), |b: String| format!("KILL QUERY {b}")), + Held::Mysql(c) => run!(&mut ***c, Some("SELECT CONNECTION_ID()"), |b: String| format!("KILL QUERY {b}")), Held::SqlServer(c) => { // Tiberius has no attention/cancel API; stop reading and discard the connection. tokio::select! { @@ -1016,6 +1285,7 @@ impl Held { Held::Oracle(c) => crate::oracle::run(c, sql, skip, max_rows, Cells::Display, cancel).await?, Held::Clickhouse(_) => unreachable!(), }; + self.after_statement(sql); let (columns, rows, affected, has_more, first) = fetched; let elapsed_ms = started.elapsed().as_secs_f64() * 1000.0; Ok(QueryResult { @@ -1055,7 +1325,7 @@ mod tests { #[test] fn discarded_sqlserver_connections_are_broken() { let manager = MssqlConnectionManager(tiberius::Config::new()); - assert!(manager.has_broken(&mut MssqlConnection(None))); + assert!(manager.has_broken(&mut MssqlConnection(None, Context::new(None)))); } #[tokio::test] @@ -1191,6 +1461,7 @@ mod tests { assert_eq!(setup.charset.as_deref(), Some("utf8mb4")); assert_eq!(setup.connect_timeout, Some(Duration::from_secs(5))); assert!(Setup::of(&DataSource { properties: vec![prop("a; drop", "1")], ..s.clone() }).is_err()); + assert!(Setup::of(&DataSource { properties: vec![prop("connect_timeout", "0")], ..s.clone() }).is_err()); assert!(Setup::of(&DataSource { properties: vec![prop("connect_timeout", "soon")], ..s }).is_err()); } diff --git a/crates/ide-db/src/lib.rs b/crates/ide-db/src/lib.rs index dabe98f99..79913f3dc 100644 --- a/crates/ide-db/src/lib.rs +++ b/crates/ide-db/src/lib.rs @@ -315,7 +315,7 @@ fn keep_alive(pool: &Arc, every: Duration) { async fn end_session(session: &Session) { let held = session.lock().await.take(); if let Some(mut conn) = held - && conn.command("ROLLBACK").await.is_err() + && (conn.command("ROLLBACK").await.is_err() || conn.restore_default().await.is_err()) { conn.discard().await; } @@ -704,8 +704,15 @@ impl DbManager { let introspection = self.source(id)?.introspection; let system = introspection.show_system; let pool = self.pool(id).await?; + if pool.driver() == DbDriver::SqlServer { + let first = + |rows: Vec>>| rows.into_iter().filter_map(|r| r.into_iter().next().flatten()); + let current = first(pool.rows("SELECT DB_NAME()", 1).await?).next().unwrap_or_default(); + let names = first(pool.rows(sql::sqlserver_databases(), MAX_META_ROWS).await?).collect(); + return Ok(Databases { current, names }); + } if pool.driver() != DbDriver::Postgres { - return Err(IdeError::invalid("only PostgreSQL data sources list databases")); + return Err(IdeError::invalid("this data source does not list databases")); } let first = |rows: Vec>>| rows.into_iter().filter_map(|r| r.into_iter().next().flatten()); let current = first(pool.rows("SELECT current_database()", 1).await?).next().unwrap_or_default(); @@ -1038,11 +1045,13 @@ impl DbManager { max_rows: usize, cancel: &CancellationToken, ) -> IdeResult { - self.execute_from(id, statement, 0, max_rows, cancel).await + self.execute_from(id, statement, 0, max_rows, cancel, None).await } /// [`execute`](Self::execute), reading past the first `skip` rows without keeping /// them: a later page of a result, where only that page is decoded and returned. + /// `schema` is a query console's active context: the connection that runs the + /// statement is put in it first, so unqualified names resolve there. pub async fn execute_from( &self, id: &str, @@ -1050,8 +1059,9 @@ impl DbManager { skip: usize, max_rows: usize, cancel: &CancellationToken, + schema: Option<&str>, ) -> IdeResult { - let r = self.run(id, statement, skip, max_rows, cancel).await; + let r = self.run(id, statement, skip, max_rows, cancel, schema).await; let ran = match &r { Ok(_) => true, Err(e) => !matches!(e.kind, ErrorKind::Cancelled | ErrorKind::PermissionDenied), @@ -1083,8 +1093,15 @@ impl DbManager { .map_err(|e| IdeError::internal(e.to_string()))? } - /// The execution plan of `statement` (explained, not run). - pub async fn explain(&self, id: &str, statement: &str, cancel: &CancellationToken) -> IdeResult> { + /// The execution plan of `statement` (explained, not run), in the console's + /// active `schema` when it has one. + pub async fn explain( + &self, + id: &str, + statement: &str, + cancel: &CancellationToken, + schema: Option<&str>, + ) -> IdeResult> { let driver = self.pool(id).await?.driver(); if driver == DbDriver::Oracle { // EXPLAIN PLAN writes PLAN_TABLE; the plan is read (and cleaned up) separately. @@ -1102,14 +1119,16 @@ impl DbManager { 0, MAX_META_ROWS, cancel, + schema, ) .await?; - let result = self.run(id, &plan::oracle_plan_query(&plan_id), 0, MAX_META_ROWS, cancel).await?; + let result = self.run(id, &plan::oracle_plan_query(&plan_id), 0, MAX_META_ROWS, cancel, None).await?; let nodes = plan::parse(driver, &result); - let _ = self.run(id, &plan::oracle_plan_cleanup(&plan_id), 0, 0, cancel).await; + let _ = self.run(id, &plan::oracle_plan_cleanup(&plan_id), 0, 0, cancel, None).await; return nodes; } - let result = self.run(id, &plan::explain_statement(driver, statement), 0, MAX_META_ROWS, cancel).await?; + let result = + self.run(id, &plan::explain_statement(driver, statement), 0, MAX_META_ROWS, cancel, schema).await?; plan::parse(driver, &result) } @@ -1120,11 +1139,25 @@ impl DbManager { skip: usize, max_rows: usize, cancel: &CancellationToken, + schema: Option<&str>, ) -> IdeResult { let pool = self.pool(id).await?; let session = self.sessions.lock().get(id).cloned(); let Some(session) = session else { - return pool.execute_from(statement, skip, max_rows, cancel).await; + if schema.is_none() { + return pool.execute_from(statement, skip, max_rows, cancel).await; + } + // Like `Pool::execute_from`, with the schema applied to the connection first. + let mut conn = Held::acquire(&pool).await?; + let r = match conn.apply_context(schema).await { + Ok(()) => conn.execute_from(&pool, statement, skip, max_rows, cancel).await, + Err(e) => Err(e), + }; + if cancel.is_cancelled() { + conn.discard().await; + return Err(IdeError::cancelled()); + } + return r; }; if pool.driver() == DbDriver::Clickhouse { return Err(IdeError::invalid("ClickHouse manual transactions are unsupported")); @@ -1132,10 +1165,12 @@ impl DbManager { let mut held = session.lock().await; if held.is_none() { let mut conn = Held::acquire(&pool).await?; + conn.apply_context(schema).await?; conn.command(conn::begin_statement(pool.driver())).await?; *held = Some(conn); } let conn = held.as_mut().expect("opened above"); + conn.apply_context(schema).await?; let r = conn.execute_from(&pool, statement, skip, max_rows, cancel).await; if cancel.is_cancelled() { // The server rolls the transaction back with the connection. @@ -1171,7 +1206,12 @@ impl DbManager { let Some(session) = session else { return Ok(()) }; let mut held = session.lock().await; let Some(mut conn) = held.take() else { return Ok(()) }; - match conn.command(if commit { "COMMIT" } else { "ROLLBACK" }).await { + let result = async { + conn.command(if commit { "COMMIT" } else { "ROLLBACK" }).await?; + conn.restore_default().await + } + .await; + match result { Ok(()) => Ok(()), Err(e) => { conn.discard().await; diff --git a/crates/ide-db/src/oracle.rs b/crates/ide-db/src/oracle.rs index 5ece541a7..842010a55 100644 --- a/crates/ide-db/src/oracle.rs +++ b/crates/ide-db/src/oracle.rs @@ -181,7 +181,7 @@ impl OraclePool { .await .map_err(|e| IdeError::internal(format!("Oracle connect task: {e}")))? .map_err(oracle_err)?; - let _ = pool.add(OracleHandle { conn: Some(opened) }); + let _ = pool.add(opened); Ok(Self { state, pool: Arc::new(pool) }) } @@ -208,6 +208,7 @@ fn pool_err(e: bb8::RunError) -> OracleError { /// discarded (a cancelled statement leaves it in the blocking task). pub struct OracleHandle { conn: Option, + context: crate::conn::Context, } pub struct OracleConnectionManager(Arc); @@ -222,7 +223,7 @@ impl ManageConnection for OracleConnectionManager { .await .map_err(|e| OracleError::message(format!("Oracle connect task: {e}")))?; let conn = opened.map_err(OracleError::new)?; - Ok(OracleHandle { conn: Some(conn) }) + Ok(conn) } async fn is_valid(&self, conn: &mut Self::Connection) -> Result<(), Self::Error> { @@ -271,6 +272,10 @@ pub struct OracleConnection { } impl OracleConnection { + pub(crate) fn context(&mut self) -> &mut crate::conn::Context { + &mut self.pooled.context + } + pub fn is_closed(&self) -> bool { self.pooled.conn.is_none() } @@ -282,7 +287,7 @@ impl OracleConnection { } } -fn open(state: &State) -> oracle::Result { +fn open(state: &State) -> oracle::Result { let mut connector = oracle::Connector::new(state.user.clone(), state.password.clone(), state.connect.clone()); connector.driver_name("OpenIDEA"); if let Some(n) = state.statement_cache { @@ -293,7 +298,12 @@ fn open(state: &State) -> oracle::Result { for sql in &state.init { statement(&mut conn, sql, 0, 0, Cells::Text).map_err(|e| message_error(&e.message))?; } - Ok(conn) + let user: String = conn.query_row_as("SELECT SYS_CONTEXT('USERENV', 'SESSION_USER') FROM DUAL", &[])?; + let default = crate::sql::use_schema(DbDriver::Oracle, &user); + let current: String = conn.query_row_as("SELECT SYS_CONTEXT('USERENV', 'CURRENT_SCHEMA') FROM DUAL", &[])?; + let mut context = crate::conn::Context::new(default); + context.set_current(crate::sql::use_schema(DbDriver::Oracle, ¤t)); + Ok(OracleHandle { conn: Some(conn), context }) } /// Runs `f` with the connection on the blocking pool. Cancelling drops the @@ -558,7 +568,7 @@ mod tests { ) .unwrap(), )); - let mut broken = OracleHandle { conn: None }; + let mut broken = OracleHandle { conn: None, context: crate::conn::Context::new(None) }; assert!(manager.has_broken(&mut broken)); // A connection that cannot be opened reports the failure without panicking // (a client-load error, or ORA-12541 when a client is available). diff --git a/crates/ide-db/src/sql.rs b/crates/ide-db/src/sql.rs index 9e2172610..013833c2a 100644 --- a/crates/ide-db/src/sql.rs +++ b/crates/ide-db/src/sql.rs @@ -63,6 +63,11 @@ pub const ORACLE_SYSTEM_SCHEMAS: &[&str] = &[ "XS$NULL", ]; +/// Databases available to this login; offline/restoring databases cannot be selected. +pub fn sqlserver_databases() -> &'static str { + "SELECT name FROM sys.databases WHERE state = 0 AND HAS_DBACCESS(name) = 1 ORDER BY name" +} + pub fn schemas(driver: DbDriver, system: bool) -> &'static str { match driver { DbDriver::Sqlite => "SELECT name FROM pragma_database_list WHERE name <> 'temp' ORDER BY seq", @@ -396,6 +401,18 @@ pub fn current_schema(driver: DbDriver) -> Option<&'static str> { } } +/// The statement that makes unqualified names resolve in `schema` on the +/// connection that runs it (a query console's active schema; `None`: the driver +/// has no schema to switch — SQLite's `main`, ClickHouse's `USE` differs). +pub fn use_schema(driver: DbDriver, schema: &str) -> Option { + match driver { + DbDriver::Sqlite | DbDriver::Clickhouse => None, + DbDriver::Postgres => Some(format!("SET search_path TO {}", quote_ident(driver, schema))), + DbDriver::Mysql | DbDriver::SqlServer => Some(format!("USE {}", quote_ident(driver, schema))), + DbDriver::Oracle => Some(format!("ALTER SESSION SET CURRENT_SCHEMA = {}", quote_ident(driver, schema))), + } +} + #[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize, Type)] #[serde(rename_all = "camelCase")] #[serde(rename = "DbSortKey")] @@ -994,6 +1011,20 @@ mod tests { assert!(object_ddl(DbDriver::Postgres, "s", K::Sequence, "q").unwrap().0.contains("sequencename = 'q'")); } + #[test] + fn schema_switch_statement_per_dialect() { + assert_eq!(use_schema(DbDriver::Postgres, "app data"), Some("SET search_path TO \"app data\"".into())); + assert_eq!(use_schema(DbDriver::Mysql, "a`b"), Some("USE `a``b`".into())); + assert_eq!(use_schema(DbDriver::SqlServer, "app] data"), Some("USE [app]] data]".into())); + assert_eq!(use_schema(DbDriver::Postgres, "a\"b"), Some("SET search_path TO \"a\"\"b\"".into())); + assert_eq!(use_schema(DbDriver::Oracle, "a\"b"), Some("ALTER SESSION SET CURRENT_SCHEMA = \"a\"\"b\"".into())); + assert!(sqlserver_databases().contains("state = 0 AND HAS_DBACCESS(name) = 1")); + assert_eq!(use_schema(DbDriver::Oracle, "hr"), Some("ALTER SESSION SET CURRENT_SCHEMA = \"hr\"".into())); + // Single-schema (SQLite) and not-yet-supported drivers: no dropdown, no switch. + assert_eq!(use_schema(DbDriver::Sqlite, "main"), None); + assert_eq!(use_schema(DbDriver::Clickhouse, "default"), None); + } + #[test] fn postgres_ddl() { let cols = [ diff --git a/crates/ide-db/tests/clickhouse.rs b/crates/ide-db/tests/clickhouse.rs index d19a9ff16..33e57afa3 100644 --- a/crates/ide-db/tests/clickhouse.rs +++ b/crates/ide-db/tests/clickhouse.rs @@ -103,7 +103,7 @@ async fn clickhouse_cloud_manager_catalogs_and_session_table() -> IdeResult<()> "completion columns differ from catalog", )?; } - let plan = db.explain("clickhouse", &format!("SELECT id FROM {table}"), &cancel).await?; + let plan = db.explain("clickhouse", &format!("SELECT id FROM {table}"), &cancel, None).await?; check(!plan.is_empty(), "query plan missing")?; let value = r"path\O'Brien"; let result = db diff --git a/crates/ide-db/tests/mysql.rs b/crates/ide-db/tests/mysql.rs index a6ab2431a..25360eb51 100644 --- a/crates/ide-db/tests/mysql.rs +++ b/crates/ide-db/tests/mysql.rs @@ -85,3 +85,101 @@ async fn metadata_names_are_text_whatever_the_column_collation() -> IdeResult<() let raw = manager(&tmp, "SET character_set_results = binary").await.unwrap(); assert_metadata_is_text(&raw).await } + +// A console's active schema: the connection that runs each autocommit statement +// is put in it first, so unqualified names resolve there (and not elsewhere). +// Its own schema names: tests of this file share one server but run in parallel. +#[tokio::test] +async fn active_schema_context_drives_the_executing_connection() -> IdeResult<()> { + const A: &str = "oi_ctx_meta"; + const B: &str = "oi_ctx_data"; + let tmp = tempfile::tempdir().unwrap(); + let Some(db) = manager(&tmp, "").await else { return Ok(()) }; + let c = CancellationToken::new(); + for s in [A, B] { + for stmt in [ + format!("DROP DATABASE IF EXISTS `{s}`"), + format!("CREATE DATABASE `{s}` CHARACTER SET utf8mb4 COLLATE utf8mb4_bin"), + format!("CREATE TABLE `{s}`.`用户` (id INT PRIMARY KEY, name VARCHAR(20))"), + format!("INSERT INTO `{s}`.`用户` VALUES (1, '{s}')"), + ] { + db.execute("m", &stmt, 10, &c).await.unwrap(); + } + } + + // The source connects with no default database; the console's pick decides. + let r = db.execute_from("m", "SELECT DATABASE()", 0, 1, &c, Some(A)).await?; + assert_eq!(r.rows, vec![vec![Some(A.into())]]); + // Both schemas hold `用户`; the context resolves the unqualified name in its own, + // on every pooled connection the autocommit statements may land on. + for _ in 0..6 { + for (schema, name) in [(A, A), (B, B)] { + let r = db.execute_from("m", "SELECT name FROM 用户", 0, 1, &c, Some(schema)).await?; + assert_eq!(r.rows, vec![vec![Some(name.into())]], "of {schema}"); + let r = db.execute_from("m", "SELECT DATABASE()", 0, 1, &c, Some("")).await?; + assert_eq!(r.rows, vec![vec![None]]); + assert_eq!(db.completion("m").await?.default_schema, ""); + } + } + // A statement without a selection returns to the source's absent default. + let r = db.execute_from("m", "SELECT DATABASE()", 0, 1, &c, None).await?; + assert_eq!(r.rows, vec![vec![None]]); + + // A manual-commit console switches schemas on its session connection. + db.set_manual_commit("m", true).await?; + // Starting a no-database manual session must keep its transaction open. + db.execute("m", &format!("INSERT INTO `{A}`.`用户` VALUES (2, 'rollback')"), 1, &c).await?; + db.finish_transaction("m", false).await?; + assert!(db.execute("m", &format!("SELECT name FROM `{A}`.`用户` WHERE id = 2"), 1, &c).await?.rows.is_empty()); + let r = db.execute_from("m", "SELECT name FROM 用户", 0, 1, &c, Some(A)).await?; + assert_eq!(r.rows, vec![vec![Some(A.into())]]); + let r = db.execute_from("m", "SELECT DATABASE()", 0, 1, &c, Some(B)).await?; + assert_eq!(r.rows, vec![vec![Some(B.into())]]); + assert!(db.execute("m", "SELECT DATABASE()", 1, &c).await.is_err()); + // The rejected clear keeps the transaction and selection intact. + assert_eq!(db.execute_from("m", "SELECT DATABASE()", 0, 1, &c, Some(B)).await?.rows, vec![vec![Some(B.into())]]); + db.finish_transaction("m", true).await?; + for commit in [false, true] { + db.execute_from("m", "SELECT 1", 0, 1, &c, Some(A)).await?; + db.finish_transaction("m", commit).await?; + let r = db.execute("m", "SELECT DATABASE()", 1, &c).await?; + assert_eq!(r.rows, vec![vec![None]]); + } + db.execute_from("m", "SELECT 1", 0, 1, &c, Some(A)).await?; + db.set_manual_commit("m", false).await?; + db.execute("m", &format!("USE `{B}`"), 1, &c).await?; + let r = db.execute("m", "SELECT DATABASE()", 1, &c).await?; + assert_eq!(r.rows, vec![vec![None]]); + + // An unknown schema fails the statement instead of running it somewhere else. + let error = db.execute_from("m", "SELECT 1", 0, 1, &c, Some("oi_missing")).await.unwrap_err(); + assert!(error.message.contains("oi_missing"), "{}", error.message); + Ok(()) +} + +#[tokio::test] +async fn a_reused_physical_connection_restores_the_configured_database() -> IdeResult<()> { + use ide_db::{Held, Pool}; + let Some((host, port, user, password)) = server() else { return Ok(()) }; + let source = DataSource { + driver: DbDriver::Mysql, + host, + port: Some(port), + user, + database: "mysql".into(), + ..Default::default() + }; + let pool = Pool::connect(&source, Some(&password), None).await?; + let c = CancellationToken::new(); + let _other = [Held::acquire(&pool).await?, Held::acquire(&pool).await?, Held::acquire(&pool).await?]; + let mut conn = Held::acquire(&pool).await?; + let backend = conn.execute(&pool, "SELECT CONNECTION_ID()", 1, &c).await?.rows; + for schema in ["information_schema", "mysql", "information_schema"] { + conn.apply_context(Some(schema)).await?; + assert_eq!(conn.execute(&pool, "SELECT DATABASE()", 1, &c).await?.rows, vec![vec![Some(schema.into())]]); + } + drop(conn); + assert_eq!(pool.rows("SELECT DATABASE()", 1).await?, vec![vec![Some("mysql".into())]]); + assert_eq!(pool.execute("SELECT CONNECTION_ID()", 1, &c).await?.rows, backend); + Ok(()) +} diff --git a/crates/ide-db/tests/postgres.rs b/crates/ide-db/tests/postgres.rs new file mode 100644 index 000000000..d74515681 --- /dev/null +++ b/crates/ide-db/tests/postgres.rs @@ -0,0 +1,135 @@ +//! The active schema context against a real PostgreSQL server, skipped unless +//! `OPENIDEA_TEST_POSTGRES=host:port:user:password` names a throwaway server. + +use ide_core::{CancellationToken, IdeResult}; +use ide_db::{DataSource, DbDriver, DbManager}; + +fn server() -> Option<(String, u16, String, String)> { + let v = std::env::var("OPENIDEA_TEST_POSTGRES").ok()?; + let mut it = v.splitn(4, ':'); + Some((it.next()?.into(), it.next()?.parse().ok()?, it.next()?.into(), it.next().unwrap_or("").into())) +} + +async fn manager(tmp: &tempfile::TempDir) -> Option { + let (host, port, user, password) = server()?; + let root = tmp.path().to_path_buf(); + let ssh = std::sync::Arc::new(ide_ssh::SshManager::new(root.join("ssh-passwords.json"))); + let db = DbManager::new(root.clone(), root.join("passwords.json"), root.join("history.json"), ssh); + db.save_sources(vec![DataSource { + id: "p".into(), + name: "postgres".into(), + driver: DbDriver::Postgres, + host, + port: Some(port), + user, + ..Default::default() + }]) + .await + .unwrap(); + db.set_password("p", Some(password), false).await.unwrap(); + Some(db) +} + +async fn create_schemas(db: &DbManager) { + let c = CancellationToken::new(); + for (schema, tag) in [("oi_a", "from a"), ("oi_b", "from b")] { + for stmt in [ + format!("DROP SCHEMA IF EXISTS \"{schema}\" CASCADE"), + format!("CREATE SCHEMA \"{schema}\""), + format!("CREATE TABLE \"{schema}\".items (id INT PRIMARY KEY, tag TEXT)"), + format!("INSERT INTO \"{schema}\".items VALUES (1, '{tag}')"), + ] { + db.execute("p", &stmt, 10, &c).await.unwrap(); + } + } +} + +#[tokio::test] +async fn active_schema_context_drives_the_executing_connection() -> IdeResult<()> { + let tmp = tempfile::tempdir().unwrap(); + let Some(db) = manager(&tmp).await else { return Ok(()) }; + create_schemas(&db).await; + let c = CancellationToken::new(); + + // The connection default stays untouched before any console picks a schema. + let r = db.execute("p", "SHOW search_path", 1, &c).await?; + assert!(r.rows[0][0].as_deref().unwrap_or("").contains("public"), "{:?}", r.rows); + + // Both schemas hold `items`; the context resolves the unqualified name in its own, + // on every pooled connection the autocommit statements may land on. + for _ in 0..6 { + for (schema, tag) in [("oi_a", "from a"), ("oi_b", "from b")] { + let r = db.execute_from("p", "SELECT tag FROM items", 0, 1, &c, Some(schema)).await?; + assert_eq!(r.rows, vec![vec![Some(tag.into())]], "of {schema}"); + let r = db.execute_from("p", "SHOW search_path", 0, 1, &c, Some("")).await?; + assert_eq!(r.rows, vec![vec![Some("\"$user\", public".into())]]); + assert_eq!(db.completion("p").await?.default_schema, "public"); + } + } + + // A manual-commit console switches schemas on its session connection. + db.set_manual_commit("p", true).await?; + let r = db.execute_from("p", "SELECT tag FROM items", 0, 1, &c, Some("oi_a")).await?; + assert_eq!(r.rows, vec![vec![Some("from a".into())]]); + let r = db.execute_from("p", "SELECT tag FROM items", 0, 1, &c, Some("oi_b")).await?; + assert_eq!(r.rows, vec![vec![Some("from b".into())]]); + db.finish_transaction("p", true).await?; + for commit in [false, true] { + db.execute_from("p", "SELECT 1", 0, 1, &c, Some("oi_a")).await?; + db.finish_transaction("p", commit).await?; + let r = db.execute("p", "SHOW search_path", 1, &c).await?; + assert_eq!(r.rows, vec![vec![Some("\"$user\", public".into())]]); + } + db.set_manual_commit("p", false).await?; + db.execute("p", "SET search_path TO oi_b", 1, &c).await?; + let r = db.execute("p", "SHOW search_path", 1, &c).await?; + assert_eq!(r.rows, vec![vec![Some("\"$user\", public".into())]]); + + // An unknown schema never resolves the statement's tables (`SET search_path` + // itself does not validate names; resolution against it does). + let error = db.execute_from("p", "SELECT tag FROM items", 0, 1, &c, Some("oi_missing")).await.unwrap_err(); + assert!(error.message.to_lowercase().contains("does not exist"), "{}", error.message); + Ok(()) +} + +#[tokio::test] +async fn a_reused_physical_connection_restores_the_full_default_path() -> IdeResult<()> { + use ide_db::{DbOptions, Held, Pool}; + let Some((host, port, user, password)) = server() else { return Ok(()) }; + let source = DataSource { + driver: DbDriver::Postgres, + host, + port: Some(port), + user, + options: DbOptions { default_schema: "public, pg_catalog".into(), ..Default::default() }, + ..Default::default() + }; + let pool = Pool::connect(&source, Some(&password), None).await?; + let c = CancellationToken::new(); + // Occupy the other three slots: every operation below must reuse the fourth. + let _other = [Held::acquire(&pool).await?, Held::acquire(&pool).await?, Held::acquire(&pool).await?]; + let mut conn = Held::acquire(&pool).await?; + let backend = conn.execute(&pool, "SELECT pg_backend_pid()", 1, &c).await?.rows; + for schema in ["pg_catalog", "public", "pg_catalog"] { + conn.apply_context(Some(schema)).await?; + let r = conn.execute(&pool, "SELECT current_schema()", 1, &c).await?; + assert_eq!(r.rows, vec![vec![Some(schema.into())]]); + } + conn.apply_context(Some("public")).await?; + conn.execute(&pool, "SELECT set_config('search_path', 'pg_catalog', false)", 1, &c).await?; + conn.apply_context(Some("public")).await?; + assert_eq!(conn.execute(&pool, "SELECT current_schema()", 1, &c).await?.rows, vec![vec![Some("public".into())]]); + drop(conn); + assert_eq!(pool.rows("SHOW search_path", 1).await?, vec![vec![Some("public, pg_catalog".into())]]); + assert_eq!(pool.execute("SELECT pg_backend_pid()", 1, &c).await?.rows, backend); + let mut conn = Held::acquire(&pool).await?; + conn.command("BEGIN").await?; + conn.apply_context(Some("pg_catalog")).await?; + conn.command("ROLLBACK").await?; + conn.restore_default().await?; + assert_eq!( + conn.execute(&pool, "SHOW search_path", 1, &c).await?.rows, + vec![vec![Some("public, pg_catalog".into())]] + ); + Ok(()) +} diff --git a/crates/ide-db/tests/sqlite.rs b/crates/ide-db/tests/sqlite.rs index 06bf741a0..c13927549 100644 --- a/crates/ide-db/tests/sqlite.rs +++ b/crates/ide-db/tests/sqlite.rs @@ -132,16 +132,16 @@ async fn a_later_page_returns_only_its_rows() -> IdeResult<()> { let sql = "SELECT id FROM users ORDER BY id"; let ids = |r: &QueryResult| r.rows.iter().map(|row| row[0].clone().unwrap()).collect::>(); - let second = db.execute_from("s1", sql, 1, 1, &c).await?; + let second = db.execute_from("s1", sql, 1, 1, &c, None).await?; assert_eq!(ids(&second), vec!["2"]); assert!(second.has_more); assert_eq!(second.columns.len(), 1); - let last = db.execute_from("s1", sql, 2, 5, &c).await?; + let last = db.execute_from("s1", sql, 2, 5, &c, None).await?; assert_eq!(ids(&last), vec!["3"]); assert!(!last.has_more); - let past = db.execute_from("s1", sql, 9, 5, &c).await?; + let past = db.execute_from("s1", sql, 9, 5, &c, None).await?; assert!(past.rows.is_empty() && !past.has_more); assert_eq!(past.columns.len(), 1, "columns of a page past the end come from the prepared statement"); Ok(()) @@ -276,11 +276,11 @@ async fn ddl_of_tables_and_views() -> IdeResult<()> { async fn explains_queries_without_running_them() -> IdeResult<()> { let (_tmp, db) = setup().await; let c = CancellationToken::new(); - let plan = db.explain("s1", "SELECT * FROM users WHERE id > 1 ORDER BY note;", &c).await?; + let plan = db.explain("s1", "SELECT * FROM users WHERE id > 1 ORDER BY note;", &c, None).await?; let ops: Vec<_> = plan.iter().map(|n| n.operation.as_str()).collect(); assert_eq!(ops, vec!["Index Scan", "Order By"]); assert!(plan[0].params.starts_with("table: users;"), "{}", plan[0].params); - db.explain("s1", "DELETE FROM users", &c).await?; + db.explain("s1", "DELETE FROM users", &c, None).await?; assert_eq!(count(&db).await, "3"); Ok(()) } diff --git a/crates/ide-host/src/ops/db.rs b/crates/ide-host/src/ops/db.rs index 079542f26..5ab776d8a 100644 --- a/crates/ide-host/src/ops/db.rs +++ b/crates/ide-host/src/ops/db.rs @@ -103,7 +103,8 @@ pub(crate) fn register(d: &mut Dispatcher) { d.add("db_explain", |c| async move { let (request, id, statement) = (c.arg::("request")?, c.arg::("id")?, c.arg::("statement")?); - explain(&c.host, &c.window, request, &id, &statement).await + let schema = c.arg::>("schema")?; + explain(&c.host, &c.window, request, &id, &statement, schema.as_deref()).await }); d.add("db_history", |c| async move { let id = c.arg::("id")?; @@ -130,7 +131,8 @@ pub(crate) fn register(d: &mut Dispatcher) { c.arg::>("skip")?.unwrap_or(0), c.arg::("maxRows")?, ); - execute(&c.host, &c.window, request, &id, statements, skip, max_rows).await + let schema = c.arg::>("schema")?; + execute(&c.host, &c.window, request, &id, statements, skip, max_rows, schema.as_deref()).await }); d.add("db_table_page", |c| async move { let (request, id, schema, table) = ( @@ -289,10 +291,11 @@ pub async fn explain( request: RequestId, id: &str, statement: &str, + schema: Option<&str>, ) -> IdeResult> { let db = host.project(window)?.db.clone(); let cancel = host.begin_request(request); - let result = db.explain(id, statement, &cancel).await; + let result = db.explain(id, statement, &cancel, schema).await; host.end_request(request); result } @@ -327,7 +330,9 @@ pub fn statements(host: &Host, window: &str, id: &str, text: &str, from: u32, to /// Runs `statements` in order on data source `id`, stopping at the first /// error. Each result reads past its first `skip` rows (a later page) and keeps -/// `max_rows`. Cancel with `cancel_request(request)`. +/// `max_rows`. `schema` is the console's active context: the connection running +/// each statement is put in it first. Cancel with `cancel_request(request)`. +#[allow(clippy::too_many_arguments)] pub async fn execute( host: &Host, window: &str, @@ -336,12 +341,13 @@ pub async fn execute( statements: Vec, skip: u32, max_rows: u32, + schema: Option<&str>, ) -> IdeResult> { let db = host.project(window)?.db.clone(); let cancel = host.begin_request(request); let mut out = Vec::new(); for statement in statements { - match db.execute_from(id, &statement.text, skip as usize, max_rows.max(1) as usize, &cancel).await { + match db.execute_from(id, &statement.text, skip as usize, max_rows.max(1) as usize, &cancel, schema).await { Ok(result) => out.push(StatementOutcome { statement, result: Some(result), error: None }), Err(e) if e.kind == ide_core::error::ErrorKind::Cancelled => { host.end_request(request); diff --git a/src-tauri/src/commands/db.rs b/src-tauri/src/commands/db.rs index 247716196..ca13a42a2 100644 --- a/src-tauri/src/commands/db.rs +++ b/src-tauri/src/commands/db.rs @@ -181,8 +181,9 @@ pub async fn db_explain( request: RequestId, id: String, statement: String, + schema: Option, ) -> IdeResult> { - ide_host::ops::db::explain(&state.host, window.label(), request, &id, &statement).await + ide_host::ops::db::explain(&state.host, window.label(), request, &id, &statement, schema.as_deref()).await } #[tauri::command] @@ -244,6 +245,7 @@ pub fn db_statements( #[tauri::command] #[specta::specta] +#[allow(clippy::too_many_arguments)] pub async fn db_execute( window: Window, state: State<'_, AppState>, @@ -252,8 +254,10 @@ pub async fn db_execute( statements: Vec, skip: u32, max_rows: u32, + schema: Option, ) -> IdeResult> { - ide_host::ops::db::execute(&state.host, window.label(), request, &id, statements, skip, max_rows).await + ide_host::ops::db::execute(&state.host, window.label(), request, &id, statements, skip, max_rows, schema.as_deref()) + .await } #[tauri::command] diff --git a/src/bindings.ts b/src/bindings.ts index 733358c83..3522b6952 100644 --- a/src/bindings.ts +++ b/src/bindings.ts @@ -1375,7 +1375,7 @@ export const commands = { dbDdl: (id: string, schema: string, table: string) => typedError(__TAURI_INVOKE("db_ddl", { id, schema, table })), dbSchemaObjects: (id: string, schema: string) => typedError(__TAURI_INVOKE("db_schema_objects", { id, schema })), dbObjectDdl: (id: string, schema: string, kind: DbObjectKind, name: string) => typedError(__TAURI_INVOKE("db_object_ddl", { id, schema, kind, name })), - dbExplain: (request: RequestId, id: string, statement: string) => typedError(__TAURI_INVOKE("db_explain", { request, id, statement })), + dbExplain: (request: RequestId, id: string, statement: string, schema: string | null) => typedError(__TAURI_INVOKE("db_explain", { request, id, statement, schema })), dbHistory: (id: string) => typedError(__TAURI_INVOKE("db_history", { id })), dbRemoveHistory: (id: string, texts: string[] | null) => typedError(__TAURI_INVOKE("db_remove_history", { id, texts })), dbCompletion: (id: string) => typedError(__TAURI_INVOKE("db_completion", { id })), @@ -1388,7 +1388,7 @@ export const commands = { dbSetManualCommit: (id: string, manual: boolean) => typedError(__TAURI_INVOKE("db_set_manual_commit", { id, manual })), dbFinishTransaction: (id: string, commit: boolean) => typedError(__TAURI_INVOKE("db_finish_transaction", { id, commit })), dbStatements: (id: string, text: string, from: number, to: number) => typedError(__TAURI_INVOKE("db_statements", { id, text, from, to })), - dbExecute: (request: RequestId, id: string, statements: SqlStatement[], skip: number, maxRows: number) => typedError(__TAURI_INVOKE("db_execute", { request, id, statements, skip, maxRows })), + dbExecute: (request: RequestId, id: string, statements: SqlStatement[], skip: number, maxRows: number, schema: string | null) => typedError(__TAURI_INVOKE("db_execute", { request, id, statements, skip, maxRows, schema })), dbTablePage: (request: RequestId, id: string, schema: string, table: string, filter: string | null, sort: DbSortKey[], offset: number, limit: number) => typedError(__TAURI_INVOKE("db_table_page", { request, id, schema, table, filter, sort, offset, limit })), dbApplyEdits: (id: string, schema: string, table: string, edits: DbRowEdit[]) => typedError(__TAURI_INVOKE("db_apply_edits", { id, schema, table, edits })), dbConsolePath: (id: string) => typedError(__TAURI_INVOKE("db_console_path", { id })), diff --git a/src/features/database/Database.module.css b/src/features/database/Database.module.css index 4d59b7dcd..924f93121 100644 --- a/src/features/database/Database.module.css +++ b/src/features/database/Database.module.css @@ -238,6 +238,12 @@ margin: 0 4px; background: var(--border); } +.treeFilter { + flex: 1; + min-width: 80px; + max-width: 280px; + margin-left: 4px; +} .treeHost { flex: 1; min-height: 0; diff --git a/src/features/database/DatabaseView.tsx b/src/features/database/DatabaseView.tsx index 4a71fabbb..afd2627d8 100644 --- a/src/features/database/DatabaseView.tsx +++ b/src/features/database/DatabaseView.tsx @@ -15,13 +15,23 @@ import { Icon, IconButton, Menu, + SearchField, Tree, showContextMenu, tooltipAttrs, type TreeHandle, T, } from "@/ui"; -import { colorVar, databaseKey, passesObjectFilter, shownNames, sourceIcon, splitTarget, targetOf } from "./model"; +import { + colorVar, + databaseKey, + passesObjectFilter, + passesTreeFilter, + shownNames, + sourceIcon, + splitTarget, + targetOf, +} from "./model"; import { nodeTarget, schemaObjectOf, tableFolder, type DbNode as Node, type Folder, type Group } from "./nodes"; import { connect, @@ -111,10 +121,25 @@ export function DatabaseView() { [st.sources], ); - /** Schemas tab › Object filter of the data source of target `id`. */ + // While a quick filter is on, the schemas of expanded sources load their objects, so + // tables / routines matching the filter show up without expanding every schema node. + useEffect(() => { + if (!st.treeFilter.trim()) return; + const s = useDatabase.getState(); + for (const [target, list] of Object.entries(s.schemas)) { + if (list.status !== "ok") continue; + const { id, database } = splitTarget(target); + if (!s.treeExpanded.has(database ? databaseKey(id, database) : id)) continue; + for (const schema of list.value) { + if (!s.tables[tablesKey(target, schema)]) void loadTables(target, schema); + } + } + }, [st.treeFilter, st.treeExpanded, st.schemas]); + + /** Schemas tab › Object filter, plus the toolbar's quick filter, of the data source of target `id`. */ const shown = (id: string) => { const filter = sourceOf(id)?.introspection?.objectFilter; - return (o: { name: string }) => passesObjectFilter(o.name, filter); + return (o: { name: string }) => passesObjectFilter(o.name, filter) && passesTreeFilter(o.name, st.treeFilter); }; const tablesOf = (id: string, schema: string) => { const t = st.tables[tablesKey(id, schema)]; @@ -146,9 +171,18 @@ export function DatabaseView() { return shownNames(v.names, i?.allDatabases, i?.databases, v.current); }; + /** Under a quick filter a schema shows when its name matches, or its loaded objects do. */ + const schemaVisible = (id: string, schema: string) => + passesTreeFilter(schema, st.treeFilter) || + !!tablesOf(id, schema)?.length || + routinesOf(id, schema).length > 0 || + sequencesOf(id, schema).length > 0; + const schemasOf = (key: string, id: string) => status(key, st.schemas[id], (v) => - shownSchemas(id, v).map((schema) => ({ kind: "schema", key: `${id}/${schema}`, id, schema }) as Node), + shownSchemas(id, v) + .filter((schema) => schemaVisible(id, schema)) + .map((schema) => ({ kind: "schema", key: `${id}/${schema}`, id, schema }) as Node), ); const children = (n: Node): readonly Node[] | undefined => { @@ -484,6 +518,7 @@ function DatabaseToolbar() { const node = useDatabase((s) => s.selectedNode); const connected = useDatabase((s) => s.connected); const sources = useDatabase((s) => s.sources); + const treeFilter = useDatabase((s) => s.treeFilter); const plus = useRef(null); const newOpen = useDatabase((s) => s.newMenuOpen); const setNewOpen = (v: boolean) => useDatabase.setState({ newMenuOpen: v }); @@ -529,6 +564,13 @@ function DatabaseToolbar() { > DDL + useDatabase.setState({ treeFilter })} + placeholder={tx("Filter objects")} + aria-label={tx("Filter objects")} + /> ); } diff --git a/src/features/database/console.tsx b/src/features/database/console.tsx index 25d838ca7..2c17e3205 100644 --- a/src/features/database/console.tsx +++ b/src/features/database/console.tsx @@ -36,13 +36,14 @@ import { iconNode } from "@/editor/lsp/iconMarkup"; import type { DbDriver, DbStatementOutcome } from "@/platform/ipc"; import { shortcutLabel } from "@/platform/keymap"; import { IconButton, Select } from "@/ui"; -import { sourceLabel, splitTarget, targetOf } from "./model"; +import { activeSchema, schemaSwitchable, sourceLabel, splitTarget, targetOf } from "./model"; import { analyze, templateContexts, valuesSlot } from "./sql/analyze"; import { shouldAutoCompleteSql, sqlCandidates, type CandidateKind, type SqlCandidate } from "./sql/complete"; import { sqlSupport } from "./sqlSupport"; import { bindConsole, cancelConsole, + consoleSchema, consoleSource, executeConsole, explainConsole, @@ -51,7 +52,9 @@ import { loadCompletion, listsDatabases, loadDatabases, + loadSchemas, loadSources, + setConsoleSchema, setManualCommit, sourceById, targetLabel, @@ -232,10 +235,35 @@ function ConsoleBar({ path, view }: { path: string; view: EditorView }) { const databases = useDatabase((s) => s.databases); const boundId = bound ? splitTarget(bound).id : null; const supportsTransactions = sourceById(boundId)?.driver !== "clickhouse"; + const sqlServer = sourceById(boundId)?.driver === "sqlServer"; + const switchable = schemaSwitchable(sourceById(boundId)?.driver); + // The console's active schema: its pick, else the connection's default (`completion.defaultSchema`). + const pickedSchema = useDatabase((s) => s.consoleSchemas[path]); + const defaultSchema = useDatabase((s) => (bound ? s.completion[bound]?.defaultSchema : undefined)); + const schemaList = useDatabase((s) => (bound ? s.schemas[bound] : undefined)); + const schema = activeSchema(pickedSchema, defaultSchema); useEffect(() => { - if (boundId && listsDatabases(sourceById(boundId)) && !useDatabase.getState().databases[boundId]) + if (boundId && (sqlServer || listsDatabases(sourceById(boundId))) && !useDatabase.getState().databases[boundId]) void loadDatabases(boundId); - }, [boundId]); + }, [boundId, sqlServer]); + useEffect(() => { + if (bound && switchable && !sqlServer && !useDatabase.getState().schemas[bound]) void loadSchemas(bound); + }, [bound, switchable, sqlServer]); + const schemaOptions = useMemo(() => { + if (!switchable) return []; + const databaseList = boundId ? databases[boundId] : undefined; + const names = sqlServer + ? databaseList?.status === "ok" + ? databaseList.value.names + : [] + : schemaList?.status === "ok" + ? schemaList.value + : []; + const out = names.map((n) => ({ value: n, label: n })); + // The connection's default is offered even before its schema list loads. + if (defaultSchema && !names.includes(defaultSchema)) out.unshift({ value: defaultSchema, label: defaultSchema }); + return out; + }, [switchable, schemaList, defaultSchema, sqlServer, boundId, databases]); const targets = useMemo(() => { const out: Array<{ value: string; label: string; detail?: string; icon: string }> = []; for (const s of sources) { @@ -325,6 +353,17 @@ function ConsoleBar({ path, view }: { path: string; view: EditorView }) { )} + {switchable && ( + { expect(shownNames(names, false, ["app*"], "postgres")).toEqual(["app", "app_test", "postgres"]); expect(shownNames(names, false, [])).toEqual([]); }); + + it("quick-filters tree objects by a case-insensitive substring", () => { + expect(passesTreeFilter("users", "user")).toBe(true); + expect(passesTreeFilter("Users", "user")).toBe(true); + expect(passesTreeFilter("order_items", " IT ")).toBe(true); + expect(passesTreeFilter("users", "orders")).toBe(false); + expect(passesTreeFilter("users", "")).toBe(true); + expect(passesTreeFilter("users", " ")).toBe(true); + expect(passesTreeFilter("users", null)).toBe(true); + }); +}); + +describe("active schema", () => { + it("follows the console's pick, then the connection default", () => { + expect(activeSchema("app", "public")).toBe("app"); + expect(activeSchema(undefined, "public")).toBe("public"); + expect(activeSchema("", "public")).toBe("public"); + expect(activeSchema(undefined, undefined)).toBe(""); + expect(activeSchema(null, null)).toBe(""); + }); + + it("offers the schema dropdown to multi-schema drivers only", () => { + for (const driver of ["postgres", "mysql", "sqlServer", "oracle"] as const) + expect(schemaSwitchable(driver)).toBe(true); + for (const driver of ["sqlite", "clickhouse"] as const) expect(schemaSwitchable(driver)).toBe(false); + expect(schemaSwitchable(undefined)).toBe(false); + }); }); describe("Options tab", () => { diff --git a/src/features/database/model.ts b/src/features/database/model.ts index 98b2470df..12b68ceb9 100644 --- a/src/features/database/model.ts +++ b/src/features/database/model.ts @@ -2,6 +2,7 @@ import type { DbCellValue, DbColor, + DbDriver, DbProperty, DbRowEdit, DbSortKey, @@ -203,6 +204,28 @@ export function passesObjectFilter(name: string, filter: string | null | undefin return !patterns.length || patterns.some((p) => matchesPattern(name.toLowerCase(), p.toLowerCase())); } +/** + * The tree's quick filter (the toolbar's field): `name` contains it, + * case-insensitively (`*` / `?` still work through `matchesPattern` callers). + */ +export function passesTreeFilter(name: string, filter: string | null | undefined): boolean { + const q = (filter ?? "").trim().toLowerCase(); + return !q || name.toLowerCase().includes(q); +} + +/** + * Whether a console of `driver` gets the active-schema dropdown (a server with + * several schemas / databases to switch between; SQLite has one database only). + */ +export function schemaSwitchable(driver: DbDriver | null | undefined): boolean { + return driver === "postgres" || driver === "mysql" || driver === "sqlServer" || driver === "oracle"; +} + +/** The schema a console's statements run in: its pick, else the connection's default. */ +export function activeSchema(picked: string | null | undefined, defaultSchema: string | null | undefined): string { + return picked || defaultSchema || ""; +} + /** * The databases / schemas the tree shows (Schemas tab): all of them, or the * ones matching a chosen name or pattern; `keep` (the connected database) always stays. diff --git a/src/features/database/sqlSupport.ts b/src/features/database/sqlSupport.ts index 9799c2053..ba17e3798 100644 --- a/src/features/database/sqlSupport.ts +++ b/src/features/database/sqlSupport.ts @@ -20,7 +20,7 @@ import { offsetToPos } from "@/editor/lsp/positions"; import type { DbCompletion, DbCompletionTable, DbDriver } from "@/platform/ipc"; import { formatSql } from "./sql/format"; import { SqlResolver, type SqlSymbol } from "./sql/resolve"; -import { consoleSource, revealInTree, sourceById, useDatabase } from "./store"; +import { consoleSchema, consoleSource, revealInTree, sourceById, useDatabase } from "./store"; /** Producer id of the console's diagnostics (Problems view, hover). */ const INSPECTIONS = "sql-inspections"; @@ -28,10 +28,13 @@ const INSPECT_DELAY = 400; function sourceOf(path: string): { id: string | null; driver: DbDriver; model: DbCompletion | null } { const id = consoleSource(path); + const model = id ? (useDatabase.getState().completion[id] ?? null) : null; + // Unqualified names resolve in the console's active schema (IDEA's schema dropdown). + const schema = consoleSchema(path); return { id, driver: sourceById(id)?.driver ?? "sqlite", - model: id ? (useDatabase.getState().completion[id] ?? null) : null, + model: model && schema ? { ...model, defaultSchema: schema } : model, }; } @@ -64,9 +67,10 @@ function inspections(path: string): Extension { timer = setTimeout(run, INSPECT_DELAY); }; schedule(); - // The schema model arrives (or the console is bound to another source) after the editor opens. + // The schema model arrives (or the console is bound to another source / schema) after the editor opens. const unsubscribe = useDatabase.subscribe((s, prev) => { - if (s.completion !== prev.completion || s.consoles !== prev.consoles) schedule(); + if (s.completion !== prev.completion || s.consoles !== prev.consoles || s.consoleSchemas !== prev.consoleSchemas) + schedule(); }); return { update: (u) => { diff --git a/src/features/database/store.ts b/src/features/database/store.ts index 038646894..712267729 100644 --- a/src/features/database/store.ts +++ b/src/features/database/store.ts @@ -33,7 +33,16 @@ import { } from "@/workbench/contributions"; import { workflowSettings } from "@/platform/workflowSettings"; import { ask, confirm } from "@/workbench/messages"; -import { colorVar, connectRemedy, databaseKey, effectiveSsh, idleTargets, newSourceId, splitTarget } from "./model"; +import { + activeSchema, + colorVar, + connectRemedy, + databaseKey, + effectiveSsh, + idleTargets, + newSourceId, + splitTarget, +} from "./model"; import { tableFolder, type DbNode } from "./nodes"; import type { ExportData, ExportFormat } from "./export"; import type { DbObjectRef } from "./objects"; @@ -46,6 +55,7 @@ import { tx } from "@/platform/i18n"; export const DB_TOOL_WINDOW = "database"; export const RESULTS_VIEW = "dbResults"; const BINDINGS_KEY = "openidea.database.consoles"; +const CONSOLE_SCHEMAS_KEY = "openidea.database.consoleSchemas"; const MAX_TABS = 12; export interface ResultTab { @@ -127,6 +137,10 @@ interface DatabaseState { objects: Record; /** Console / SQL file path → data source id. */ consoles: Record; + /** Console path → its active schema (unqualified names resolve there; see `consoleSchema`). */ + consoleSchemas: Record; + /** The Database tree's quick object filter (empty: everything shows). */ + treeFilter: string; completion: Record; results: ResultTab[]; /** Result tab id, or `OUTPUT_TAB`. */ @@ -163,6 +177,15 @@ function readBindings(): Record { } } +function readConsoleSchemas(): Record { + try { + const raw = localStorage.getItem(CONSOLE_SCHEMAS_KEY); + return raw ? (JSON.parse(raw) as Record) : {}; + } catch { + return {}; + } +} + export const useDatabase = create(() => ({ sources: [], loaded: false, @@ -173,6 +196,8 @@ export const useDatabase = create(() => ({ details: {}, objects: {}, consoles: readBindings(), + consoleSchemas: readConsoleSchemas(), + treeFilter: "", completion: {}, results: [], activeResult: null, @@ -593,7 +618,39 @@ export function bindConsole(path: string, id: string | null) { } catch { // Per-viewer convenience only. } - return { consoles }; + // Re-bound to another data source: the console starts from its default schema again. + const consoleSchemas = { ...s.consoleSchemas }; + if (consoles[path] !== s.consoles[path]) { + delete consoleSchemas[path]; + try { + localStorage.setItem(CONSOLE_SCHEMAS_KEY, JSON.stringify(consoleSchemas)); + } catch { + // Per-viewer convenience only. + } + } + return { consoles, consoleSchemas }; + }); +} + +/** The schema a console's statements run in: its pick, else the connection's default. */ +export function consoleSchema(path: string): string { + const s = useDatabase.getState(); + const id = consoleSource(path); + return activeSchema(s.consoleSchemas[path], id ? s.completion[id]?.defaultSchema : undefined); +} + +/** Sets (or clears, for `null`) the console's active schema and remembers it per console file. */ +export function setConsoleSchema(path: string, schema: string | null) { + useDatabase.setState((s) => { + const consoleSchemas = { ...s.consoleSchemas }; + if (schema) consoleSchemas[path] = schema; + else delete consoleSchemas[path]; + try { + localStorage.setItem(CONSOLE_SCHEMAS_KEY, JSON.stringify(consoleSchemas)); + } catch { + // Per-viewer convenience only. + } + return { consoleSchemas }; }); } @@ -653,6 +710,8 @@ function promptOf(id: string): string { /** * Runs `statements` on `id`, logging to Output; errors stop the batch (backend). * Results keep `maxRows` rows after the first `skip`, which the backend reads past. + * `schema` is a console's active context: the backend puts the connection that runs + * each statement in it, so unqualified names resolve there. */ async function run( id: string, @@ -660,12 +719,16 @@ async function run( request: RequestId, maxRows = useDatabase.getState().maxRows, skip = 0, + schema?: string, ): Promise { const wasConnected = useDatabase.getState().connected.has(id); - const outcomes = await withConnection(id, () => call(commands.dbExecute(request, id, statements, skip, maxRows))); + const outcomes = await withConnection(id, () => + call(commands.dbExecute(request, id, statements, skip, maxRows, schema || null)), + ); if (!wasConnected) log("info", "Connected"); + const prompt = schema || promptOf(id); for (const o of outcomes) { - log("statement", o.statement.text, `${promptOf(id)}>`); + log("statement", o.statement.text, `${prompt}>`); log(o.error ? "error" : "result", o.error ? `[${targetLabel(id)}] ${o.error}` : outcomeSummary(o)); } if (useDatabase.getState().manualCommit.has(id) && outcomes.some((o) => !o.error && !READS.test(o.statement.text))) @@ -806,7 +869,14 @@ export async function executeConsole( useDatabase.setState((s) => ({ running: { ...s.running, [path]: request } })); try { const pageSize = useDatabase.getState().maxRows; - const outcomes = await run(id, statements, request, pageSize); + const outcomes = await run( + id, + statements, + request, + pageSize, + 0, + useDatabase.getState().consoleSchemas[path] || undefined, + ); const tabs: ResultTab[] = outcomes .filter((o) => o.result?.columns.length) .map((outcome) => { @@ -870,9 +940,11 @@ function showPlan(path: string, id: string, statement: SqlStatement, plan: DbPla showServiceView(RESULTS_VIEW, false); } -async function explain(id: string, statement: SqlStatement): Promise { - const plan = await withConnection(id, () => call(commands.dbExplain(nextRequestId(), id, statement.text))); - log("statement", `explain ${statement.text}`, `${promptOf(id)}>`); +async function explain(id: string, statement: SqlStatement, schema?: string): Promise { + const plan = await withConnection(id, () => + call(commands.dbExplain(nextRequestId(), id, statement.text, schema || null)), + ); + log("statement", `explain ${statement.text}`, `${schema || promptOf(id)}>`); log("result", `${plan.length} plan node${plan.length === 1 ? "" : "s"}`); return plan; } @@ -884,7 +956,12 @@ export async function explainConsole(path: string, text: string, from: number, t try { const [statement] = await consoleStatements(id, text, from, to); if (!statement) return; - showPlan(path, id, statement, await explain(id, statement)); + showPlan( + path, + id, + statement, + await explain(id, statement, useDatabase.getState().consoleSchemas[path] || undefined), + ); } catch (e) { if (isCancelled(e)) return; log("error", `[${targetLabel(id)}] ${message(e)}`); @@ -930,7 +1007,14 @@ export async function rerunResult(tabId: number): Promise { if (tab.plan) { patchResult(tabId, { busy: true }); try { - patchResult(tabId, { plan: await explain(tab.sourceId, tab.outcome.statement), busy: false }); + patchResult(tabId, { + plan: await explain( + tab.sourceId, + tab.outcome.statement, + useDatabase.getState().consoleSchemas[tab.path] || undefined, + ), + busy: false, + }); } catch (e) { patchResult(tabId, { busy: false }); if (!isCancelled(e)) log("error", `[${targetLabel(tab.sourceId)}] ${message(e)}`); @@ -961,7 +1045,14 @@ export async function loadResultPage(tabId: number, page: number, pageSize: numb patchResult(tabId, { busy: true }); const { maxRows, skip } = pageFetch(page, pageSize); try { - const [outcome] = await run(tab.sourceId, [tab.outcome.statement], nextRequestId(), maxRows, skip); + const [outcome] = await run( + tab.sourceId, + [tab.outcome.statement], + nextRequestId(), + maxRows, + skip, + useDatabase.getState().consoleSchemas[tab.path] || undefined, + ); const r = outcome?.result; if (r) { const sizeChanged = pageSize !== tab.pageSize; @@ -999,7 +1090,14 @@ export async function countResult(tabId: number): Promise { const text = countQuery(tab.outcome.statement.text)!; patchResult(tabId, { busy: true }); try { - const [outcome] = await run(tab.sourceId, [{ from: 0, to: text.length, text }], nextRequestId(), 1); + const [outcome] = await run( + tab.sourceId, + [{ from: 0, to: text.length, text }], + nextRequestId(), + 1, + 0, + useDatabase.getState().consoleSchemas[tab.path] || undefined, + ); const total = Number(outcome?.result?.rows[0]?.[0]); const known = outcome?.error || !Number.isFinite(total) ? null : total; patchResult(tabId, { busy: false, ...(known === null ? {} : { total: known }) }); diff --git a/src/platform/locales/text/en.json b/src/platform/locales/text/en.json index e72e108db..5a4d0189e 100644 --- a/src/platform/locales/text/en.json +++ b/src/platform/locales/text/en.json @@ -1,4 +1,9 @@ { + "this data source does not list databases": "this data source does not list databases", + "commit or roll back before clearing the selected database": "commit or roll back before clearing the selected database", + "connect_timeout must be greater than zero": "connect_timeout must be greater than zero", + "database connection is closed": "database connection is closed", + "Database": "Database", "Run `{command}` in a terminal.": "Run `{command}` in a terminal.", "Run `{command}` in a terminal and sign in.": "Run `{command}` in a terminal and sign in.", "{name} needs you to sign in. {hint}": "{name} needs you to sign in. {hint}", diff --git a/src/platform/locales/text/zh.json b/src/platform/locales/text/zh.json index 7b713571d..d09ef0908 100644 --- a/src/platform/locales/text/zh.json +++ b/src/platform/locales/text/zh.json @@ -1,4 +1,8 @@ { + "this data source does not list databases": "此数据源不支持列出数据库", + "commit or roll back before clearing the selected database": "清除所选数据库前,请先提交或回滚事务", + "connect_timeout must be greater than zero": "connect_timeout 必须大于零", + "database connection is closed": "数据库连接已关闭", "Run `{command}` in a terminal.": "在终端中运行 `{command}`。", "Run `{command}` in a terminal and sign in.": "在终端中运行 `{command}` 并登录。", "{name} needs you to sign in. {hint}": "{name} 需要你登录。{hint}", @@ -934,6 +938,7 @@ "Filter by name, PID or command line": "按名称、PID 或命令行过滤", "Filter by tool or language": "按工具或语言过滤", "Filter language servers": "过滤语言服务器", + "Filter objects": "过滤对象", "Filter plugins": "过滤插件", "Filter processes": "过滤进程", "Filter recent files": "过滤最近的文件", @@ -3935,6 +3940,7 @@ "Explain Plan": "解释计划", "Browse Query History": "浏览查询历史记录", "": "<数据源>", + "": "<架构>", "Cannot explain the statement": "无法解释该语句", "Cannot get the definition of {name}": "无法获取 {name} 的定义", "Script finished: {length} statement": "脚本已完成:{length} 条语句", @@ -6161,6 +6167,7 @@ "Default project-wide schema": "默认项目范围架构", "the configuration is code and is not evaluated": "配置为代码,不会被求值", "schema is empty": "架构为空", + "Schema": "架构", "Schema discovery summary": "架构发现摘要", "Schema errors": "架构错误", "No errors found": "未发现错误",