feat(database): select active schema per SQL console (#1807)

Closes #447. Each SQL console gets an active-schema dropdown (PostgreSQL / MySQL / Oracle schemas; SQL Server selects the database, since T-SQL has no per-session default schema), remembered per console file; hidden for SQLite and ClickHouse. Every statement run through the pool first asserts the console's context, or the connection's default when none is selected (skipped when the tracked context already matches), and manual-commit sessions restore the default before the connection returns to the pool, so one console's schema never leaks into another.

Co-authored-by: Inf1mum <[email protected]>
Цей коміт міститься в:
Vadim LazukтаInf1mum authored and GitHub committed 2026-11-08 00:32:52 +08:00
1 джерело daad91777f
коміт f407880204
20 files changed
+991 -139

No files matched your search

+355 -84
Переглянути файл
@@ -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<String>,
current: Option<String>,
known: bool,
}
impl Context {
pub(crate) fn new(default: Option<String>) -> Self {
Self { current: default.clone(), default, known: true }
}
pub(crate) fn set_current(&mut self, current: Option<String>) {
self.current = current;
}
fn wanted(&self, driver: DbDriver, schema: Option<&str>) -> Option<String> {
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<Vec<String>>,
timeout: Duration,
}
pub struct $handle {
conn: Option<$connection>,
context: Context,
options: $options,
init: Arc<Vec<String>>,
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<Self::Connection, Self::Error> {
tokio::time::timeout(self.timeout, async {
let mut conn = <$connection>::connect_with(&self.options).await?;
let context_sql = |value: Option<String>| {
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<M: ManageConnection> {
inner: Bb8Pool<M>,
closed: AtomicBool,
}
impl<M: ManageConnection> ServerPool<M> {
async fn get(&self) -> IdeResult<bb8::PooledConnection<'static, M>>
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<tokio_util::compat::Compat<tokio::net::TcpStream>>;
pub struct MssqlConnection(Option<MssqlClient>);
pub struct MssqlConnection(Option<MssqlClient>, 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<PgManager>),
Mysql(ServerPool<MysqlManager>),
SqlServer(Bb8Pool<MssqlConnectionManager>),
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<Fetched> = 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<u64> {
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<Sqlite>),
Postgres(PoolConnection<Postgres>),
Mysql(PoolConnection<MySql>),
Postgres(Box<bb8::PooledConnection<'static, PgManager>>),
Mysql(Box<bb8::PooledConnection<'static, MysqlManager>>),
SqlServer(Box<bb8::PooledConnection<'static, MssqlConnectionManager>>),
Clickhouse(crate::clickhouse::ClickhousePool),
Oracle(crate::oracle::OracleConnection),
@@ -889,8 +1040,8 @@ impl Held {
pub async fn acquire(pool: &Pool) -> IdeResult<Self> {
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<Fetched> {
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());
}
+51 -11
Переглянути файл
@@ -315,7 +315,7 @@ fn keep_alive(pool: &Arc<Pool>, 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<Vec<Option<String>>>| 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<Vec<Option<String>>>| 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<QueryResult> {
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<QueryResult> {
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<Vec<PlanNode>> {
/// 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<Vec<PlanNode>> {
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<QueryResult> {
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;
+15 -5
Переглянути файл
@@ -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>) -> OracleError {
/// discarded (a cancelled statement leaves it in the blocking task).
pub struct OracleHandle {
conn: Option<oracle::Connection>,
context: crate::conn::Context,
}
pub struct OracleConnectionManager(Arc<State>);
@@ -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<oracle::Connection> {
fn open(state: &State) -> oracle::Result<OracleHandle> {
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<oracle::Connection> {
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, &current));
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).
+31
Переглянути файл
@@ -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<String> {
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 = [
+1 -1
Переглянути файл
@@ -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
+98
Переглянути файл
@@ -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(())
}
+135
Переглянути файл
@@ -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<DbManager> {
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(())
}
+5 -5
Переглянути файл
@@ -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::<Vec<_>>();
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(())
}
+11 -5
Переглянути файл
@@ -103,7 +103,8 @@ pub(crate) fn register(d: &mut Dispatcher) {
d.add("db_explain", |c| async move {
let (request, id, statement) =
(c.arg::<RequestId>("request")?, c.arg::<String>("id")?, c.arg::<String>("statement")?);
explain(&c.host, &c.window, request, &id, &statement).await
let schema = c.arg::<Option<String>>("schema")?;
explain(&c.host, &c.window, request, &id, &statement, schema.as_deref()).await
});
d.add("db_history", |c| async move {
let id = c.arg::<String>("id")?;
@@ -130,7 +131,8 @@ pub(crate) fn register(d: &mut Dispatcher) {
c.arg::<Option<u32>>("skip")?.unwrap_or(0),
c.arg::<u32>("maxRows")?,
);
execute(&c.host, &c.window, request, &id, statements, skip, max_rows).await
let schema = c.arg::<Option<String>>("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<Vec<PlanNode>> {
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<SqlStatement>,
skip: u32,
max_rows: u32,
schema: Option<&str>,
) -> IdeResult<Vec<StatementOutcome>> {
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);
+6 -2
Переглянути файл
@@ -181,8 +181,9 @@ pub async fn db_explain(
request: RequestId,
id: String,
statement: String,
schema: Option<String>,
) -> IdeResult<Vec<ide_db::PlanNode>> {
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<SqlStatement>,
skip: u32,
max_rows: u32,
schema: Option<String>,
) -> IdeResult<Vec<StatementOutcome>> {
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]
+2 -2
Переглянути файл
@@ -1375,7 +1375,7 @@ export const commands = {
dbDdl: (id: string, schema: string, table: string) => typedError<string, IdeError>(__TAURI_INVOKE("db_ddl", { id, schema, table })),
dbSchemaObjects: (id: string, schema: string) => typedError<DbSchemaObjects, IdeError>(__TAURI_INVOKE("db_schema_objects", { id, schema })),
dbObjectDdl: (id: string, schema: string, kind: DbObjectKind, name: string) => typedError<string, IdeError>(__TAURI_INVOKE("db_object_ddl", { id, schema, kind, name })),
dbExplain: (request: RequestId, id: string, statement: string) => typedError<DbPlanNode[], IdeError>(__TAURI_INVOKE("db_explain", { request, id, statement })),
dbExplain: (request: RequestId, id: string, statement: string, schema: string | null) => typedError<DbPlanNode[], IdeError>(__TAURI_INVOKE("db_explain", { request, id, statement, schema })),
dbHistory: (id: string) => typedError<DbHistoryEntry[], IdeError>(__TAURI_INVOKE("db_history", { id })),
dbRemoveHistory: (id: string, texts: string[] | null) => typedError<null, IdeError>(__TAURI_INVOKE("db_remove_history", { id, texts })),
dbCompletion: (id: string) => typedError<DbCompletion, IdeError>(__TAURI_INVOKE("db_completion", { id })),
@@ -1388,7 +1388,7 @@ export const commands = {
dbSetManualCommit: (id: string, manual: boolean) => typedError<null, IdeError>(__TAURI_INVOKE("db_set_manual_commit", { id, manual })),
dbFinishTransaction: (id: string, commit: boolean) => typedError<null, IdeError>(__TAURI_INVOKE("db_finish_transaction", { id, commit })),
dbStatements: (id: string, text: string, from: number, to: number) => typedError<SqlStatement[], IdeError>(__TAURI_INVOKE("db_statements", { id, text, from, to })),
dbExecute: (request: RequestId, id: string, statements: SqlStatement[], skip: number, maxRows: number) => typedError<DbStatementOutcome[], IdeError>(__TAURI_INVOKE("db_execute", { request, id, statements, skip, maxRows })),
dbExecute: (request: RequestId, id: string, statements: SqlStatement[], skip: number, maxRows: number, schema: string | null) => typedError<DbStatementOutcome[], IdeError>(__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<DbTablePage, IdeError>(__TAURI_INVOKE("db_table_page", { request, id, schema, table, filter, sort, offset, limit })),
dbApplyEdits: (id: string, schema: string, table: string, edits: DbRowEdit[]) => typedError<number | null, IdeError>(__TAURI_INVOKE("db_apply_edits", { id, schema, table, edits })),
dbConsolePath: (id: string) => typedError<IpcPath, IdeError>(__TAURI_INVOKE("db_console_path", { id })),
+6
Переглянути файл
@@ -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;
+46 -4
Переглянути файл
@@ -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<HTMLButtonElement>(null);
const newOpen = useDatabase((s) => s.newMenuOpen);
const setNewOpen = (v: boolean) => useDatabase.setState({ newMenuOpen: v });
@@ -529,6 +564,13 @@ function DatabaseToolbar() {
>
<T>DDL</T>
</button>
<SearchField
className={styles.treeFilter}
value={treeFilter}
onChange={(treeFilter) => useDatabase.setState({ treeFilter })}
placeholder={tx("Filter objects")}
aria-label={tx("Filter objects")}
/>
</div>
);
}
+46 -4
Переглянути файл
@@ -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 }) {
</>
)}
<span className={styles.spacer} />
{switchable && (
<Select
aria-label={sqlServer ? "Database" : "Schema"}
value={schema || null}
placeholder="<schema>"
options={schemaOptions}
searchable={schemaOptions.length > 10}
onChange={(v) => setConsoleSchema(path, v)}
disabled={!bound}
/>
)}
<Select
aria-label="Data source"
value={bound ?? null}
@@ -426,6 +465,9 @@ function completionSource(path: string): CompletionSource {
const driver: DbDriver = source?.driver ?? "sqlite";
const model = id ? (useDatabase.getState().completion[id] ?? null) : null;
if (id && !model) void loadCompletion(id);
// Unqualified names resolve in the console's active schema (IDEA's schema dropdown).
const schema = consoleSchema(path);
const scoped = model && schema ? { ...model, defaultSchema: schema } : model;
// A selected template placeholder (`alias`) is about to be replaced: analyze without it.
const sel = ctx.state.selection.main;
const placeholder = !sel.empty && (ctx.pos === sel.from || ctx.pos === sel.to);
@@ -434,7 +476,7 @@ function completionSource(path: string): CompletionSource {
? analyze(doc.slice(0, sel.from) + doc.slice(sel.to), sel.from, driver)
: analyze(doc, ctx.pos, driver);
if (!ctx.explicit && !shouldAutoCompleteSql(analysis)) return null;
const candidates = sqlCandidates({ model, analysis, driver, sourceName: source?.name ?? undefined });
const candidates = sqlCandidates({ model: scoped, analysis, driver, sourceName: source?.name ?? undefined });
if (!candidates.length) return null;
return {
from: analysis.prefix.from,
+30
Переглянути файл
@@ -1,6 +1,7 @@
import { describe, expect, it } from "vitest";
import { exportData, type ExportData } from "./export";
import {
activeSchema,
buildRowEdits,
colorVar,
connectRemedy,
@@ -9,6 +10,8 @@ import {
idleTargets,
matchesPattern,
parseSourceUrl,
passesTreeFilter,
schemaSwitchable,
shownNames,
sourceLabel,
splitTarget,
@@ -294,6 +297,33 @@ describe("Schemas tab", () => {
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", () => {
+23
Переглянути файл
@@ -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.
+8 -4
Переглянути файл
@@ -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) => {
+110 -12
Переглянути файл
@@ -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<string, DbSchemaObjects>;
/** Console / SQL file path → data source id. */
consoles: Record<string, string>;
/** Console path → its active schema (unqualified names resolve there; see `consoleSchema`). */
consoleSchemas: Record<string, string>;
/** The Database tree's quick object filter (empty: everything shows). */
treeFilter: string;
completion: Record<string, DbCompletion>;
results: ResultTab[];
/** Result tab id, or `OUTPUT_TAB`. */
@@ -163,6 +177,15 @@ function readBindings(): Record<string, string> {
}
}
function readConsoleSchemas(): Record<string, string> {
try {
const raw = localStorage.getItem(CONSOLE_SCHEMAS_KEY);
return raw ? (JSON.parse(raw) as Record<string, string>) : {};
} catch {
return {};
}
}
export const useDatabase = create<DatabaseState>(() => ({
sources: [],
loaded: false,
@@ -173,6 +196,8 @@ export const useDatabase = create<DatabaseState>(() => ({
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<DbStatementOutcome[]> {
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<DbPlanNode[]> {
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<DbPlanNode[]> {
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<void> {
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<number | null> {
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 }) });
+5
Переглянути файл
@@ -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}",
+7
Переглянути файл
@@ -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": "浏览查询历史记录",
"<data source>": "<数据源>",
"<schema>": "<架构>",
"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": "未发现错误",