xiaoxue_admin
/
dbx
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity
dbx/crates/dbx-core/src/db/sqlserver.rs

3377 lines
140 KiB
Rust
Raw Blame History

This file contains ambiguous Unicode characters

This file contains Unicode characters that might be confused with other characters. If you think that this is intentional, you can safely ignore this warning. Use the Escape button to reveal them.

use crate::query::MAX_ROWS;
use crate::sql::starts_with_executable_sql_keyword;
use crate::types::{
ColumnInfo, DatabaseInfo, ForeignKeyInfo, IndexInfo, LinkedServerInfo, ObjectStatistics, QueryResult, TableInfo,
TriggerInfo,
};
use futures::{FutureExt, TryStreamExt};
use sqlparser::ast::{Expr, SelectItem, SetExpr, Statement};
use sqlparser::dialect::MsSqlDialect;
use sqlparser::parser::Parser;
use std::future::Future;
use std::panic::AssertUnwindSafe;
use std::sync::{Arc as StdArc, Mutex as StdMutex};
use std::time::{Duration, Instant};
use tiberius::{AuthMethod, Client, ColumnData, ColumnType, Config, FromSql, QueryItem, QueryStream, Row, SqlBrowser};
use tokio::net::TcpStream;
use tokio_util::compat::{Compat, TokioAsyncWriteCompatExt};
use tokio_util::sync::CancellationToken;
use tracing::instrument::WithSubscriber;
use tracing::{Event, Level, Subscriber};
use tracing_subscriber::layer::{Context as LayerContext, SubscriberExt};
use tracing_subscriber::Layer;
pub type SqlServerClient = Client<Compat<TcpStream>>;
pub const SQLSERVER_DRIVER_PANIC_ERROR_PREFIX: &str = "SQL Server driver panic:";
pub const SQLSERVER_LEGACY_DRIVER_PROFILE: &str = "sqlserver-legacy";
pub const SQLSERVER_LEGACY_DRIVER_LABEL: &str = "SQL Server legacy compatibility component";
const SQLSERVER_DEFAULT_PORT: u16 = 1433;
const SQLSERVER_UNSAFE_PROBE_BLOCK_ERROR_PREFIX: &str = "SQL Server unsafe result type:";
const SQLSERVER_RESULT_TYPE_PROBE_SQL: &str = "\
DECLARE @dbx_use_describe_dmv bit = CASE \
WHEN @P2 = 0 AND OBJECT_ID(N'sys.dm_exec_describe_first_result_set') IS NOT NULL THEN 1 ELSE 0 END; \
SELECT @dbx_use_describe_dmv AS dbx_use_describe_dmv; \
IF @dbx_use_describe_dmv = 1 \
BEGIN \
EXEC sys.sp_executesql \
N'SELECT name, system_type_name, user_type_schema, user_type_name \
FROM sys.dm_exec_describe_first_result_set(@sql, NULL, 0) \
WHERE error_number IS NULL AND is_hidden = 0 \
ORDER BY column_ordinal', \
N'@sql nvarchar(max)', @sql = @P1 \
END \
ELSE \
BEGIN \
DECLARE @dbx_probe_table sysname = N'##dbx_result_type_probe_' + CONVERT(nvarchar(12), @@SPID); \
DECLARE @dbx_probe_object nvarchar(258) = N'tempdb..' + QUOTENAME(@dbx_probe_table); \
DECLARE @dbx_probe_sql nvarchar(max); \
BEGIN TRY \
SET @dbx_probe_sql = N'SELECT TOP (0) * INTO ' + QUOTENAME(@dbx_probe_table) + \
N' FROM ' + @P3; \
EXEC sys.sp_executesql @dbx_probe_sql; \
SELECT c.name, TYPE_NAME(c.system_type_id) AS system_type_name, \
SCHEMA_NAME(t.schema_id) AS user_type_schema, t.name AS user_type_name \
FROM tempdb.sys.columns c \
JOIN tempdb.sys.types t ON c.user_type_id = t.user_type_id \
WHERE c.object_id = OBJECT_ID(@dbx_probe_object) \
ORDER BY c.column_id; \
SET @dbx_probe_sql = N'DROP TABLE ' + QUOTENAME(@dbx_probe_table); \
EXEC sys.sp_executesql @dbx_probe_sql; \
END TRY \
BEGIN CATCH \
IF OBJECT_ID(@dbx_probe_object) IS NOT NULL \
BEGIN \
SET @dbx_probe_sql = N'DROP TABLE ' + QUOTENAME(@dbx_probe_table); \
EXEC sys.sp_executesql @dbx_probe_sql; \
END \
DECLARE @dbx_probe_error nvarchar(2048) = ERROR_MESSAGE(); \
RAISERROR(@dbx_probe_error, 16, 1); \
END CATCH \
END";
const SIMPLE_QUERY_MODULE_KEYWORDS: &[&str] = &["FUNCTION", "PROC", "PROCEDURE", "TRIGGER", "VIEW"];
// Match JDBC/tiberius `encrypt=false`: encrypt only login, then drop back to raw TDS.
const SQLSERVER_LEGACY_ENCRYPTION_LEVEL: tiberius::EncryptionLevel = tiberius::EncryptionLevel::Off;
// Some very old SQL Server setups only accepted DBX <= 0.5.48 because the fallback
// advertised no encryption support at all. Keep it as the last-resort compatibility path.
const SQLSERVER_UNSUPPORTED_ENCRYPTION_LEVEL: tiberius::EncryptionLevel = tiberius::EncryptionLevel::NotSupported;
const SQLSERVER_LEGACY_ENCRYPTION_FALLBACKS: [(&str, tiberius::EncryptionLevel); 2] = [
("login-only encryption", SQLSERVER_LEGACY_ENCRYPTION_LEVEL),
("no-encryption compatibility fallback", SQLSERVER_UNSUPPORTED_ENCRYPTION_LEVEL),
];
#[derive(Debug, Clone, serde::Serialize, serde::Deserialize)]
pub struct SqlServerColumnMetadata {
#[serde(flatten)]
pub column: ColumnInfo,
pub is_identity: bool,
pub is_computed: bool,
pub is_hidden: bool,
pub generated_always_type: i32,
}
#[derive(Debug, PartialEq, Eq)]
struct SqlServerEndpoint<'a> {
host: &'a str,
instance_name: Option<&'a str>,
}
fn sqlserver_endpoint(host: &str) -> SqlServerEndpoint<'_> {
if let Some((server, instance)) = host.split_once('\\') {
if !server.trim().is_empty() && !instance.trim().is_empty() {
return SqlServerEndpoint { host: server.trim(), instance_name: Some(instance.trim()) };
}
}
SqlServerEndpoint { host: host.trim(), instance_name: None }
}
fn sqlserver_uses_named_instance_resolution(endpoint: &SqlServerEndpoint<'_>, port: u16, port_explicit: bool) -> bool {
endpoint.instance_name.is_some() && (port == 0 || (port == SQLSERVER_DEFAULT_PORT && !port_explicit))
}
fn query_result_row_limit(max_rows: Option<usize>) -> usize {
max_rows.unwrap_or(MAX_ROWS).max(1)
}
pub async fn connect(
host: &str,
port: u16,
user: &str,
pass: &str,
database: Option<&str>,
_url_params: Option<&str>,
timeout: Duration,
) -> Result<SqlServerClient, String> {
connect_with_port_explicit(host, port, false, user, pass, database, timeout).await
}
pub async fn connect_with_port_explicit(
host: &str,
port: u16,
port_explicit: bool,
user: &str,
pass: &str,
database: Option<&str>,
timeout: Duration,
) -> Result<SqlServerClient, String> {
match try_connect(host, port, port_explicit, user, pass, database, tiberius::EncryptionLevel::Required, timeout)
.await
{
Ok(client) => Ok(client),
Err(encrypted_error) => {
try_connect_legacy_sqlserver_encryption(host, port, port_explicit, user, pass, database, timeout)
.await
.map_err(|plain_error| {
if is_sqlserver_tls_handshake_error(&encrypted_error) {
format!(
"{encrypted_error}\n\nThis may be caused by an old SQL Server TLS/encryption configuration. \
If you are connecting to SQL Server 2008/2008 R2/2012 or another legacy instance, \
try SQL Server legacy compatibility mode. It first behaves like encrypt=false and, \
when explicitly enabled, DBX can also fall back to the SQL Server legacy compatibility \
driver for TLS 1.0 encrypted transport. Only use this mode on trusted networks, VPNs, \
or SSH tunnels.\n\n\
Automatic native legacy fallback also failed: {plain_error}"
)
} else {
plain_error
}
})
}
}
}
async fn try_connect_legacy_sqlserver_encryption(
host: &str,
port: u16,
port_explicit: bool,
user: &str,
pass: &str,
database: Option<&str>,
timeout: Duration,
) -> Result<SqlServerClient, String> {
let mut errors = Vec::new();
for (label, encryption) in SQLSERVER_LEGACY_ENCRYPTION_FALLBACKS {
match try_connect(host, port, port_explicit, user, pass, database, encryption, timeout).await {
Ok(client) => return Ok(client),
Err(error) => errors.push(format!("{label} failed: {error}")),
}
}
Err(errors.join("\n"))
}
pub fn sqlserver_native_encryption_disabled(url_params: Option<&str>) -> bool {
let Some(params) = url_params.map(str::trim).filter(|params| !params.is_empty()) else {
return false;
};
params.trim_start_matches('?').split(['&', ';']).filter_map(|pair| pair.split_once('=')).any(|(key, value)| {
let key = key.trim();
let value = value.trim().to_ascii_lowercase();
let disabled = matches!(value.as_str(), "disabled" | "disable" | "false" | "0" | "off" | "no");
(key.eq_ignore_ascii_case("sqlserverEncryption") || key.eq_ignore_ascii_case("encrypt")) && disabled
})
}
fn is_sqlserver_tls_handshake_error(error: &str) -> bool {
let error = error.to_ascii_lowercase();
error.contains("tls") && (error.contains("handshake") || error.contains("eof") || error.contains("performing i/o"))
}
async fn try_connect(
host: &str,
port: u16,
port_explicit: bool,
user: &str,
pass: &str,
database: Option<&str>,
encryption: tiberius::EncryptionLevel,
timeout: Duration,
) -> Result<SqlServerClient, String> {
let mut config = Config::new();
let endpoint = sqlserver_endpoint(host);
let uses_named_instance_resolution = sqlserver_uses_named_instance_resolution(&endpoint, port, port_explicit);
config.host(endpoint.host);
if let Some(instance_name) = endpoint.instance_name.filter(|_| uses_named_instance_resolution) {
config.instance_name(instance_name);
} else {
config.port(port);
}
config.authentication(AuthMethod::sql_server(user, pass));
if let Some(db) = database {
config.database(db);
}
config.trust_cert();
config.encryption(encryption);
let tcp = if uses_named_instance_resolution {
tokio::time::timeout(timeout, TcpStream::connect_named(&config))
.await
.map_err(|_| format!("SQL Server connection timed out ({}s)", timeout.as_secs()))?
.map_err(|e| format!("SQL Server connection failed: {e}"))?
} else {
tokio::time::timeout(timeout, TcpStream::connect(config.get_addr()))
.await
.map_err(|_| format!("SQL Server connection timed out ({}s)", timeout.as_secs()))?
.map_err(|e| format!("SQL Server connection failed: {e}"))?
};
tokio::time::timeout(timeout, Client::connect(config, tcp.compat_write()))
.await
.map_err(|_| format!("SQL Server handshake timed out ({}s)", timeout.as_secs()))?
.map_err(|e| format!("SQL Server connection failed: {e}"))
}
fn row_to_json(row: &tiberius::Row) -> Vec<serde_json::Value> {
row.cells().map(|(_, cell)| sqlserver_cell_to_json(cell)).collect()
}
fn columns_from_metadata(metadata: &tiberius::ResultMetadata) -> Vec<String> {
metadata.columns().iter().map(|c| c.name().to_string()).collect()
}
fn sqlserver_column_type_name(column: &tiberius::Column) -> String {
match column.column_type() {
ColumnType::BigVarChar => "varchar".to_string(),
ColumnType::BigChar => "char".to_string(),
ColumnType::BigVarBin => "varbinary".to_string(),
ColumnType::BigBinary => "binary".to_string(),
column_type => format!("{column_type:?}").to_lowercase(),
}
}
fn column_types_from_metadata(metadata: &tiberius::ResultMetadata) -> Vec<String> {
metadata.columns().iter().map(sqlserver_column_type_name).collect()
}
const SQLSERVER_MESSAGE_COLUMN: &str = "Message";
// Tiberius 0.12.3 emits both user-visible INFO tokens and internal ENVCHANGE
// tokens at the same tracing target/level. Keep only the TokenInfo callsite;
// update this guard together with the pinned driver when Tiberius is upgraded.
const TIBERIUS_INFO_TOKEN_EVENT_LINE: u32 = 194;
#[derive(Clone, Default)]
struct SqlServerMessageLayer {
messages: StdArc<StdMutex<Vec<String>>>,
}
impl<S> Layer<S> for SqlServerMessageLayer
where
S: Subscriber,
{
fn on_event(&self, event: &Event<'_>, _ctx: LayerContext<'_, S>) {
let metadata = event.metadata();
if metadata.level() != &Level::INFO
|| metadata.target() != "tiberius::tds::stream::token"
|| metadata.line() != Some(TIBERIUS_INFO_TOKEN_EVENT_LINE)
{
return;
}
let mut visitor = SqlServerMessageVisitor::default();
event.record(&mut visitor);
if let Some(message) = visitor.message.filter(|message| !message.trim().is_empty()) {
if let Ok(mut messages) = self.messages.lock() {
messages.push(message);
}
}
}
}
#[derive(Default)]
struct SqlServerMessageVisitor {
message: Option<String>,
}
impl tracing::field::Visit for SqlServerMessageVisitor {
fn record_debug(&mut self, field: &tracing::field::Field, value: &dyn std::fmt::Debug) {
if field.name() == "message" {
self.message = Some(format!("{value:?}").trim_matches('"').to_string());
}
}
}
async fn capture_sqlserver_messages<F, T>(future: F) -> (T, Vec<String>)
where
F: Future<Output = T>,
{
let layer = SqlServerMessageLayer::default();
let messages = layer.messages.clone();
// Tiberius consumes TDS INFO tokens internally and exposes them only as
// tracing events. Scope collection to this future to isolate connections.
let output = future.with_subscriber(tracing_subscriber::registry().with(layer)).await;
let messages = messages.lock().map(|messages| messages.clone()).unwrap_or_default();
(output, messages)
}
fn query_result_with_server_messages(mut result: QueryResult, messages: Vec<String>) -> QueryResult {
if messages.is_empty() || !result.columns.is_empty() || !result.rows.is_empty() {
return result;
}
result.columns = vec![SQLSERVER_MESSAGE_COLUMN.to_string()];
result.column_types = vec!["nvarchar".to_string()];
result.rows = messages.into_iter().map(|message| vec![serde_json::Value::String(message)]).collect();
result
}
fn server_messages_query_result(messages: Vec<String>, start: Instant) -> Option<QueryResult> {
if messages.is_empty() {
return None;
}
Some(query_result_with_server_messages(
QueryResult {
columns: vec![],
column_types: vec![],
column_sortables: vec![],
rows: vec![],
affected_rows: 0,
execution_time_ms: start.elapsed().as_millis(),
truncated: false,
session_id: None,
has_more: false,
elasticsearch_raw_body: None,
},
messages,
))
}
async fn collect_first_result_limited(
mut stream: QueryStream<'_>,
start: Instant,
max_rows: Option<usize>,
) -> Result<QueryResult, String> {
let row_limit = query_result_row_limit(max_rows);
let mut columns: Vec<String> = vec![];
let mut column_types: Vec<String> = vec![];
let mut rows: Vec<Vec<serde_json::Value>> = Vec::new();
let mut truncated = false;
while let Some(item) = stream.try_next().await.map_err(|e| e.to_string())? {
match item {
QueryItem::Metadata(metadata) if metadata.result_index() == 0 => {
columns = columns_from_metadata(&metadata);
column_types = column_types_from_metadata(&metadata);
}
QueryItem::Metadata(_) => {}
QueryItem::Row(row) if row.result_index() == 0 => {
if rows.len() < row_limit {
rows.push(row_to_json(&row));
} else {
truncated = true;
}
}
QueryItem::Row(_) => {}
}
}
Ok(QueryResult {
columns,
column_types,
column_sortables: vec![],
rows,
affected_rows: 0,
execution_time_ms: start.elapsed().as_millis(),
truncated,
session_id: None,
has_more: false,
elasticsearch_raw_body: None,
})
}
struct SqlServerResultSet {
columns: Vec<String>,
column_types: Vec<String>,
rows: Vec<Vec<serde_json::Value>>,
truncated: bool,
}
pub struct SqlServerStreamExportSummary {
pub columns: Vec<String>,
pub rows_exported: u64,
}
pub enum SqlServerStreamItem<'a> {
Columns { columns: &'a [String], column_types: &'a [String] },
Row(&'a [serde_json::Value]),
}
#[derive(Debug, Clone, PartialEq, Eq)]
struct SqlServerDescribedColumn {
name: Option<String>,
system_type_name: Option<String>,
user_type_schema: Option<String>,
user_type_name: Option<String>,
}
#[derive(Debug, PartialEq, Eq)]
struct SqlServerLegacyProbe {
source_sql: String,
output_names: Option<Vec<Option<String>>>,
}
async fn sqlserver_driver_result<T, E, F>(future: F) -> Result<T, String>
where
F: Future<Output = Result<T, E>>,
E: ToString,
{
match AssertUnwindSafe(future).catch_unwind().await {
Ok(result) => result.map_err(|e| e.to_string()),
Err(_) => Err(format!(
"{SQLSERVER_DRIVER_PANIC_ERROR_PREFIX} the current client will be rebuilt. \
Unsupported columns may need to be cast to text."
)),
}
}
pub fn is_driver_panic_error(error: &str) -> bool {
error.starts_with(SQLSERVER_DRIVER_PANIC_ERROR_PREFIX)
}
async fn describe_sqlserver_result_set(
client: &mut SqlServerClient,
sql: &str,
legacy_probe: &SqlServerLegacyProbe,
) -> Result<Vec<SqlServerDescribedColumn>, String> {
describe_sqlserver_result_set_with_mode(client, sql, legacy_probe, false).await
}
async fn describe_sqlserver_result_set_with_mode(
client: &mut SqlServerClient,
sql: &str,
legacy_probe: &SqlServerLegacyProbe,
force_legacy: bool,
) -> Result<Vec<SqlServerDescribedColumn>, String> {
// SQL Server 2008 has no first-result-set DMV. Keep one probe round trip by
// selecting the modern DMV path server-side and using metadata-only execution otherwise.
let force_legacy = i32::from(force_legacy);
let mut stream = sqlserver_driver_result(
client.query(SQLSERVER_RESULT_TYPE_PROBE_SQL, &[&sql, &force_legacy, &legacy_probe.source_sql]),
)
.await?;
let mut active_result_index = None;
let mut uses_describe_dmv = None;
let mut rows = Vec::new();
loop {
let item = match sqlserver_driver_result(stream.try_next()).await {
Ok(item) => item,
Err(error) if uses_describe_dmv == Some(false) => {
return Err(format!("{SQLSERVER_UNSAFE_PROBE_BLOCK_ERROR_PREFIX} {error}"));
}
Err(error) => return Err(error),
};
let Some(item) = item else {
break;
};
match item {
QueryItem::Metadata(result_metadata) => {
active_result_index = Some(result_metadata.result_index());
}
QueryItem::Row(row) if active_result_index == Some(0) => {
uses_describe_dmv = row.try_get::<bool, _>(0).ok().flatten();
}
QueryItem::Row(row) => rows.push(row),
}
}
if uses_describe_dmv.is_none() {
return Err("SQL Server result type probe did not report its compatibility mode".to_string());
}
let mut columns = rows.iter().map(sqlserver_described_column_from_row).collect::<Vec<_>>();
if uses_describe_dmv == Some(false) {
if let Some(output_names) = &legacy_probe.output_names {
for (column, output_name) in columns.iter_mut().zip(output_names) {
column.name.clone_from(output_name);
}
}
}
Ok(columns)
}
fn sqlserver_described_column_from_row(row: &Row) -> SqlServerDescribedColumn {
SqlServerDescribedColumn {
name: row.try_get::<&str, _>(0).ok().flatten().map(str::to_string),
system_type_name: row.try_get::<&str, _>(1).ok().flatten().map(str::to_string),
user_type_schema: row.try_get::<&str, _>(2).ok().flatten().map(str::to_string),
user_type_name: row.try_get::<&str, _>(3).ok().flatten().map(str::to_string),
}
}
fn is_blocking_sqlserver_unsafe_probe_error(error: &str) -> bool {
error.starts_with(SQLSERVER_UNSAFE_PROBE_BLOCK_ERROR_PREFIX)
}
async fn sqlserver_unsafe_type_query(client: &mut SqlServerClient, sql: &str) -> Result<Option<String>, String> {
if !is_single_sqlserver_select(sql) {
return Ok(None);
}
let Some(legacy_probe) = sqlserver_legacy_probe(sql) else {
return Ok(None);
};
let columns = describe_sqlserver_result_set(client, sql, &legacy_probe).await?;
Ok(build_sqlserver_unsafe_type_query(sql, &columns))
}
fn sqlserver_legacy_probe(sql: &str) -> Option<SqlServerLegacyProbe> {
let statement = normalized_sqlserver_select_statement(sql)?;
let output_names = sqlserver_projection_output_names(&statement);
let source_alias = quote_sqlserver_identifier("dbx_probe_source");
let source_sql = if let Some(names) = &output_names {
let aliases = (0..names.len())
.map(sqlserver_source_column_name)
.map(|name| quote_sqlserver_identifier(&name))
.collect::<Vec<_>>()
.join(", ");
format!("({statement}) AS {source_alias}({aliases})")
} else {
format!("({statement}) AS {source_alias}")
};
Some(SqlServerLegacyProbe { source_sql, output_names })
}
fn sqlserver_projection_output_names(statement: &str) -> Option<Vec<Option<String>>> {
let statements = Parser::parse_sql(&MsSqlDialect {}, statement).ok()?;
let [Statement::Query(query)] = statements.as_slice() else {
return None;
};
let SetExpr::Select(select) = query.body.as_ref() else {
return None;
};
select
.projection
.iter()
.map(|item| match item {
SelectItem::ExprWithAlias { alias, .. } => Some(Some(alias.value.clone())),
SelectItem::UnnamedExpr(Expr::Identifier(identifier)) => {
Some((!identifier.value.starts_with('@')).then(|| identifier.value.clone()))
}
SelectItem::UnnamedExpr(Expr::CompoundIdentifier(identifiers)) => {
Some(identifiers.last().map(|identifier| identifier.value.clone()))
}
SelectItem::UnnamedExpr(_) => Some(None),
SelectItem::ExprWithAliases { .. } | SelectItem::QualifiedWildcard(_, _) | SelectItem::Wildcard(_) => None,
})
.collect()
}
fn build_sqlserver_unsafe_type_query(sql: &str, columns: &[SqlServerDescribedColumn]) -> Option<String> {
if columns.is_empty() || !columns.iter().any(is_sqlserver_unsafe_column) {
return None;
}
let statement = normalized_sqlserver_select_statement(sql)?;
let source_alias = quote_sqlserver_identifier("dbx_unsafe_source");
let source_columns = (0..columns.len()).map(sqlserver_source_column_name).collect::<Vec<_>>();
let source_alias_list =
source_columns.iter().map(|name| quote_sqlserver_identifier(name)).collect::<Vec<_>>().join(", ");
let select_list = columns
.iter()
.enumerate()
.map(|(index, column)| {
let output_name = sqlserver_output_column_name(column, index);
let quoted_output = quote_sqlserver_identifier(&output_name);
let source_column = quote_sqlserver_identifier(&source_columns[index]);
let value_ref = format!("{source_alias}.{source_column}");
if is_sqlserver_spatial_column(column) {
format!("{quoted_output} = CASE WHEN {value_ref} IS NULL THEN NULL ELSE {value_ref}.STAsText() END")
} else if is_sqlserver_variant_column(column) {
format!("{quoted_output} = CAST({value_ref} AS NVARCHAR(MAX))")
} else {
format!("{quoted_output} = {value_ref}")
}
})
.collect::<Vec<_>>()
.join(", ");
Some(format!("SELECT {select_list} FROM ({statement}) AS {source_alias}({source_alias_list})"))
}
fn is_sqlserver_unsafe_column(column: &SqlServerDescribedColumn) -> bool {
is_sqlserver_spatial_column(column) || is_sqlserver_variant_column(column)
}
fn is_sqlserver_spatial_column(column: &SqlServerDescribedColumn) -> bool {
[&column.system_type_name, &column.user_type_name].into_iter().flatten().any(|name| {
let normalized = name.trim().trim_matches(['[', ']']).to_ascii_lowercase();
normalized == "geometry"
|| normalized == "geography"
|| normalized.ends_with(".geometry")
|| normalized.ends_with(".geography")
})
}
fn is_sqlserver_variant_column(column: &SqlServerDescribedColumn) -> bool {
[&column.system_type_name, &column.user_type_name].into_iter().flatten().any(|name| {
let normalized = name.trim().trim_matches(['[', ']']).to_ascii_lowercase();
normalized == "sql_variant" || normalized.ends_with(".sql_variant")
})
}
fn normalized_sqlserver_select_statement(sql: &str) -> Option<String> {
let statement = trim_sqlserver_statement(sql);
let trimmed = statement.trim_start();
if trimmed.is_empty() || !trimmed.get(..6).is_some_and(|prefix| prefix.eq_ignore_ascii_case("SELECT")) {
return None;
}
if has_top_level_select_into(trimmed) {
return None;
}
// Strip trailing ORDER BY so the statement can be used as a derived table
// subquery. SQL Server requires TOP / OFFSET / FOR XML alongside ORDER BY
// in subqueries, none of which are version-safe across 2008–2022.
let mut statement = trimmed.to_string();
let tokens = top_level_sqlserver_tokens(&statement);
for index in (0..tokens.len().saturating_sub(1)).rev() {
if tokens[index].text == "ORDER" && tokens.get(index + 1).is_some_and(|token| token.text == "BY") {
statement.truncate(tokens[index].start);
statement = statement.trim_end().to_string();
break;
}
}
Some(statement)
}
fn trim_sqlserver_statement(sql: &str) -> String {
let mut statement = sql.trim();
while let Some(stripped) = statement.strip_suffix(';') {
statement = stripped.trim_end();
}
statement.to_string()
}
fn is_single_sqlserver_select(sql: &str) -> bool {
let statements = crate::sql::split_sql_statements(sql);
if statements.len() != 1 {
return false;
}
let statement = statements[0].trim_start();
statement.get(..6).is_some_and(|prefix| prefix.eq_ignore_ascii_case("SELECT"))
}
fn sqlserver_source_column_name(index: usize) -> String {
format!("dbx_col_{}", index + 1)
}
fn sqlserver_output_column_name(column: &SqlServerDescribedColumn, index: usize) -> String {
column
.name
.as_deref()
.map(str::trim)
.filter(|name| !name.is_empty())
.map(str::to_string)
.unwrap_or_else(|| format!("column_{}", index + 1))
}
fn quote_sqlserver_identifier(identifier: &str) -> String {
format!("[{}]", identifier.replace(']', "]]"))
}
#[derive(Debug, Clone, PartialEq, Eq)]
struct SqlServerToken {
text: String,
start: usize,
}
fn top_level_sqlserver_tokens(sql: &str) -> Vec<SqlServerToken> {
let mut tokens = Vec::new();
let mut i = 0;
let mut depth = 0usize;
while i < sql.len() {
let ch = next_char(sql, i);
let next = next_char_at(sql, i + ch.len_utf8());
if ch == '-' && next == Some('-') {
i += 2;
while i < sql.len() && next_char(sql, i) != '\n' {
i += next_char(sql, i).len_utf8();
}
continue;
}
if ch == '/' && next == Some('*') {
i += 2;
while i < sql.len() {
let current = next_char(sql, i);
let following = next_char_at(sql, i + current.len_utf8());
if current == '*' && following == Some('/') {
i += 2;
break;
}
i += current.len_utf8();
}
continue;
}
if matches!(ch, '\'' | '"') {
i = skip_sqlserver_quoted(sql, i, ch);
continue;
}
if ch == '[' {
i = skip_sqlserver_bracket_identifier(sql, i);
continue;
}
if ch == '(' {
depth += 1;
i += ch.len_utf8();
continue;
}
if ch == ')' {
depth = depth.saturating_sub(1);
i += ch.len_utf8();
continue;
}
if depth == 0 && is_sqlserver_token_start(ch) {
let start = i;
i += ch.len_utf8();
while i < sql.len() && is_sqlserver_token_part(next_char(sql, i)) {
i += next_char(sql, i).len_utf8();
}
tokens.push(SqlServerToken { text: sql[start..i].to_ascii_uppercase(), start });
continue;
}
i += ch.len_utf8();
}
tokens
}
fn has_top_level_select_into(sql: &str) -> bool {
let tokens = top_level_sqlserver_tokens(sql);
let Some(select_index) = tokens.iter().position(|token| token.text == "SELECT") else {
return false;
};
let from_index = tokens
.iter()
.enumerate()
.find(|(index, token)| *index > select_index && token.text == "FROM")
.map(|(index, _)| index)
.unwrap_or(tokens.len());
tokens[select_index + 1..from_index].iter().any(|token| token.text == "INTO")
}
fn skip_sqlserver_quoted(sql: &str, pos: usize, quote: char) -> usize {
let mut i = pos + quote.len_utf8();
while i < sql.len() {
let ch = next_char(sql, i);
let next = next_char_at(sql, i + ch.len_utf8());
if ch == quote {
if next == Some(quote) {
i += ch.len_utf8() + quote.len_utf8();
continue;
}
return i + ch.len_utf8();
}
i += ch.len_utf8();
}
sql.len()
}
fn skip_sqlserver_bracket_identifier(sql: &str, pos: usize) -> usize {
let mut i = pos + 1;
while i < sql.len() {
let ch = next_char(sql, i);
let next = next_char_at(sql, i + ch.len_utf8());
if ch == ']' {
if next == Some(']') {
i += 2;
continue;
}
return i + 1;
}
i += ch.len_utf8();
}
sql.len()
}
fn is_sqlserver_token_start(ch: char) -> bool {
ch.is_ascii_alphabetic() || ch == '_'
}
fn is_sqlserver_token_part(ch: char) -> bool {
ch.is_ascii_alphanumeric() || matches!(ch, '_' | '$' | '#')
}
fn next_char(sql: &str, index: usize) -> char {
sql[index..].chars().next().unwrap_or('\0')
}
fn next_char_at(sql: &str, index: usize) -> Option<char> {
if index >= sql.len() {
None
} else {
sql[index..].chars().next()
}
}
fn push_sqlserver_result_set(results: &mut Vec<QueryResult>, result: Option<SqlServerResultSet>, start: Instant) {
if let Some(result) = result {
if result.rows.is_empty() && result.columns.is_empty() {
return;
}
results.push(QueryResult {
columns: result.columns,
column_types: result.column_types,
column_sortables: vec![],
rows: result.rows,
affected_rows: 0,
execution_time_ms: start.elapsed().as_millis(),
truncated: result.truncated,
session_id: None,
has_more: false,
elasticsearch_raw_body: None,
});
}
}
async fn collect_result_sets_limited(
mut stream: QueryStream<'_>,
start: Instant,
max_rows: Option<usize>,
) -> Result<Vec<QueryResult>, String> {
let row_limit = query_result_row_limit(max_rows);
let mut results = Vec::new();
let mut current: Option<SqlServerResultSet> = None;
while let Some(item) = stream.try_next().await.map_err(|e| e.to_string())? {
match item {
QueryItem::Metadata(metadata) => {
push_sqlserver_result_set(&mut results, current.take(), start);
current = Some(SqlServerResultSet {
columns: columns_from_metadata(&metadata),
column_types: column_types_from_metadata(&metadata),
rows: Vec::new(),
truncated: false,
});
}
QueryItem::Row(row) => {
let result = current.get_or_insert_with(|| SqlServerResultSet {
columns: row.columns().iter().map(|c| c.name().to_string()).collect(),
column_types: row.columns().iter().map(sqlserver_column_type_name).collect(),
rows: Vec::new(),
truncated: false,
});
if result.rows.len() < row_limit {
result.rows.push(row_to_json(&row));
} else {
result.truncated = true;
}
}
}
}
push_sqlserver_result_set(&mut results, current, start);
Ok(results)
}
pub async fn stream_first_result_set(
client: &mut SqlServerClient,
sql: &str,
row_limit: Option<usize>,
cancel_token: Option<CancellationToken>,
mut on_item: impl for<'a> FnMut(SqlServerStreamItem<'a>) -> Result<(), String>,
) -> Result<SqlServerStreamExportSummary, String> {
let query_sql = match sqlserver_unsafe_type_query(client, sql).await {
Ok(Some(sql)) => sql,
Ok(None) => sql.to_string(),
Err(error) if is_blocking_sqlserver_unsafe_probe_error(&error) => return Err(error),
Err(_) => sql.to_string(),
};
let mut stream = sqlserver_driver_result(client.query(query_sql.as_str(), &[])).await?;
let mut active_result_index: Option<usize> = None;
let mut columns: Vec<String> = Vec::new();
let mut column_types: Vec<String> = Vec::new();
let mut columns_emitted = false;
let mut rows_exported = 0_u64;
loop {
if cancel_token.as_ref().is_some_and(|token| token.is_cancelled()) {
return Err(crate::query::canceled_error());
}
let item = match cancel_token.as_ref() {
Some(token) => {
tokio::select! {
biased;
_ = token.cancelled() => return Err(crate::query::canceled_error()),
item = stream.try_next() => item.map_err(|e| e.to_string())?,
}
}
None => stream.try_next().await.map_err(|e| e.to_string())?,
};
let Some(item) = item else {
break;
};
match item {
QueryItem::Metadata(metadata) => {
if active_result_index.is_none() {
active_result_index = Some(metadata.result_index());
columns = columns_from_metadata(&metadata);
column_types = column_types_from_metadata(&metadata);
on_item(SqlServerStreamItem::Columns { columns: &columns, column_types: &column_types })?;
columns_emitted = true;
}
}
QueryItem::Row(row) => {
if active_result_index.is_none() {
active_result_index = Some(row.result_index());
columns = row.columns().iter().map(|c| c.name().to_string()).collect();
column_types = row.columns().iter().map(sqlserver_column_type_name).collect();
on_item(SqlServerStreamItem::Columns { columns: &columns, column_types: &column_types })?;
columns_emitted = true;
}
if Some(row.result_index()) != active_result_index {
continue;
}
if row_limit.is_some_and(|limit| rows_exported as usize >= limit) {
break;
}
let values = row_to_json(&row);
on_item(SqlServerStreamItem::Row(&values))?;
rows_exported += 1;
}
}
}
if !columns_emitted {
on_item(SqlServerStreamItem::Columns { columns: &columns, column_types: &column_types })?;
}
Ok(SqlServerStreamExportSummary { columns, rows_exported })
}
// rust_decimal (behind tiberius's `Decimal: FromSql`) only supports scale <= 28 and a
// 96-bit mantissa, while SQL Server NUMERIC allows precision/scale up to 38; converting
// such values panics and aborts the app (issue #3648). Format the raw i128 value and
// scale manually instead.
fn format_sqlserver_numeric(value: i128, scale: u8) -> String {
if scale == 0 {
return value.to_string();
}
let digits = value.unsigned_abs().to_string();
let scale = scale as usize;
let sign = if value < 0 { "-" } else { "" };
if digits.len() > scale {
let (int_part, frac_part) = digits.split_at(digits.len() - scale);
format!("{}{}.{}", sign, int_part, frac_part)
} else {
format!("{}0.{:0>width$}", sign, digits, width = scale)
}
}
fn sqlserver_cell_to_json(cell: &ColumnData<'static>) -> serde_json::Value {
if let ColumnData::Numeric(numeric) = cell {
return match numeric {
Some(n) => serde_json::Value::String(format_sqlserver_numeric(n.value(), n.scale())),
None => serde_json::Value::Null,
};
}
if let Ok(Some(v)) = <&tiberius::xml::XmlData as FromSql>::from_sql(cell) {
return serde_json::Value::String(v.as_ref().to_string());
}
if let Ok(Some(v)) = <&str as FromSql>::from_sql(cell) {
return serde_json::Value::String(v.to_string());
}
if let Ok(Some(v)) = <chrono::NaiveDateTime as FromSql>::from_sql(cell) {
let value = match cell {
ColumnData::DateTime(_) => crate::sqlserver_temporal::format_sqlserver_datetime_display(&v),
_ => v.to_string(),
};
return serde_json::Value::String(value);
}
if let Ok(Some(v)) = <chrono::NaiveDate as FromSql>::from_sql(cell) {
return serde_json::Value::String(v.to_string());
}
if let Ok(Some(v)) = <chrono::NaiveTime as FromSql>::from_sql(cell) {
return serde_json::Value::String(v.to_string());
}
if let Ok(Some(v)) = <chrono::DateTime<chrono::FixedOffset> as FromSql>::from_sql(cell) {
return serde_json::Value::String(v.to_rfc3339());
}
if let Ok(Some(v)) = <u8 as FromSql>::from_sql(cell) {
return serde_json::Value::Number(v.into());
}
if let Ok(Some(v)) = <i16 as FromSql>::from_sql(cell) {
return serde_json::Value::Number(v.into());
}
if let Ok(Some(v)) = <i32 as FromSql>::from_sql(cell) {
return serde_json::Value::Number(v.into());
}
if let Ok(Some(v)) = <i64 as FromSql>::from_sql(cell) {
return super::safe_i64_to_json(v);
}
if let Ok(Some(v)) = <f32 as FromSql>::from_sql(cell) {
return serde_json::Number::from_f64(v as f64)
.map(serde_json::Value::Number)
.unwrap_or(serde_json::Value::Null);
}
if let Ok(Some(v)) = <f64 as FromSql>::from_sql(cell) {
return serde_json::Number::from_f64(v).map(serde_json::Value::Number).unwrap_or(serde_json::Value::Null);
}
if let Ok(Some(v)) = <bool as FromSql>::from_sql(cell) {
return serde_json::Value::Bool(v);
}
if let Ok(Some(v)) = <uuid::Uuid as FromSql>::from_sql(cell) {
return serde_json::Value::String(v.to_string());
}
if let Ok(Some(v)) = <Vec<u8> as tiberius::FromSqlOwned>::from_sql_owned(cell.clone()) {
return super::binary_value_to_json(&v);
}
serde_json::Value::Null
}
pub async fn list_databases(client: &mut SqlServerClient) -> Result<Vec<DatabaseInfo>, String> {
let stream = client
.query(
"SELECT name \
FROM sys.databases \
WHERE state = 0 \
ORDER BY name",
&[],
)
.await
.map_err(|e| e.to_string())?;
let rows = stream.into_first_result().await.map_err(|e| e.to_string())?;
Ok(rows.iter().map(|row| DatabaseInfo { name: row.get::<&str, _>(0).unwrap_or("").to_string() }).collect())
}
pub async fn test_connection(client: &mut SqlServerClient) -> Result<(), String> {
crate::db::with_connection_timeout("SQL Server", crate::db::connection_timeout(), async {
let stream = client.simple_query("SELECT 1").await.map_err(|e| e.to_string())?;
let _ = stream.into_first_result().await.map_err(|e| e.to_string())?;
Ok(())
})
.await
}
pub async fn list_linked_servers(client: &mut SqlServerClient) -> Result<Vec<LinkedServerInfo>, String> {
let stream = client
.query(
"SELECT name, product, provider, data_source \
FROM sys.servers \
WHERE is_linked = 1 \
ORDER BY name",
&[],
)
.await
.map_err(|e| e.to_string())?;
let rows = stream.into_first_result().await.map_err(|e| e.to_string())?;
Ok(rows
.iter()
.map(|row| LinkedServerInfo {
name: row.get::<&str, _>(0).unwrap_or("").to_string(),
product: row.get::<&str, _>(1).filter(|value| !value.trim().is_empty()).map(str::to_string),
provider: row.get::<&str, _>(2).filter(|value| !value.trim().is_empty()).map(str::to_string),
data_source: row.get::<&str, _>(3).filter(|value| !value.trim().is_empty()).map(str::to_string),
})
.filter(|server| !server.name.trim().is_empty())
.collect())
}
pub async fn list_linked_server_catalogs(
client: &mut SqlServerClient,
server: &str,
) -> Result<Vec<DatabaseInfo>, String> {
let stream = client.query("EXEC sp_catalogs @server_name = @P1", &[&server]).await.map_err(|e| e.to_string())?;
let rows = stream.into_first_result().await.map_err(|e| e.to_string())?;
Ok(rows
.iter()
.filter_map(|row| row.get::<&str, _>(0).map(str::trim).filter(|name| !name.is_empty()))
.map(|name| DatabaseInfo { name: name.to_string() })
.collect())
}
pub async fn list_linked_server_schemas(
client: &mut SqlServerClient,
server: &str,
catalog: &str,
) -> Result<Vec<String>, String> {
let tables = linked_server_table_rows(client, server, catalog, None, None).await?;
let mut schemas = Vec::new();
for table in tables {
if let Some(schema) = table.schema.filter(|value| !value.trim().is_empty()) {
if !schemas.iter().any(|existing: &String| existing.eq_ignore_ascii_case(&schema)) {
schemas.push(schema);
}
}
}
schemas.sort_by_key(|schema| (if schema.eq_ignore_ascii_case("dbo") { 0 } else { 1 }, schema.to_lowercase()));
Ok(schemas)
}
pub async fn list_linked_server_tables(
client: &mut SqlServerClient,
server: &str,
catalog: &str,
schema: &str,
filter: Option<&str>,
limit: Option<usize>,
offset: Option<usize>,
) -> Result<Vec<TableInfo>, String> {
let filter = filter.map(str::trim).filter(|value| !value.is_empty()).map(str::to_lowercase);
let limit = limit.unwrap_or(usize::MAX);
let offset = offset.unwrap_or(0);
let rows = linked_server_table_rows(client, server, catalog, Some(schema), None).await?;
Ok(rows
.into_iter()
.filter(|row| filter.as_ref().is_none_or(|value| row.name.to_lowercase().contains(value)))
.skip(offset)
.take(limit)
.map(|row| TableInfo {
name: row.name,
table_type: normalize_linked_server_table_type(row.table_type.as_deref()),
comment: row.comment,
parent_schema: None,
parent_name: None,
})
.collect())
}
pub async fn get_linked_server_columns(
client: &mut SqlServerClient,
server: &str,
catalog: &str,
schema: &str,
table: &str,
) -> Result<Vec<ColumnInfo>, String> {
let stream = client
.query(
"EXEC sp_columns_ex \
@table_server = @P1, \
@table_name = @P2, \
@table_schema = @P3, \
@table_catalog = @P4",
&[&server, &table, &schema, &catalog],
)
.await
.map_err(|e| e.to_string())?;
let rows = stream.into_first_result().await.map_err(|e| e.to_string())?;
Ok(rows
.iter()
.filter_map(|row| {
let name = row.get::<&str, _>(3)?.trim();
if name.is_empty() {
return None;
}
let base_type = row.get::<&str, _>(5).unwrap_or("").trim();
let column_size = linked_i32(row, 6);
let numeric_scale = linked_i32(row, 8);
let nullable = linked_i32(row, 10).unwrap_or(1) != 0;
let data_type = linked_server_column_type(base_type, column_size, numeric_scale);
Some(ColumnInfo {
name: name.to_string(),
data_type,
is_nullable: nullable,
column_default: row.get::<&str, _>(12).filter(|value| !value.trim().is_empty()).map(str::to_string),
is_primary_key: false,
extra: None,
comment: row.get::<&str, _>(11).filter(|value| !value.trim().is_empty()).map(str::to_string),
numeric_precision: column_size,
numeric_scale,
character_maximum_length: linked_i32(row, 15),
enum_values: None,
..Default::default()
})
})
.collect())
}
struct LinkedServerTableRow {
schema: Option<String>,
name: String,
table_type: Option<String>,
comment: Option<String>,
}
async fn linked_server_table_rows(
client: &mut SqlServerClient,
server: &str,
catalog: &str,
schema: Option<&str>,
table_name: Option<&str>,
) -> Result<Vec<LinkedServerTableRow>, String> {
let sql = format!(
"EXEC sp_tables_ex \
@table_server = {}, \
@table_name = {}, \
@table_schema = {}, \
@table_catalog = {}, \
@table_type = '''TABLE'',''VIEW''', \
@fUsePattern = 0",
sqlserver_nstring_literal(server),
sqlserver_optional_nstring_literal(table_name),
sqlserver_optional_nstring_literal(schema),
sqlserver_nstring_literal(catalog),
);
let stream = client.query(sql.as_str(), &[]).await.map_err(|e| e.to_string())?;
let rows = stream.into_first_result().await.map_err(|e| e.to_string())?;
Ok(rows
.iter()
.filter_map(|row| {
let name = row.get::<&str, _>(2)?.trim();
if name.is_empty() {
return None;
}
Some(LinkedServerTableRow {
schema: row.get::<&str, _>(1).filter(|value| !value.trim().is_empty()).map(str::to_string),
name: name.to_string(),
table_type: row.get::<&str, _>(3).filter(|value| !value.trim().is_empty()).map(str::to_string),
comment: row.get::<&str, _>(4).filter(|value| !value.trim().is_empty()).map(str::to_string),
})
})
.collect())
}
fn sqlserver_optional_nstring_literal(value: Option<&str>) -> String {
value.filter(|value| !value.trim().is_empty()).map(sqlserver_nstring_literal).unwrap_or_else(|| "NULL".to_string())
}
fn sqlserver_nstring_literal(value: &str) -> String {
format!("N'{}'", value.replace('\'', "''"))
}
fn normalize_linked_server_table_type(value: Option<&str>) -> String {
let upper = value.unwrap_or("TABLE").to_ascii_uppercase();
if upper.contains("VIEW") {
"VIEW".to_string()
} else {
"BASE TABLE".to_string()
}
}
fn linked_server_column_type(base_type: &str, size: Option<i32>, scale: Option<i32>) -> String {
let lower = base_type.to_ascii_lowercase();
if matches!(lower.as_str(), "varchar" | "nvarchar" | "char" | "nchar" | "binary" | "varbinary") {
if let Some(size) = size {
if size > 0 {
return format!("{base_type}({size})");
}
}
}
if matches!(lower.as_str(), "decimal" | "numeric") {
if let (Some(size), Some(scale)) = (size, scale) {
return format!("{base_type}({size},{scale})");
}
}
base_type.to_string()
}
fn linked_i32(row: &tiberius::Row, index: usize) -> Option<i32> {
row.try_get::<i32, _>(index).ok().flatten().or_else(|| row.try_get::<i16, _>(index).ok().flatten().map(i32::from))
}
pub async fn list_schemas(client: &mut SqlServerClient) -> Result<Vec<String>, String> {
let sql = sqlserver_list_schemas_sql();
let stream = client.query(&*sql, &[]).await.map_err(|e| e.to_string())?;
let rows = stream.into_first_result().await.map_err(|e| e.to_string())?;
Ok(rows.iter().map(|row| row.get::<&str, _>(0).unwrap_or("").to_string()).collect())
}
fn sqlserver_list_schemas_sql() -> String {
let excluded_schemas =
sqlserver_hidden_schema_names().iter().map(|name| format!("'{name}'")).collect::<Vec<_>>().join(",");
format!(
"SELECT s.name \
FROM sys.schemas s \
WHERE s.name NOT IN ({excluded_schemas}) \
ORDER BY CASE WHEN s.name = 'dbo' THEN 0 ELSE 1 END, s.name"
)
}
pub async fn list_tables(
client: &mut SqlServerClient,
schema: &str,
filter: Option<&str>,
limit: Option<usize>,
offset: Option<usize>,
) -> Result<Vec<TableInfo>, String> {
let sql = sqlserver_list_tables_sql(schema, filter, limit, offset);
let stream = client.query(&*sql, &[]).await.map_err(|e| e.to_string())?;
let rows = stream.into_first_result().await.map_err(|e| e.to_string())?;
Ok(rows
.iter()
.map(|row| TableInfo {
name: row.get::<&str, _>(0).unwrap_or("").to_string(),
table_type: row.get::<&str, _>(1).unwrap_or("BASE TABLE").to_string(),
comment: row.get::<&str, _>(2).filter(|s: &&str| !s.is_empty()).map(|s: &str| s.to_string()),
parent_schema: None,
parent_name: None,
})
.collect())
}
pub async fn completion_assistant_search(
client: &mut SqlServerClient,
request: &crate::types::CompletionAssistantRequest,
) -> Result<crate::types::CompletionAssistantResponse, String> {
let limit = request.max_results.unwrap_or(100).clamp(1, 1000);
let sql = sqlserver_completion_assistant_sql(request, limit);
let stream = client.query(&*sql, &[]).await.map_err(|e| e.to_string())?;
let rows = stream.into_first_result().await.map_err(|e| e.to_string())?;
let candidates = rows
.iter()
.map(|row| {
let object_type = row.get::<&str, _>(2).unwrap_or("OBJECT");
crate::types::CompletionAssistantCandidate {
name: row.get::<&str, _>(0).unwrap_or("").to_string(),
kind: sqlserver_completion_candidate_kind(object_type),
database: Some(request.database.clone()),
schema: row.get::<&str, _>(1).map(str::to_string),
parent_schema: row.get::<&str, _>(3).map(str::to_string),
parent_name: row.get::<&str, _>(4).map(str::to_string),
comment: row.get::<&str, _>(5).filter(|s: &&str| !s.is_empty()).map(|s| (*s).to_string()),
data_type: row.get::<&str, _>(6).map(str::to_string),
}
})
.collect::<Vec<_>>();
Ok(crate::types::CompletionAssistantResponse {
incomplete: candidates.len() >= limit,
candidates,
fallback_used: false,
})
}
fn sqlserver_completion_candidate_kind(object_type: &str) -> crate::types::CompletionAssistantCandidateKind {
match object_type.to_ascii_uppercase().as_str() {
"SCHEMA" => crate::types::CompletionAssistantCandidateKind::Schema,
"TABLE" | "BASE TABLE" => crate::types::CompletionAssistantCandidateKind::Table,
"VIEW" => crate::types::CompletionAssistantCandidateKind::View,
"PROCEDURE" => crate::types::CompletionAssistantCandidateKind::Procedure,
"FUNCTION" => crate::types::CompletionAssistantCandidateKind::Function,
"COLUMN" => crate::types::CompletionAssistantCandidateKind::Column,
_ => crate::types::CompletionAssistantCandidateKind::Object,
}
}
fn sqlserver_completion_assistant_sql(request: &crate::types::CompletionAssistantRequest, limit: usize) -> String {
let object_kinds = if request.object_kinds.is_empty() {
vec![crate::types::CompletionAssistantObjectKind::Table, crate::types::CompletionAssistantObjectKind::View]
} else {
request.object_kinds.clone()
};
let mask = request.mask.trim();
let like_pattern = completion_like_pattern(mask, request.match_mode.as_ref());
let like_clause = if like_pattern == "%" {
String::new()
} else {
format!(" AND LOWER({}) LIKE LOWER('{like_pattern}') ESCAPE '\\' ", "name_expr")
};
let schema_filter = request
.schema
.as_deref()
.or(request.parent_schema.as_deref())
.filter(|schema| !schema.trim().is_empty())
.map(|schema| format!(" AND s.name = '{}' ", schema.replace('\'', "''")))
.unwrap_or_default();
let mut queries = Vec::new();
if (mask.starts_with('#') || mask.starts_with("%#"))
&& object_kinds.iter().any(crate::types::CompletionAssistantObjectKind::is_table_like)
{
let object_like = sqlserver_completion_object_search_clause(request, &like_pattern);
queries.push(format!(
"SELECT TOP ({limit}) o.name, s.name AS schema_name, 'TABLE' AS object_type, CAST(NULL AS NVARCHAR(128)) AS parent_schema, CAST(NULL AS NVARCHAR(128)) AS parent_name, CAST(NULL AS NVARCHAR(MAX)) AS object_comment, CAST(NULL AS NVARCHAR(128)) AS data_type \
FROM tempdb.sys.all_objects o \
JOIN tempdb.sys.schemas s ON s.schema_id = o.schema_id \
WHERE o.type = 'U' {object_like}"
));
return format!("SELECT * FROM ({}) AS dbx_completion ORDER BY name", queries.remove(0));
}
if object_kinds.iter().any(|kind| matches!(kind, crate::types::CompletionAssistantObjectKind::Schema)) {
let schema_like = like_clause.replace("name_expr", "s.name");
queries.push(format!(
"SELECT TOP ({limit}) s.name, s.name AS schema_name, 'SCHEMA' AS object_type, CAST(NULL AS NVARCHAR(128)) AS parent_schema, CAST(NULL AS NVARCHAR(128)) AS parent_name, CAST(NULL AS NVARCHAR(MAX)) AS object_comment, CAST(NULL AS NVARCHAR(128)) AS data_type \
FROM sys.schemas s \
WHERE s.name NOT IN ('guest','INFORMATION_SCHEMA','sys') {schema_like}"
));
}
if object_kinds.iter().any(crate::types::CompletionAssistantObjectKind::is_table_like)
|| object_kinds.iter().any(crate::types::CompletionAssistantObjectKind::is_routine_like)
{
let mut type_ids = Vec::new();
if object_kinds.iter().any(|kind| matches!(kind, crate::types::CompletionAssistantObjectKind::Table)) {
type_ids.push("'U'");
}
if object_kinds.iter().any(|kind| matches!(kind, crate::types::CompletionAssistantObjectKind::View)) {
type_ids.push("'V'");
}
if object_kinds.iter().any(|kind| {
matches!(
kind,
crate::types::CompletionAssistantObjectKind::Procedure
| crate::types::CompletionAssistantObjectKind::Routine
)
}) {
type_ids.push("'P'");
}
if object_kinds.iter().any(|kind| {
matches!(
kind,
crate::types::CompletionAssistantObjectKind::Function
| crate::types::CompletionAssistantObjectKind::Routine
)
}) {
type_ids.extend(["'FN'", "'IF'", "'TF'", "'FS'", "'FT'"]);
}
let object_like = sqlserver_completion_object_search_clause(request, &like_pattern);
let object_visibility = sqlserver_visible_object_predicate();
let data_type = if object_kinds.iter().any(crate::types::CompletionAssistantObjectKind::is_routine_like) {
"CASE WHEN o.type IN ('IF','TF','FT') THEN 'table' ELSE (SELECT TOP (1) TYPE_NAME(p.user_type_id) FROM sys.parameters p WHERE p.object_id = o.object_id AND p.parameter_id = 0) END"
} else {
"CAST(NULL AS NVARCHAR(128))"
};
queries.push(format!(
"SELECT TOP ({limit}) o.name, s.name AS schema_name, \
CASE o.type WHEN 'U' THEN 'TABLE' WHEN 'V' THEN 'VIEW' WHEN 'P' THEN 'PROCEDURE' WHEN 'FN' THEN 'FUNCTION' WHEN 'IF' THEN 'FUNCTION' WHEN 'TF' THEN 'FUNCTION' WHEN 'FS' THEN 'FUNCTION' WHEN 'FT' THEN 'FUNCTION' ELSE o.type_desc END AS object_type, \
CAST(NULL AS NVARCHAR(128)) AS parent_schema, CAST(NULL AS NVARCHAR(128)) AS parent_name, ep.value AS object_comment, {data_type} AS data_type \
FROM sys.objects o \
JOIN sys.schemas s ON s.schema_id = o.schema_id \
OUTER APPLY (SELECT CAST(ep.value AS NVARCHAR(MAX)) AS value FROM sys.extended_properties ep WHERE ep.major_id = o.object_id AND ep.minor_id = 0 AND ep.name = N'MS_Description') ep \
WHERE o.type IN ({}) AND {object_visibility} {schema_filter} {object_like}",
type_ids.join(",")
));
}
if object_kinds.iter().any(|kind| matches!(kind, crate::types::CompletionAssistantObjectKind::Column)) {
let column_like = like_clause.replace("name_expr", "c.name");
let parent_table_filter = request
.parent_name
.as_deref()
.filter(|table| !table.trim().is_empty())
.map(|table| format!(" AND o.name = '{}' ", table.replace('\'', "''")))
.unwrap_or_default();
let object_visibility = sqlserver_visible_object_predicate();
queries.push(format!(
"SELECT TOP ({limit}) c.name, s.name AS schema_name, 'COLUMN' AS object_type, s.name AS parent_schema, o.name AS parent_name, CAST(NULL AS NVARCHAR(MAX)) AS object_comment, TYPE_NAME(c.user_type_id) AS data_type \
FROM sys.columns c \
JOIN sys.objects o ON o.object_id = c.object_id \
JOIN sys.schemas s ON s.schema_id = o.schema_id \
WHERE o.type IN ('U','V') AND {object_visibility} {schema_filter} {parent_table_filter} {column_like}"
));
}
if queries.is_empty() {
"SELECT TOP (0) CAST('' AS NVARCHAR(128)) AS name, CAST('' AS NVARCHAR(128)) AS schema_name, CAST('' AS NVARCHAR(60)) AS object_type, CAST(NULL AS NVARCHAR(128)) AS parent_schema, CAST(NULL AS NVARCHAR(128)) AS parent_name, CAST(NULL AS NVARCHAR(MAX)) AS object_comment, CAST(NULL AS NVARCHAR(128)) AS data_type".to_string()
} else if queries.len() == 1 {
format!("SELECT * FROM ({}) AS dbx_completion ORDER BY name", queries.remove(0))
} else {
format!("SELECT TOP ({limit}) * FROM ({}) AS dbx_completion ORDER BY name", queries.join(" UNION ALL "))
}
}
fn sqlserver_completion_object_search_clause(
request: &crate::types::CompletionAssistantRequest,
like_pattern: &str,
) -> String {
if like_pattern == "%" {
return String::new();
}
let mut predicates = vec![format!("LOWER(o.name) LIKE LOWER('{like_pattern}') ESCAPE '\\'")];
if request.search_in_comments {
predicates.push(format!("LOWER(COALESCE(ep.value, '')) LIKE LOWER('{like_pattern}') ESCAPE '\\'"));
}
if request.search_in_definitions {
predicates.push(format!(
"LOWER(COALESCE(OBJECT_DEFINITION(o.object_id), '')) LIKE LOWER('{like_pattern}') ESCAPE '\\'"
));
}
format!(" AND ({}) ", predicates.join(" OR "))
}
fn completion_like_pattern(mask: &str, mode: Option<&crate::types::CompletionAssistantMatchMode>) -> String {
if mask.is_empty() || mask == "%" {
return "%".to_string();
}
let has_wildcard = mask.contains('%');
if has_wildcard {
return mask.split('%').map(escape_like_literal).collect::<Vec<_>>().join("%");
}
let escaped = escape_like_literal(mask);
match mode.unwrap_or(&crate::types::CompletionAssistantMatchMode::Prefix) {
crate::types::CompletionAssistantMatchMode::Prefix => format!("{escaped}%"),
crate::types::CompletionAssistantMatchMode::Contains => format!("%{escaped}%"),
}
}
fn sqlserver_list_tables_sql(
schema: &str,
filter: Option<&str>,
limit: Option<usize>,
offset: Option<usize>,
) -> String {
let filter_clause = filter
.filter(|value| !value.trim().is_empty())
.map(|value| {
let contains_pattern = format!("%{}%", escape_like_literal(value.trim()));
if crate::sql::fuzzy_filter_enabled(value) {
let fuzzy_pattern =
crate::sql::fuzzy_like_pattern_with_escape(value.trim(), escape_like_literal);
format!(
" AND (LOWER(o.name) LIKE LOWER('{contains_pattern}') ESCAPE '\\' OR LOWER(o.name) LIKE LOWER('{fuzzy_pattern}') ESCAPE '\\') "
)
} else {
format!(" AND LOWER(o.name) LIKE LOWER('{contains_pattern}') ESCAPE '\\' ")
}
})
.unwrap_or_default();
let base_columns = "o.name, CASE WHEN o.type = 'V' THEN 'VIEW' ELSE 'BASE TABLE' END, ep.value AS TABLE_COMMENT";
let base_from = "FROM sys.objects o \
JOIN sys.schemas s ON s.schema_id = o.schema_id \
OUTER APPLY (SELECT CAST(ep.value AS NVARCHAR(MAX)) AS value FROM sys.extended_properties ep \
WHERE ep.major_id = o.object_id AND ep.minor_id = 0 AND ep.name = N'MS_Description') ep";
let object_visibility = sqlserver_visible_object_predicate();
let schema_filter = sqlserver_schema_name_predicate(schema, "s.name");
let base_where = format!("WHERE {schema_filter} AND o.type IN ('U','V') AND {object_visibility} {filter_clause}");
let order_by = "ORDER BY o.name";
// Use SELECT TOP for broad SQL Server version compatibility.
// OFFSET / FETCH NEXT is only available in SQL Server 2012+.
match (limit, offset) {
(Some(limit), Some(offset)) if offset > 0 => {
let end = offset + limit.min(1000);
format!(
"SELECT * FROM (\
SELECT {base_columns}, ROW_NUMBER() OVER ({order_by}) AS __dbx_rn \
{base_from} {base_where}\
) AS __dbx_page WHERE __dbx_rn > {offset} AND __dbx_rn <= {end} ORDER BY __dbx_rn"
)
}
(Some(limit), _) => {
format!("SELECT TOP ({}) {base_columns} {base_from} {base_where} {order_by}", limit.min(1000))
}
_ => {
format!("SELECT {base_columns} {base_from} {base_where} {order_by}")
}
}
}
fn sqlserver_schema_name_predicate(schema: &str, schema_name_expression: &str) -> String {
if schema.trim().is_empty() {
// SQL Server resolves unqualified objects through the user's default schema.
// A configured default can be missing, so match DBeaver and fall back to dbo.
return format!(
"{schema_name_expression} = COALESCE((SELECT default_schema.name FROM sys.schemas default_schema WHERE default_schema.name = SCHEMA_NAME()), N'dbo')"
);
}
format!("{schema_name_expression} = N'{}'", schema.replace('\'', "''"))
}
fn escape_like_literal(value: &str) -> String {
value.replace('\\', "\\\\").replace('\'', "''").replace('%', "\\%").replace('_', "\\_").replace('[', "\\[")
}
fn sqlserver_visible_object_predicate() -> &'static str {
"(o.is_ms_shipped = 0 OR s.name = 'cdc')"
}
fn sqlserver_hidden_schema_names() -> &'static [&'static str] {
&[
"guest",
"INFORMATION_SCHEMA",
"sys",
"db_owner",
"db_accessadmin",
"db_securityadmin",
"db_ddladmin",
"db_backupoperator",
"db_datareader",
"db_datawriter",
"db_denydatareader",
"db_denydatawriter",
]
}
pub async fn list_objects(client: &mut SqlServerClient, schema: &str) -> Result<Vec<crate::types::ObjectInfo>, String> {
let sql = sqlserver_list_objects_sql(schema);
let stream = client.query(&*sql, &[]).await.map_err(|e| e.to_string())?;
let rows = stream.into_first_result().await.map_err(|e| e.to_string())?;
Ok(rows
.iter()
.map(|row| crate::types::ObjectInfo {
name: row.get::<&str, _>(0).unwrap_or("").to_string(),
object_type: row.get::<&str, _>(1).unwrap_or("TABLE").to_string(),
schema: Some(schema.to_string()),
valid: None,
signature: None,
comment: row.get::<&str, _>(4).filter(|s: &&str| !s.is_empty()).map(|s: &str| s.to_string()),
created_at: row.get::<chrono::NaiveDateTime, _>(2).map(|value| value.to_string()),
updated_at: row.get::<chrono::NaiveDateTime, _>(3).map(|value| value.to_string()),
parent_schema: None,
parent_name: None,
})
.collect())
}
fn sqlserver_list_objects_sql(schema: &str) -> String {
let s = schema.replace('\'', "''");
let object_visibility = sqlserver_visible_object_predicate();
format!(
"SELECT o.name, \
CASE o.type \
WHEN 'U' THEN 'TABLE' \
WHEN 'V' THEN 'VIEW' \
WHEN 'P' THEN 'PROCEDURE' \
WHEN 'FN' THEN 'FUNCTION' \
WHEN 'IF' THEN 'FUNCTION' \
WHEN 'TF' THEN 'FUNCTION' \
WHEN 'FS' THEN 'FUNCTION' \
WHEN 'FT' THEN 'FUNCTION' \
ELSE o.type_desc \
END AS object_type, \
o.create_date, \
o.modify_date, \
ep.value AS object_comment \
FROM sys.objects o \
JOIN sys.schemas s ON s.schema_id = o.schema_id \
OUTER APPLY (SELECT CAST(ep.value AS NVARCHAR(MAX)) AS value FROM sys.extended_properties ep WHERE ep.major_id = o.object_id AND ep.minor_id = 0 AND ep.name = N'MS_Description') ep \
WHERE s.name = '{s}' \
AND o.type IN ('U','V','P','FN','IF','TF','FS','FT') \
AND {object_visibility} \
ORDER BY CASE o.type \
WHEN 'U' THEN 0 \
WHEN 'V' THEN 1 \
WHEN 'P' THEN 2 \
ELSE 3 \
END, o.name"
)
}
pub async fn list_object_statistics(
client: &mut SqlServerClient,
schema: &str,
) -> Result<Vec<ObjectStatistics>, String> {
let s = schema.replace('\'', "''");
let object_visibility = sqlserver_visible_object_predicate();
let sql = format!(
"SELECT o.name, \
SUM(CASE WHEN ps.index_id IN (0, 1) THEN ps.row_count ELSE 0 END) AS estimated_rows, \
SUM(ps.reserved_page_count) * 8192 AS total_bytes \
FROM sys.objects o \
JOIN sys.schemas s ON s.schema_id = o.schema_id \
JOIN sys.dm_db_partition_stats ps ON ps.object_id = o.object_id \
WHERE s.name = '{s}' AND o.type = 'U' AND {object_visibility} \
GROUP BY o.object_id, o.name \
ORDER BY o.name"
);
let stream = client.query(&*sql, &[]).await.map_err(|e| e.to_string())?;
let rows = stream.into_first_result().await.map_err(|e| e.to_string())?;
Ok(rows
.iter()
.map(|row| ObjectStatistics {
name: row.get::<&str, _>(0).unwrap_or("").to_string(),
schema: Some(schema.to_string()),
estimated_rows: row
.try_get::<i64, _>(1)
.ok()
.flatten()
.or_else(|| row.try_get::<i32, _>(1).ok().flatten().map(i64::from)),
total_bytes: row
.try_get::<i64, _>(2)
.ok()
.flatten()
.or_else(|| row.try_get::<i32, _>(2).ok().flatten().map(i64::from)),
})
.filter(|stat| !stat.name.is_empty())
.collect())
}
pub async fn get_columns(client: &mut SqlServerClient, schema: &str, table: &str) -> Result<Vec<ColumnInfo>, String> {
Ok(get_column_metadata(client, schema, table).await?.into_iter().map(|metadata| metadata.column).collect())
}
pub async fn get_column_metadata(
client: &mut SqlServerClient,
schema: &str,
table: &str,
) -> Result<Vec<SqlServerColumnMetadata>, String> {
let sql = sqlserver_columns_sql(schema, table);
let stream = client.query(&*sql, &[]).await.map_err(|e| e.to_string())?;
let rows = stream.into_first_result().await.map_err(|e| e.to_string())?;
Ok(rows.iter().map(sqlserver_column_metadata_from_row).collect())
}
fn sqlserver_column_metadata_from_row(row: &Row) -> SqlServerColumnMetadata {
let base = row.get::<&str, _>(1).unwrap_or("").to_string();
let max_len = row
.try_get::<i32, _>(7)
.ok()
.flatten()
.or_else(|| row.try_get::<i16, _>(7).ok().flatten().map(|v| v as i32))
.or_else(|| row.try_get::<u8, _>(7).ok().flatten().map(|v| v as i32));
let dt_prec = row
.try_get::<i32, _>(8)
.ok()
.flatten()
.or_else(|| row.try_get::<i16, _>(8).ok().flatten().map(|v| v as i32))
.or_else(|| row.try_get::<u8, _>(8).ok().flatten().map(|v| v as i32));
let num_prec = row
.try_get::<i32, _>(5)
.ok()
.flatten()
.or_else(|| row.try_get::<i16, _>(5).ok().flatten().map(|v| v as i32))
.or_else(|| row.try_get::<u8, _>(5).ok().flatten().map(|v| v as i32));
let num_scale = row
.try_get::<i32, _>(6)
.ok()
.flatten()
.or_else(|| row.try_get::<i16, _>(6).ok().flatten().map(|v| v as i32))
.or_else(|| row.try_get::<u8, _>(6).ok().flatten().map(|v| v as i32));
let data_type = match base.to_lowercase().as_str() {
"varchar" => match max_len {
Some(-1) => "varchar(max)".to_string(),
Some(n) => format!("varchar({n})"),
None => "varchar".to_string(),
},
"nvarchar" => match max_len {
Some(-1) => "nvarchar(max)".to_string(),
Some(n) => format!("nvarchar({n})"),
None => "nvarchar".to_string(),
},
"varbinary" => match max_len {
Some(-1) => "varbinary(max)".to_string(),
Some(n) if n > 0 => format!("varbinary({n})"),
_ => "varbinary".to_string(),
},
"char" | "nchar" | "binary" => match max_len {
Some(n) if n > 0 => format!("{base}({n})"),
_ => base,
},
"decimal" | "numeric" => match (num_prec, num_scale) {
(Some(p), Some(s)) => format!("{base}({p},{s})"),
_ => base,
},
"datetime2" | "datetimeoffset" | "time" => match dt_prec {
Some(p) => format!("{base}({p})"),
_ => base,
},
_ => base,
};
let column = ColumnInfo {
name: row.get::<&str, _>(0).unwrap_or("").to_string(),
data_type,
is_nullable: row.get::<&str, _>(2).unwrap_or("NO") == "YES",
column_default: row.get::<&str, _>(3).map(|s| s.to_string()),
is_primary_key: row.get::<i32, _>(4).unwrap_or(0) == 1,
extra: row.get::<&str, _>(9).filter(|s: &&str| !s.is_empty()).map(|s: &str| s.to_string()),
comment: row.get::<&str, _>(10).filter(|s: &&str| !s.is_empty()).map(|s: &str| s.to_string()),
numeric_precision: num_prec,
numeric_scale: num_scale,
character_maximum_length: max_len,
enum_values: None,
..Default::default()
};
SqlServerColumnMetadata {
column,
is_identity: row.get::<i32, _>(11).unwrap_or(0) == 1,
is_computed: row.get::<i32, _>(12).unwrap_or(0) == 1,
is_hidden: row.get::<i32, _>(13).unwrap_or(0) == 1,
generated_always_type: row.get::<i32, _>(14).unwrap_or(0),
}
}
fn sqlserver_columns_sql(schema: &str, table: &str) -> String {
let t = table.replace('\'', "''");
let schema_filter = sqlserver_schema_name_predicate(schema, "s.name");
// COLUMNPROPERTY keeps hidden/generated flags separate and returns NULL on
// SQL Server versions that do not expose a newer property.
format!(
"SELECT c.name AS COLUMN_NAME, \
ty.name AS DATA_TYPE, \
CASE WHEN c.is_nullable = 1 THEN 'YES' ELSE 'NO' END AS IS_NULLABLE, \
dc.definition AS COLUMN_DEFAULT, \
CASE WHEN pk.column_id IS NOT NULL THEN 1 ELSE 0 END AS IS_PK, \
CONVERT(INT, c.precision) AS NUMERIC_PRECISION, \
CONVERT(INT, c.scale) AS NUMERIC_SCALE, \
CASE \
WHEN ty.name IN ('nchar','nvarchar') AND c.max_length > 0 THEN CONVERT(INT, c.max_length / 2) \
WHEN c.max_length = -1 THEN -1 \
ELSE CONVERT(INT, c.max_length) \
END AS CHARACTER_MAXIMUM_LENGTH, \
CONVERT(INT, c.scale) AS DATETIME_PRECISION, \
CASE \
WHEN c.is_computed = 1 THEN 'computed' \
WHEN ic.column_id IS NOT NULL THEN 'identity(' + CONVERT(VARCHAR(38), ic.seed_value) + ',' + CONVERT(VARCHAR(38), ic.increment_value) + ')' \
ELSE NULL \
END AS COLUMN_EXTRA, \
ep.value AS COLUMN_COMMENT, \
CONVERT(INT, COLUMNPROPERTY(c.object_id, c.name, 'IsIdentity')) AS IS_IDENTITY, \
CONVERT(INT, COLUMNPROPERTY(c.object_id, c.name, 'IsComputed')) AS IS_COMPUTED, \
CONVERT(INT, COLUMNPROPERTY(c.object_id, c.name, 'IsHidden')) AS IS_HIDDEN, \
CONVERT(INT, COLUMNPROPERTY(c.object_id, c.name, 'GeneratedAlwaysType')) AS GENERATED_ALWAYS_TYPE \
FROM sys.objects o \
JOIN sys.schemas s ON s.schema_id = o.schema_id \
JOIN sys.columns c ON c.object_id = o.object_id \
JOIN sys.types ty ON ty.user_type_id = c.user_type_id \
LEFT JOIN sys.default_constraints dc ON dc.object_id = c.default_object_id \
LEFT JOIN sys.identity_columns ic ON ic.object_id = c.object_id AND ic.column_id = c.column_id \
LEFT JOIN ( \
SELECT ic.object_id, ic.column_id \
FROM sys.indexes i \
JOIN sys.index_columns ic ON ic.object_id = i.object_id AND ic.index_id = i.index_id \
WHERE i.is_primary_key = 1 \
) pk ON pk.object_id = c.object_id AND pk.column_id = c.column_id \
OUTER APPLY (SELECT CAST(ep.value AS NVARCHAR(MAX)) AS value FROM sys.extended_properties ep WHERE ep.major_id = c.object_id AND ep.minor_id = c.column_id AND ep.name = N'MS_Description') ep \
WHERE {schema_filter} AND o.name = '{t}' AND o.type IN ('U','V') \
ORDER BY c.column_id"
)
}
pub async fn list_indexes(client: &mut SqlServerClient, schema: &str, table: &str) -> Result<Vec<IndexInfo>, String> {
let sql = sqlserver_indexes_sql(schema, table);
let stream = client.query(&*sql, &[]).await.map_err(|e| e.to_string())?;
let rows = stream.into_first_result().await.map_err(|e| e.to_string())?;
Ok(rows
.iter()
.map(|row| {
let cols_str = row.get::<&str, _>(1).unwrap_or("");
let inc_str = row.get::<&str, _>(5).unwrap_or("");
IndexInfo {
name: row.get::<&str, _>(0).unwrap_or("").to_string(),
columns: cols_str.split(',').filter(|s| !s.is_empty()).map(|s| s.to_string()).collect(),
is_unique: row.get::<bool, _>(2).unwrap_or(false),
is_primary: row.get::<bool, _>(3).unwrap_or(false),
filter: row.get::<&str, _>(6).map(|s| s.to_string()),
index_type: row.get::<&str, _>(4).map(|s| s.to_string()),
included_columns: if inc_str.is_empty() {
None
} else {
Some(inc_str.split(',').map(|s| s.to_string()).collect())
},
comment: row.get::<&str, _>(7).filter(|s: &&str| !s.is_empty()).map(|s: &str| s.to_string()),
}
})
.collect())
}
fn sqlserver_indexes_sql(schema: &str, table: &str) -> String {
format!(
"SELECT i.name, \
STUFF((SELECT ',' + c2.name \
FROM sys.index_columns ic2 \
JOIN sys.columns c2 ON ic2.object_id = c2.object_id AND ic2.column_id = c2.column_id \
WHERE ic2.object_id = i.object_id AND ic2.index_id = i.index_id AND ic2.is_included_column = 0 \
ORDER BY ic2.key_ordinal \
FOR XML PATH(''), TYPE).value('.', 'nvarchar(max)'), 1, 1, '') AS columns, \
i.is_unique, i.is_primary_key, i.type_desc, \
STUFF((SELECT ',' + c3.name \
FROM sys.index_columns ic3 \
JOIN sys.columns c3 ON ic3.object_id = c3.object_id AND ic3.column_id = c3.column_id \
WHERE ic3.object_id = i.object_id AND ic3.index_id = i.index_id AND ic3.is_included_column = 1 \
ORDER BY ic3.index_column_id \
FOR XML PATH(''), TYPE).value('.', 'nvarchar(max)'), 1, 1, '') AS included_cols, \
i.filter_definition, \
ep.value AS index_comment \
FROM sys.indexes i \
OUTER APPLY (SELECT CAST(ep.value AS NVARCHAR(MAX)) AS value FROM sys.extended_properties ep WHERE ep.major_id = i.object_id AND ep.minor_id = i.index_id AND ep.name = N'MS_Description' AND ep.class = 7) ep \
WHERE i.object_id = OBJECT_ID('{s}.{t}') AND i.name IS NOT NULL \
ORDER BY i.name",
s = schema.replace('\'', "''"),
t = table.replace('\'', "''")
)
}
pub async fn list_foreign_keys(
client: &mut SqlServerClient,
schema: &str,
table: &str,
) -> Result<Vec<ForeignKeyInfo>, String> {
let sql = format!(
"SELECT fk.name, c.name, SCHEMA_NAME(rt.schema_id), rt.name, rc.name \
FROM sys.foreign_keys fk \
JOIN sys.foreign_key_columns fkc ON fk.object_id = fkc.constraint_object_id \
JOIN sys.columns c ON fkc.parent_object_id = c.object_id AND fkc.parent_column_id = c.column_id \
JOIN sys.tables rt ON fkc.referenced_object_id = rt.object_id \
JOIN sys.columns rc ON fkc.referenced_object_id = rc.object_id AND fkc.referenced_column_id = rc.column_id \
WHERE fk.parent_object_id = OBJECT_ID('{s}.{t}') \
ORDER BY fk.name, fkc.constraint_column_id",
s = schema.replace('\'', "''"),
t = table.replace('\'', "''")
);
let stream = client.query(&*sql, &[]).await.map_err(|e| e.to_string())?;
let rows = stream.into_first_result().await.map_err(|e| e.to_string())?;
Ok(rows
.iter()
.map(|row| ForeignKeyInfo {
name: row.get::<&str, _>(0).unwrap_or("").to_string(),
column: row.get::<&str, _>(1).unwrap_or("").to_string(),
ref_schema: Some(row.get::<&str, _>(2).unwrap_or("").to_string()),
ref_table: row.get::<&str, _>(3).unwrap_or("").to_string(),
ref_column: row.get::<&str, _>(4).unwrap_or("").to_string(),
on_update: None,
on_delete: None,
})
.collect())
}
pub async fn get_table_comment(
client: &mut SqlServerClient,
schema: &str,
table: &str,
) -> Result<Option<String>, String> {
let sql = sqlserver_table_comment_sql(schema, table);
let stream = client.query(&*sql, &[]).await.map_err(|e| e.to_string())?;
let rows = stream.into_first_result().await.map_err(|e| e.to_string())?;
Ok(rows.first().and_then(|row| row.get::<&str, _>(0)).filter(|s| !s.is_empty()).map(|s| s.to_string()))
}
fn sqlserver_table_comment_sql(schema: &str, table: &str) -> String {
let s = schema.replace('\'', "''");
let t = table.replace('\'', "''");
format!(
"SELECT CAST(ep.value AS NVARCHAR(MAX)) \
FROM sys.extended_properties ep \
WHERE ep.major_id = OBJECT_ID(QUOTENAME('{s}') + '.' + QUOTENAME('{t}')) \
AND ep.minor_id = 0 \
AND ep.name = N'MS_Description'"
)
}
pub async fn list_triggers(
client: &mut SqlServerClient,
schema: &str,
table: &str,
) -> Result<Vec<TriggerInfo>, String> {
let sql = format!(
"SELECT t.name, te.type_desc, CASE WHEN t.is_instead_of_trigger = 1 THEN 'INSTEAD OF' ELSE 'AFTER' END \
FROM sys.triggers t \
JOIN sys.trigger_events te ON t.object_id = te.object_id \
WHERE t.parent_id = OBJECT_ID('{s}.{t}') \
ORDER BY t.name",
s = schema.replace('\'', "''"),
t = table.replace('\'', "''")
);
let stream = client.query(&*sql, &[]).await.map_err(|e| e.to_string())?;
let rows = stream.into_first_result().await.map_err(|e| e.to_string())?;
Ok(rows
.iter()
.map(|row| TriggerInfo {
name: row.get::<&str, _>(0).unwrap_or("").to_string(),
event: row.get::<&str, _>(1).unwrap_or("").to_string(),
timing: row.get::<&str, _>(2).unwrap_or("AFTER").to_string(),
statement: None,
})
.collect())
}
pub async fn execute_query(client: &mut SqlServerClient, sql: &str) -> Result<QueryResult, String> {
execute_query_with_max_rows(client, sql, None).await
}
pub async fn execute_query_with_max_rows(
client: &mut SqlServerClient,
sql: &str,
max_rows: Option<usize>,
) -> Result<QueryResult, String> {
let start = Instant::now();
if starts_with_executable_sql_keyword(sql, &["SELECT", "EXEC", "WITH", "TABLE"]) {
let query_sql = match sqlserver_unsafe_type_query(client, sql).await {
Ok(Some(sql)) => sql,
Ok(None) => sql.to_string(),
Err(error) if is_blocking_sqlserver_unsafe_probe_error(&error) => return Err(error),
Err(_) => sql.to_string(),
};
let (result, messages) = capture_sqlserver_messages(async {
let stream = sqlserver_driver_result(client.query(query_sql.as_str(), &[])).await?;
sqlserver_driver_result(collect_first_result_limited(stream, start, max_rows)).await
})
.await;
let mut result = query_result_with_server_messages(result?, messages);
strip_dbx_sqlserver_row_number_column(&mut result, sql);
Ok(result)
} else if requires_simple_query_batch(sql) || contains_transaction_control(sql) {
let (result, messages) = capture_sqlserver_messages(async {
let stream = sqlserver_driver_result(client.simple_query(sql)).await?;
sqlserver_driver_result(collect_result_sets_limited(stream, start, max_rows)).await
})
.await;
let _ = result?;
Ok(query_result_with_server_messages(
QueryResult {
columns: vec![],
column_types: Vec::new(),
column_sortables: vec![],
rows: vec![],
affected_rows: 0,
execution_time_ms: start.elapsed().as_millis(),
truncated: false,
session_id: None,
has_more: false,
elasticsearch_raw_body: None,
},
messages,
))
} else {
let (result, messages) = capture_sqlserver_messages(sqlserver_driver_result(client.execute(sql, &[]))).await;
let result = result?;
Ok(query_result_with_server_messages(
QueryResult {
columns: vec![],
column_types: Vec::new(),
column_sortables: vec![],
rows: vec![],
affected_rows: result.rows_affected().iter().sum::<u64>(),
execution_time_ms: start.elapsed().as_millis(),
truncated: false,
session_id: None,
has_more: false,
elasticsearch_raw_body: None,
},
messages,
))
}
}
pub async fn execute_batch(client: &mut SqlServerClient, sql: &str) -> Result<Vec<QueryResult>, String> {
execute_batch_with_max_rows(client, sql, None).await
}
pub async fn execute_batch_with_max_rows(
client: &mut SqlServerClient,
sql: &str,
max_rows: Option<usize>,
) -> Result<Vec<QueryResult>, String> {
let start = Instant::now();
if sqlserver_batch_can_use_execute(sql) {
let (result, messages) = capture_sqlserver_messages(sqlserver_driver_result(client.execute(sql, &[]))).await;
let result = result?;
return Ok(vec![query_result_with_server_messages(
QueryResult {
columns: vec![],
column_types: Vec::new(),
column_sortables: vec![],
rows: vec![],
affected_rows: result.rows_affected().iter().sum::<u64>(),
execution_time_ms: start.elapsed().as_millis(),
truncated: false,
session_id: None,
has_more: false,
elasticsearch_raw_body: None,
},
messages,
)]);
}
if is_single_sqlserver_select(sql) {
match sqlserver_unsafe_type_query(client, sql).await {
Ok(Some(query_sql)) => {
let (result, messages) = capture_sqlserver_messages(async {
let stream = sqlserver_driver_result(client.query(query_sql.as_str(), &[])).await?;
sqlserver_driver_result(collect_first_result_limited(stream, start, max_rows)).await
})
.await;
return result.map(|result| {
let mut result = query_result_with_server_messages(result, messages);
strip_dbx_sqlserver_row_number_column(&mut result, sql);
vec![result]
});
}
Err(error) if is_blocking_sqlserver_unsafe_probe_error(&error) => return Err(error),
Ok(None) | Err(_) => {}
}
}
execute_simple_batch_with_max_rows(client, sql, max_rows).await
}
/// Execute a SQL Server batch directly through TDS simple-query mode.
///
/// This intentionally bypasses result-set type probing and SQL rewriting. It is
/// required while `SHOWPLAN_XML` is enabled because any probe issued on the same
/// session is itself affected by SHOWPLAN state.
pub async fn execute_simple_batch_with_max_rows(
client: &mut SqlServerClient,
sql: &str,
max_rows: Option<usize>,
) -> Result<Vec<QueryResult>, String> {
let start = Instant::now();
let (results, messages) = capture_sqlserver_messages(async {
let stream = sqlserver_driver_result(client.simple_query(sql)).await?;
sqlserver_driver_result(collect_result_sets_limited(stream, start, max_rows)).await
})
.await;
let mut results = results?;
for result in &mut results {
strip_dbx_sqlserver_row_number_column(result, sql);
}
if let Some(message_result) = server_messages_query_result(messages, start) {
results.push(message_result);
} else if results.is_empty() {
results.push(QueryResult {
columns: vec![],
column_types: Vec::new(),
column_sortables: vec![],
rows: vec![],
affected_rows: 0,
execution_time_ms: start.elapsed().as_millis(),
truncated: false,
session_id: None,
has_more: false,
elasticsearch_raw_body: None,
});
}
Ok(results)
}
fn strip_dbx_sqlserver_row_number_column(result: &mut QueryResult, sql: &str) {
if !is_dbx_sqlserver_row_number_page_sql(sql) {
return;
}
if !result.columns.last().is_some_and(|column| column.eq_ignore_ascii_case("__dbx_row_num")) {
return;
}
result.columns.pop();
if result.column_types.len() > result.columns.len() {
result.column_types.pop();
}
if result.column_sortables.len() > result.columns.len() {
result.column_sortables.pop();
}
for row in &mut result.rows {
if row.len() > result.columns.len() {
row.pop();
}
}
}
fn is_dbx_sqlserver_row_number_page_sql(sql: &str) -> bool {
let normalized = sql.to_ascii_uppercase();
normalized.contains("ROW_NUMBER() OVER")
&& normalized.contains("[__DBX_ROW_NUM]")
&& normalized.contains("DBX_PAGE_SOURCE.*")
}
fn sqlserver_batch_can_use_execute(sql: &str) -> bool {
!requires_simple_query_batch(sql)
// Tiberius executes this path through an RPC. SQL Server rejects an RPC
// that changes @@TRANCOUNT with error 266, while a regular batch is allowed
// to leave an explicit transaction open for a later COMMIT or ROLLBACK.
&& !contains_transaction_control(sql)
&& !sqlserver_batch_may_return_result_set(sql)
&& !sqlserver_dml_output_returns_rows(sql)
}
fn sqlserver_batch_may_return_result_set(sql: &str) -> bool {
let tokens = top_level_sqlserver_tokens(sql);
let starts_with_cte_dml = tokens.first().is_some_and(|token| token.text == "WITH")
&& tokens.iter().any(|token| matches!(token.text.as_str(), "INSERT" | "UPDATE" | "DELETE" | "MERGE"))
&& crate::sql::split_sql_statements(sql).len() == 1;
tokens.iter().any(|token| {
matches!(token.text.as_str(), "SELECT" | "EXEC" | "EXECUTE" | "TABLE")
|| (token.text == "WITH" && !starts_with_cte_dml)
})
}
fn sqlserver_dml_output_returns_rows(sql: &str) -> bool {
crate::sql::split_sql_statements(sql).iter().any(|statement| {
let tokens = top_level_sqlserver_tokens(statement);
let contains_dml = starts_with_executable_sql_keyword(statement, &["INSERT", "UPDATE", "DELETE", "MERGE"])
|| (tokens.first().is_some_and(|token| token.text == "WITH")
&& tokens.iter().any(|token| matches!(token.text.as_str(), "INSERT" | "UPDATE" | "DELETE" | "MERGE")));
if !contains_dml {
return false;
}
tokens.iter().enumerate().any(|(output_index, token)| {
if token.text != "OUTPUT" || (token.start > 0 && statement.as_bytes()[token.start - 1] == b'@') {
return false;
}
// SQL Server may combine OUTPUT ... INTO with a second OUTPUT clause.
// Only an OUTPUT whose rows are not routed before the next clause reaches the client.
!tokens[output_index + 1..].iter().take_while(|next| next.text != "OUTPUT").any(|next| next.text == "INTO")
})
})
}
fn contains_transaction_control(sql: &str) -> bool {
let tokens = top_level_sqlserver_tokens(sql);
tokens.iter().enumerate().any(|(index, token)| {
if matches!(token.text.as_str(), "COMMIT" | "ROLLBACK") {
return true;
}
if token.text == "BEGIN" {
if tokens.get(index + 1).is_some_and(|next| matches!(next.text.as_str(), "TRANSACTION" | "TRAN")) {
return true;
}
if tokens.get(index + 1).is_some_and(|next| next.text == "DISTRIBUTED")
&& tokens.get(index + 2).is_some_and(|next| matches!(next.text.as_str(), "TRANSACTION" | "TRAN"))
{
return true;
}
}
token.text == "SET"
&& tokens.get(index + 1).is_some_and(|next| next.text == "IMPLICIT_TRANSACTIONS")
&& tokens.get(index + 2).is_some_and(|next| next.text == "ON")
})
}
fn requires_simple_query_batch(sql: &str) -> bool {
let tokens = first_sql_tokens(sql, 4);
if tokens.len() >= 2 && tokens[0].eq_ignore_ascii_case("SET") && tokens[1].eq_ignore_ascii_case("SHOWPLAN_XML") {
return true;
}
if tokens.len() >= 2 && tokens[0].eq_ignore_ascii_case("CREATE") && tokens[1].eq_ignore_ascii_case("SCHEMA") {
return true;
}
if tokens.len() >= 4
&& tokens[0].eq_ignore_ascii_case("CREATE")
&& tokens[1].eq_ignore_ascii_case("OR")
&& tokens[2].eq_ignore_ascii_case("ALTER")
{
return SIMPLE_QUERY_MODULE_KEYWORDS.iter().any(|keyword| tokens[3].eq_ignore_ascii_case(keyword));
}
if tokens.len() >= 2 && (tokens[0].eq_ignore_ascii_case("CREATE") || tokens[0].eq_ignore_ascii_case("ALTER")) {
return SIMPLE_QUERY_MODULE_KEYWORDS.iter().any(|keyword| tokens[1].eq_ignore_ascii_case(keyword));
}
false
}
fn first_sql_tokens(sql: &str, limit: usize) -> Vec<String> {
let bytes = sql.as_bytes();
let mut tokens = Vec::new();
let mut i = 0;
while i < bytes.len() && tokens.len() < limit {
while i < bytes.len() && bytes[i].is_ascii_whitespace() {
i += 1;
}
if i + 1 < bytes.len() && bytes[i] == b'-' && bytes[i + 1] == b'-' {
i += 2;
while i < bytes.len() && bytes[i] != b'\n' {
i += 1;
}
continue;
}
if i + 1 < bytes.len() && bytes[i] == b'/' && bytes[i + 1] == b'*' {
i += 2;
while i + 1 < bytes.len() && !(bytes[i] == b'*' && bytes[i + 1] == b'/') {
i += 1;
}
i = (i + 2).min(bytes.len());
continue;
}
let start = i;
while i < bytes.len() && (bytes[i].is_ascii_alphanumeric() || bytes[i] == b'_') {
i += 1;
}
if i > start {
tokens.push(sql[start..i].to_string());
} else {
i += 1;
}
}
tokens
}
#[cfg(test)]
mod tests {
use super::{
build_sqlserver_unsafe_type_query, capture_sqlserver_messages, format_sqlserver_numeric,
is_blocking_sqlserver_unsafe_probe_error, is_sqlserver_spatial_column, is_sqlserver_variant_column,
query_result_with_server_messages, requires_simple_query_batch, sqlserver_batch_can_use_execute,
sqlserver_cell_to_json, sqlserver_columns_sql, sqlserver_completion_assistant_sql,
sqlserver_dml_output_returns_rows, sqlserver_hidden_schema_names, sqlserver_indexes_sql,
sqlserver_legacy_probe, sqlserver_list_objects_sql, sqlserver_list_schemas_sql, sqlserver_list_tables_sql,
sqlserver_schema_name_predicate, sqlserver_table_comment_sql, sqlserver_visible_object_predicate,
strip_dbx_sqlserver_row_number_column, SqlServerDescribedColumn, SqlServerResultSet,
SQLSERVER_RESULT_TYPE_PROBE_SQL,
};
use crate::types::{
CompletionAssistantMatchMode, CompletionAssistantObjectKind, CompletionAssistantRequest, QueryResult,
};
use chrono::NaiveDate;
use std::{borrow::Cow, time::Instant};
use tiberius::{Column, ColumnData, ColumnType, IntoSql};
#[tokio::test]
async fn sqlserver_ignores_non_info_tiberius_events() {
let (_, messages) = capture_sqlserver_messages(async {
tracing::event!(target: "tiberius::tds::stream::token", tracing::Level::ERROR, "permission denied");
tracing::event!(target: "dbx_core::db::sqlserver", tracing::Level::INFO, "not a TDS token");
})
.await;
assert!(messages.is_empty());
}
#[test]
fn sqlserver_server_messages_fill_only_empty_results() {
let empty = QueryResult {
columns: vec![],
column_types: vec![],
column_sortables: vec![],
rows: vec![],
affected_rows: 0,
execution_time_ms: 1,
truncated: false,
session_id: None,
has_more: false,
elasticsearch_raw_body: None,
};
let result = query_result_with_server_messages(empty, vec!["DBCC execution completed".to_string()]);
assert_eq!(result.columns, vec!["Message"]);
assert_eq!(result.rows, vec![vec![serde_json::json!("DBCC execution completed")]]);
let select = QueryResult {
columns: vec!["id".to_string()],
column_types: vec!["int".to_string()],
column_sortables: vec![],
rows: vec![vec![serde_json::json!(1)]],
affected_rows: 0,
execution_time_ms: 1,
truncated: false,
session_id: None,
has_more: false,
elasticsearch_raw_body: None,
};
let result = query_result_with_server_messages(select, vec!["informational".to_string()]);
assert_eq!(result.columns, vec!["id"]);
assert_eq!(result.rows, vec![vec![serde_json::json!(1)]]);
}
#[test]
fn sqlserver_endpoint_splits_named_instance_hosts() {
assert_eq!(
super::sqlserver_endpoint(r"192.168.1.10\SQL2022"),
super::SqlServerEndpoint { host: "192.168.1.10", instance_name: Some("SQL2022") }
);
assert_eq!(
super::sqlserver_endpoint(r" db.example.com\SQLEXPRESS "),
super::SqlServerEndpoint { host: "db.example.com", instance_name: Some("SQLEXPRESS") }
);
}
#[test]
fn sqlserver_xml_cells_are_returned_as_strings() {
let cell = ColumnData::Xml(Some(Cow::Owned(tiberius::xml::XmlData::new(
"<ShowPlanXML><RelOp NodeId=\"0\" /></ShowPlanXML>",
))));
assert_eq!(
sqlserver_cell_to_json(&cell),
serde_json::Value::String("<ShowPlanXML><RelOp NodeId=\"0\" /></ShowPlanXML>".to_string())
);
}
#[test]
fn sqlserver_uses_sql_type_names_for_tds_big_types() {
for (column_type, expected) in [
(ColumnType::BigVarChar, "varchar"),
(ColumnType::BigChar, "char"),
(ColumnType::BigVarBin, "varbinary"),
(ColumnType::BigBinary, "binary"),
] {
let column = Column::new("value".to_string(), column_type);
assert_eq!(super::sqlserver_column_type_name(&column), expected);
}
}
#[test]
fn sqlserver_endpoint_keeps_regular_hosts() {
assert_eq!(
super::sqlserver_endpoint("db.example.com"),
super::SqlServerEndpoint { host: "db.example.com", instance_name: None }
);
assert_eq!(
super::sqlserver_endpoint(r"db.example.com\"),
super::SqlServerEndpoint { host: r"db.example.com\", instance_name: None }
);
}
#[test]
fn sqlserver_named_instance_resolution_yields_to_explicit_port() {
let endpoint = super::sqlserver_endpoint(r"db.example.com\SQLEXPRESS");
assert!(super::sqlserver_uses_named_instance_resolution(&endpoint, 0, false));
assert!(super::sqlserver_uses_named_instance_resolution(&endpoint, 1433, false));
assert!(!super::sqlserver_uses_named_instance_resolution(&endpoint, 1433, true));
assert!(!super::sqlserver_uses_named_instance_resolution(&endpoint, 40030, false));
}
#[test]
fn sqlserver_connect_uses_named_instance_resolution() {
let source = include_str!("sqlserver.rs");
let try_connect = source.split("\nasync fn try_connect(").nth(1).unwrap();
let try_connect = try_connect.split("fn row_to_json").next().unwrap();
assert!(try_connect.contains("connect_named(&config)"));
}
#[test]
fn sqlserver_native_encryption_flag_accepts_dbx_and_jdbc_params() {
assert!(!super::sqlserver_native_encryption_disabled(None));
assert!(!super::sqlserver_native_encryption_disabled(Some("encrypt=true")));
assert!(super::sqlserver_native_encryption_disabled(Some("sqlserverEncryption=disabled")));
assert!(super::sqlserver_native_encryption_disabled(Some("applicationName=dbx;sqlserverEncryption=off")));
assert!(super::sqlserver_native_encryption_disabled(Some("?sqlserverEncryption=false&applicationName=dbx")));
assert!(super::sqlserver_native_encryption_disabled(Some("applicationName=dbx;encrypt=false")));
assert!(super::sqlserver_native_encryption_disabled(Some("?Encrypt=0&applicationName=dbx")));
}
#[test]
fn sqlserver_legacy_encryption_modes_cover_jdbc_and_no_encryption_fallback() {
assert_eq!(super::SQLSERVER_LEGACY_ENCRYPTION_LEVEL, tiberius::EncryptionLevel::Off);
assert_eq!(super::SQLSERVER_UNSUPPORTED_ENCRYPTION_LEVEL, tiberius::EncryptionLevel::NotSupported);
}
#[test]
fn sqlserver_automatic_fallback_preserves_v48_no_encryption_compatibility() {
let levels = super::SQLSERVER_LEGACY_ENCRYPTION_FALLBACKS.map(|(_, encryption)| encryption);
assert_eq!(levels, [tiberius::EncryptionLevel::Off, tiberius::EncryptionLevel::NotSupported]);
}
#[test]
fn sqlserver_tls_handshake_error_detection_matches_legacy_hint_cases() {
assert!(super::is_sqlserver_tls_handshake_error(
"SQL Server connection failed: An error occured during the attempt of performing I/O: tls handshake eof"
));
assert!(super::is_sqlserver_tls_handshake_error("TLS handshake failed: unexpected EOF"));
assert!(!super::is_sqlserver_tls_handshake_error("SQL Server connection failed: Login failed for user"));
}
#[test]
fn sqlserver_module_definitions_require_simple_query_batch() {
assert!(requires_simple_query_batch("SET SHOWPLAN_XML ON;"));
assert!(requires_simple_query_batch("SET SHOWPLAN_XML OFF;"));
assert!(requires_simple_query_batch("CREATE SCHEMA [analytics];"));
assert!(requires_simple_query_batch("CREATE FUNCTION dbo.fn_demo() RETURNS INT AS BEGIN RETURN 1; END;"));
assert!(requires_simple_query_batch("ALTER PROCEDURE dbo.usp_demo AS SELECT 1;"));
assert!(requires_simple_query_batch("CREATE OR ALTER VIEW dbo.vw_demo AS SELECT 1 AS id;"));
assert!(requires_simple_query_batch(
"-- comment\nALTER TRIGGER dbo.tr_demo ON dbo.t AFTER INSERT AS SELECT 1;"
));
}
#[test]
fn sqlserver_regular_ddl_can_use_execute() {
assert!(!sqlserver_batch_can_use_execute("CREATE SCHEMA [analytics];"));
assert!(!requires_simple_query_batch("ALTER TABLE dbo.t ADD name NVARCHAR(20);"));
assert!(!requires_simple_query_batch("CREATE TABLE dbo.t(id INT);"));
assert!(!requires_simple_query_batch("UPDATE dbo.t SET id = 1;"));
}
#[test]
fn sqlserver_cud_batches_use_execute_for_affected_rows() {
assert!(sqlserver_batch_can_use_execute("UPDATE dbo.users SET active = 0 WHERE id = 1;"));
assert!(sqlserver_batch_can_use_execute("INSERT INTO dbo.users(id) VALUES (1);"));
assert!(sqlserver_batch_can_use_execute("DELETE FROM dbo.users WHERE id = 1;"));
assert!(sqlserver_batch_can_use_execute(
"MERGE dbo.t AS t USING dbo.s AS s ON t.id = s.id WHEN MATCHED THEN UPDATE SET name = s.name;"
));
}
#[test]
fn sqlserver_cte_dml_batches_use_execute_for_affected_rows() {
assert!(sqlserver_batch_can_use_execute(
";WITH cte AS (SELECT 1 AS id) UPDATE dbo.users SET active = 0 FROM dbo.users JOIN cte ON cte.id = users.id;"
));
assert!(sqlserver_batch_can_use_execute(
"WITH a AS (SELECT 1 AS id), b AS (SELECT id FROM a) DELETE dbo.users FROM dbo.users JOIN b ON b.id = users.id;"
));
assert!(sqlserver_batch_can_use_execute(
"WITH source AS (SELECT 1 AS id) MERGE dbo.users AS target USING source ON target.id = source.id WHEN MATCHED THEN UPDATE SET active = 0;"
));
assert!(sqlserver_batch_can_use_execute(
"WITH cte AS (SELECT 1 AS id) UPDATE dbo.users WITH (ROWLOCK) SET note = 'SELECT OUTPUT' /* WITH SELECT OUTPUT */;"
));
assert!(sqlserver_batch_can_use_execute(
"WITH cte AS (SELECT 1 AS id) UPDATE dbo.users SET active = @output WHERE id = 1;"
));
assert!(!sqlserver_batch_can_use_execute("WITH cte AS (SELECT 1 AS id) SELECT * FROM cte;"));
assert!(!sqlserver_batch_can_use_execute("WITH cte AS (SELECT 1 AS id) (SELECT id FROM cte);"));
assert!(!sqlserver_batch_can_use_execute(
"WITH cte AS (SELECT 1 AS id) (SELECT id FROM cte) UNION (SELECT 2);"
));
assert!(!sqlserver_batch_can_use_execute(
"WITH cte AS (SELECT 1 AS id) (SELECT id FROM cte); UPDATE dbo.users SET active = 0;"
));
assert!(!sqlserver_batch_can_use_execute("WITH XMLNAMESPACES ('urn:demo' AS ns) SELECT 1 AS id;"));
assert!(!sqlserver_batch_can_use_execute(
"WITH cte AS (SELECT 1 AS id) UPDATE dbo.users SET active = 0 OUTPUT inserted.id;"
));
assert!(!sqlserver_batch_can_use_execute(
"WITH cte AS (SELECT 1 AS id) DELETE dbo.users OUTPUT deleted.id FROM dbo.users JOIN cte ON cte.id = users.id;"
));
assert!(!sqlserver_batch_can_use_execute(
"WITH source AS (SELECT 1 AS id) MERGE dbo.users AS target USING source ON target.id = source.id WHEN MATCHED THEN DELETE OUTPUT deleted.id;"
));
assert!(sqlserver_batch_can_use_execute(
"WITH cte AS (SELECT 1 AS id) UPDATE dbo.users SET active = 0 OUTPUT inserted.id INTO dbo.audit_ids;"
));
assert!(sqlserver_batch_can_use_execute(
"WITH cte AS (SELECT 1 AS id) DELETE dbo.users OUTPUT deleted.id INTO dbo.audit_ids FROM dbo.users JOIN cte ON cte.id = users.id;"
));
assert!(sqlserver_batch_can_use_execute(
"WITH source AS (SELECT 1 AS id) MERGE dbo.users AS target USING source ON target.id = source.id WHEN MATCHED THEN DELETE OUTPUT deleted.id INTO dbo.audit_ids;"
));
assert!(!sqlserver_batch_can_use_execute(
"WITH source AS (SELECT 1 AS id) MERGE dbo.users AS target USING source ON target.id = source.id WHEN MATCHED THEN UPDATE SET active = 0 OUTPUT inserted.id INTO dbo.audit_ids OUTPUT inserted.id;"
));
assert!(!sqlserver_batch_can_use_execute(
"WITH cte AS (SELECT 1 AS id) UPDATE dbo.users SET active = 0; SELECT 1;"
));
assert!(!sqlserver_batch_can_use_execute(
"WITH cte AS (SELECT 1 AS id) UPDATE dbo.users SET active = 0 OUTPUT inserted.id INTO dbo.audit_ids; SELECT 1;"
));
assert!(!sqlserver_batch_can_use_execute(
"WITH cte AS (SELECT 1 AS id) INSERT INTO dbo.users(id) SELECT id FROM cte;"
));
assert!(!sqlserver_batch_can_use_execute("RESTORE HEADERONLY FROM DISK = 'backup.bak' WITH FILE = 1;"));
assert!(sqlserver_batch_can_use_execute(
"DECLARE @output INT = 1; UPDATE dbo.users SET active = @output WHERE id = 1;"
));
}
#[test]
fn sqlserver_transaction_batches_keep_simple_query_path() {
assert!(!sqlserver_batch_can_use_execute("BEGIN TRANSACTION\nUPDATE dbo.users SET active = 0 WHERE id = 1;"));
assert!(!sqlserver_batch_can_use_execute("UPDATE dbo.users SET active = 0 WHERE id = 1;\nCOMMIT;"));
assert!(!sqlserver_batch_can_use_execute("ROLLBACK TRANSACTION;"));
assert!(!sqlserver_batch_can_use_execute(
"BEGIN DISTRIBUTED TRANSACTION\nUPDATE dbo.users SET active = 0 WHERE id = 1;"
));
assert!(!sqlserver_batch_can_use_execute(
"SET IMPLICIT_TRANSACTIONS ON\nUPDATE dbo.users SET active = 0 WHERE id = 1;"
));
assert!(sqlserver_batch_can_use_execute("BEGIN TRY\nUPDATE dbo.users SET active = 0 WHERE id = 1;\nEND TRY"));
assert!(sqlserver_batch_can_use_execute(
"UPDATE dbo.users SET note = 'BEGIN TRANSACTION; ROLLBACK' WHERE id = 1;"
));
assert!(sqlserver_batch_can_use_execute("UPDATE dbo.users SET [rollback] = 1 WHERE id = 1;"));
}
#[test]
fn sqlserver_result_returning_batches_keep_simple_query_path() {
assert!(!sqlserver_batch_can_use_execute("SELECT * FROM dbo.users;"));
assert!(!sqlserver_batch_can_use_execute("EXEC dbo.list_users;"));
assert!(!sqlserver_batch_can_use_execute("DECLARE @id INT = 1; EXEC dbo.list_users @id;"));
assert!(!sqlserver_batch_can_use_execute(
"DECLARE @id INT = 1; CREATE TABLE #t(id INT); INSERT INTO #t VALUES (@id); SELECT id FROM #t;"
));
assert!(!sqlserver_batch_can_use_execute("WITH cte AS (SELECT 1 AS id) SELECT * FROM cte;"));
assert!(!sqlserver_batch_can_use_execute("UPDATE dbo.users SET active = 0 OUTPUT inserted.id WHERE id = 1;"));
assert!(sqlserver_batch_can_use_execute(
"UPDATE dbo.users SET active = 0 OUTPUT inserted.id INTO dbo.audit_ids WHERE id = 1;"
));
assert!(sqlserver_batch_can_use_execute(
"DELETE FROM dbo.users OUTPUT deleted.id INTO @audit_ids WHERE id = 1;"
));
assert!(sqlserver_batch_can_use_execute(
"MERGE dbo.users AS target USING dbo.source AS source ON target.id = source.id WHEN MATCHED THEN DELETE OUTPUT deleted.id INTO dbo.audit_ids;"
));
assert!(!sqlserver_batch_can_use_execute(
"UPDATE dbo.users SET active = 0 OUTPUT inserted.id INTO dbo.audit_ids OUTPUT inserted.id WHERE id = 1;"
));
assert!(sqlserver_batch_can_use_execute(
"DECLARE @id INT = 1; UPDATE dbo.users SET active = 0 WHERE id = @id;"
));
assert!(sqlserver_dml_output_returns_rows("DELETE FROM dbo.users OUTPUT deleted.id WHERE id = 1;"));
}
#[test]
fn sqlserver_user_query_paths_do_not_collect_full_results_before_limiting() {
let source = include_str!("sqlserver.rs");
let execute_query = source.split("pub async fn execute_query").nth(1).unwrap();
let execute_query = execute_query.split("pub async fn execute_batch").next().unwrap();
assert!(!execute_query.contains("into_first_result"));
let execute_batch = source.split("pub async fn execute_batch").nth(1).unwrap();
let execute_batch = execute_batch.split("#[cfg(test)]").next().unwrap();
assert!(!execute_batch.contains("into_results"));
}
#[test]
fn sqlserver_explicit_simple_batch_bypasses_result_type_probing() {
let source = include_str!("sqlserver.rs");
let simple_batch = source.split("pub async fn execute_simple_batch_with_max_rows").nth(1).unwrap();
let simple_batch = simple_batch.split("fn strip_dbx_sqlserver_row_number_column").next().unwrap();
assert!(simple_batch.contains("client.simple_query(sql)"));
assert!(!simple_batch.contains("sqlserver_unsafe_type_query"));
assert!(!simple_batch.contains("describe_sqlserver_result_set"));
}
#[test]
fn sqlserver_index_metadata_sql_avoids_string_agg_for_older_compatibility_levels() {
let sql = sqlserver_indexes_sql("dbo", "DF_Rule");
assert!(!sql.contains("STRING_AGG"));
assert!(sql.contains("FOR XML PATH"));
assert!(sql.contains("OBJECT_ID('dbo.DF_Rule')"));
}
#[test]
fn sqlserver_indexes_sql_includes_index_comment_via_extended_properties() {
let sql = sqlserver_indexes_sql("dbo", "orders");
assert!(sql.contains("sys.extended_properties ep"));
assert!(sql.contains("ep.minor_id = i.index_id"));
assert!(sql.contains("MS_Description"));
}
#[test]
fn sqlserver_columns_sql_reads_column_comment_by_column_id() {
let sql = sqlserver_columns_sql("dbo", "orders");
assert!(sql.contains("FROM sys.objects o"));
assert!(sql.contains("JOIN sys.columns c ON c.object_id = o.object_id"));
assert!(sql.contains("sys.extended_properties ep"));
assert!(sql.contains("ep.major_id = c.object_id"));
assert!(sql.contains("ep.minor_id = c.column_id"));
assert!(sql.contains("MS_Description"));
assert!(sql.contains("c.is_computed = 1 THEN 'computed'"));
}
#[test]
fn sqlserver_columns_sql_exposes_structured_generation_flags() {
let sql = sqlserver_columns_sql("dbo", "orders");
assert!(sql.contains("COLUMNPROPERTY(c.object_id, c.name, 'IsIdentity')"));
assert!(sql.contains("COLUMNPROPERTY(c.object_id, c.name, 'IsComputed')"));
assert!(sql.contains("COLUMNPROPERTY(c.object_id, c.name, 'IsHidden')"));
assert!(sql.contains("COLUMNPROPERTY(c.object_id, c.name, 'GeneratedAlwaysType')"));
}
#[test]
fn sqlserver_table_comment_sql_queries_extended_properties() {
let sql = sqlserver_table_comment_sql("dbo", "users");
assert!(sql.contains("sys.extended_properties ep"));
assert!(sql.contains("ep.minor_id = 0"));
assert!(sql.contains("MS_Description"));
assert!(sql.contains("QUOTENAME('dbo')"));
assert!(sql.contains("QUOTENAME('users')"));
}
#[test]
fn sqlserver_metadata_sql_escapes_literals() {
let columns_sql = sqlserver_columns_sql("d'bo", "t'able");
let indexes_sql = sqlserver_indexes_sql("d'bo", "t'able");
assert!(columns_sql.contains("s.name = N'd''bo'"));
assert!(columns_sql.contains("o.name = 't''able'"));
assert!(columns_sql.contains("sys.identity_columns"));
assert!(indexes_sql.contains("OBJECT_ID('d''bo.t''able')"));
}
#[test]
fn sqlserver_metadata_resolves_blank_schema_to_default_with_dbo_fallback() {
let predicate = sqlserver_schema_name_predicate(" ", "s.name");
assert_eq!(
predicate,
"s.name = COALESCE((SELECT default_schema.name FROM sys.schemas default_schema WHERE default_schema.name = SCHEMA_NAME()), N'dbo')"
);
assert!(sqlserver_list_tables_sql("", None, None, None).contains(&predicate));
assert!(sqlserver_columns_sql("\t", "orders").contains(&predicate));
}
#[test]
fn sqlserver_metadata_preserves_explicit_schema_matching() {
let predicate = sqlserver_schema_name_predicate("Sales'Ops", "s.name");
assert_eq!(predicate, "s.name = N'Sales''Ops'");
assert!(sqlserver_list_tables_sql("Sales'Ops", None, None, None).contains(&predicate));
assert!(sqlserver_columns_sql("Sales'Ops", "orders").contains(&predicate));
assert!(!predicate.contains("SCHEMA_NAME()"));
}
#[test]
fn sqlserver_list_objects_sql_includes_timestamps() {
let sql = sqlserver_list_objects_sql("dbo");
assert!(sql.contains("create_date"));
assert!(sql.contains("modify_date"));
}
#[test]
fn sqlserver_metadata_allows_cdc_system_shipped_objects() {
let predicate = sqlserver_visible_object_predicate();
assert_eq!(predicate, "(o.is_ms_shipped = 0 OR s.name = 'cdc')");
assert!(sqlserver_list_tables_sql("cdc", None, Some(200), None).contains(predicate));
assert!(sqlserver_list_objects_sql("cdc").contains(predicate));
}
#[test]
fn sqlserver_list_schemas_includes_empty_user_schemas() {
let sql = sqlserver_list_schemas_sql();
assert!(!sql.contains("sys.objects"));
assert!(sql.contains("s.name NOT IN"));
assert!(sql.contains("'db_owner'"));
assert!(sql.contains("'db_datareader'"));
assert!(sqlserver_hidden_schema_names().contains(&"sys"));
}
#[test]
fn sqlserver_list_tables_filter_is_case_insensitive() {
let sql = sqlserver_list_tables_sql("dbo", Some("temp"), Some(200), None);
assert!(sql.contains("LOWER(o.name) LIKE LOWER('%temp%') ESCAPE '\\'"));
assert!(sql.contains("LOWER(o.name) LIKE LOWER('%t%e%m%p%') ESCAPE '\\'"));
assert!(sql.contains("SELECT TOP (200)"));
}
#[test]
fn sqlserver_list_tables_filter_escapes_like_literals() {
let sql = sqlserver_list_tables_sql("dbo", Some("Temp_Table[%]"), Some(200), None);
assert!(sql.contains("LOWER(o.name) LIKE LOWER('%Temp\\_Table\\[\\%]%') ESCAPE '\\'"));
assert!(sql.contains("LOWER(o.name) LIKE LOWER('%T%e%m%p%\\_%T%a%b%l%e%\\[%\\%%]%') ESCAPE '\\'"));
}
#[test]
fn sqlserver_list_tables_filter_adds_fuzzy_pattern() {
let sql = sqlserver_list_tables_sql("dbo", Some("sysu"), Some(200), None);
assert!(sql.contains("LOWER(o.name) LIKE LOWER('%sysu%') ESCAPE '\\'"));
assert!(sql.contains("LOWER(o.name) LIKE LOWER('%s%y%s%u%') ESCAPE '\\'"));
assert!(sql.contains("SELECT TOP (200)"));
}
#[test]
fn sqlserver_list_tables_filter_skips_fuzzy_pattern_for_single_character() {
let sql = sqlserver_list_tables_sql("dbo", Some("u"), Some(200), None);
assert!(sql.contains("LOWER(o.name) LIKE LOWER('%u%') ESCAPE '\\'"));
assert!(!sql.contains(" OR LOWER(o.name) LIKE"));
assert!(sql.contains("SELECT TOP (200)"));
}
#[test]
fn sqlserver_completion_assistant_searches_objects_before_limiting() {
let request = CompletionAssistantRequest {
connection_id: "c1".to_string(),
database: "app".to_string(),
schema: Some("dbo".to_string()),
object_kinds: vec![CompletionAssistantObjectKind::Table, CompletionAssistantObjectKind::View],
mask: "Temp".to_string(),
case_sensitive: false,
global_search: false,
max_results: Some(100),
search_in_comments: false,
search_in_definitions: false,
parent_schema: None,
parent_name: None,
match_mode: Some(CompletionAssistantMatchMode::Prefix),
};
let sql = sqlserver_completion_assistant_sql(&request, 100);
assert!(sql.contains("SELECT TOP (100)"));
assert!(sql.contains("FROM sys.objects o"));
assert!(sql.contains("o.type IN ('U','V')"));
assert!(sql.contains("s.name = 'dbo'"));
assert!(sql.contains("LOWER(o.name) LIKE LOWER('Temp%') ESCAPE '\\'"));
assert!(sql.contains("CAST(NULL AS NVARCHAR(128)) AS parent_schema"));
assert!(sql.contains("CAST(NULL AS NVARCHAR(128)) AS parent_name"));
assert!(sql.contains("CAST(NULL AS NVARCHAR(128)) AS data_type"));
}
#[test]
fn sqlserver_completion_assistant_searches_columns_by_parent_table() {
let request = CompletionAssistantRequest {
connection_id: "c1".to_string(),
database: "app".to_string(),
schema: Some("dbo".to_string()),
object_kinds: vec![CompletionAssistantObjectKind::Column],
mask: "id".to_string(),
case_sensitive: false,
global_search: false,
max_results: Some(50),
search_in_comments: false,
search_in_definitions: false,
parent_schema: Some("dbo".to_string()),
parent_name: Some("Users".to_string()),
match_mode: Some(CompletionAssistantMatchMode::Contains),
};
let sql = sqlserver_completion_assistant_sql(&request, 50);
assert!(sql.contains("FROM sys.columns c"));
assert!(sql.contains("o.name = 'Users'"));
assert!(sql.contains("LOWER(c.name) LIKE LOWER('%id%') ESCAPE '\\'"));
assert!(sql.contains("CAST(NULL AS NVARCHAR(MAX)) AS object_comment"));
}
#[test]
fn sqlserver_completion_assistant_returns_function_result_types() {
let request = CompletionAssistantRequest {
connection_id: "c1".to_string(),
database: "app".to_string(),
schema: Some("dbo".to_string()),
object_kinds: vec![CompletionAssistantObjectKind::Routine],
mask: "fn_".to_string(),
case_sensitive: false,
global_search: false,
max_results: Some(50),
search_in_comments: false,
search_in_definitions: false,
parent_schema: Some("dbo".to_string()),
parent_name: None,
match_mode: Some(CompletionAssistantMatchMode::Prefix),
};
let sql = sqlserver_completion_assistant_sql(&request, 50);
assert!(sql.contains("WHEN o.type IN ('IF','TF','FT') THEN 'table'"));
assert!(sql.contains("p.parameter_id = 0"));
assert!(sql.contains("TYPE_NAME(p.user_type_id)"));
}
#[test]
fn sqlserver_completion_assistant_searches_tempdb_for_temp_table_masks() {
let request = CompletionAssistantRequest {
connection_id: "c1".to_string(),
database: "app".to_string(),
schema: Some("dbo".to_string()),
object_kinds: vec![CompletionAssistantObjectKind::Table],
mask: "#Temp".to_string(),
case_sensitive: false,
global_search: false,
max_results: Some(100),
search_in_comments: false,
search_in_definitions: false,
parent_schema: None,
parent_name: None,
match_mode: Some(CompletionAssistantMatchMode::Prefix),
};
let sql = sqlserver_completion_assistant_sql(&request, 100);
assert!(sql.contains("FROM tempdb.sys.all_objects o"));
assert!(sql.contains("o.type = 'U'"));
assert!(sql.contains("LOWER(o.name) LIKE LOWER('#Temp%') ESCAPE '\\'"));
assert!(sql.contains("CAST(NULL AS NVARCHAR(128)) AS parent_schema"));
assert!(sql.contains("CAST(NULL AS NVARCHAR(MAX)) AS object_comment"));
}
#[test]
fn sqlserver_completion_assistant_generates_scoped_search_masks() {
assert_eq!(super::completion_like_pattern("Temp", Some(&CompletionAssistantMatchMode::Prefix)), "Temp%");
assert_eq!(super::completion_like_pattern("Temp", Some(&CompletionAssistantMatchMode::Contains)), "%Temp%");
assert_eq!(
super::completion_like_pattern("dbo.Temp%", Some(&CompletionAssistantMatchMode::Prefix)),
"dbo.Temp%"
);
assert_eq!(
super::completion_like_pattern("Temp_Table", Some(&CompletionAssistantMatchMode::Prefix)),
"Temp\\_Table%"
);
}
#[test]
fn sqlserver_completion_assistant_can_search_comments_and_definitions() {
let request = CompletionAssistantRequest {
connection_id: "c1".to_string(),
database: "app".to_string(),
schema: Some("dbo".to_string()),
object_kinds: vec![CompletionAssistantObjectKind::Procedure],
mask: "audit".to_string(),
case_sensitive: false,
global_search: false,
max_results: Some(100),
search_in_comments: true,
search_in_definitions: true,
parent_schema: None,
parent_name: None,
match_mode: Some(CompletionAssistantMatchMode::Contains),
};
let sql = sqlserver_completion_assistant_sql(&request, 100);
assert!(sql.contains("COALESCE(ep.value, '')"));
assert!(sql.contains("OBJECT_DEFINITION(o.object_id)"));
assert!(sql.contains("LOWER('%audit%')"));
}
#[test]
fn sqlserver_completion_assistant_casts_schema_and_empty_result_placeholders() {
let schema_request = CompletionAssistantRequest {
connection_id: "c1".to_string(),
database: "app".to_string(),
schema: None,
object_kinds: vec![CompletionAssistantObjectKind::Schema],
mask: "d".to_string(),
case_sensitive: false,
global_search: false,
max_results: Some(100),
search_in_comments: false,
search_in_definitions: false,
parent_schema: None,
parent_name: None,
match_mode: Some(CompletionAssistantMatchMode::Prefix),
};
let schema_sql = sqlserver_completion_assistant_sql(&schema_request, 100);
assert!(schema_sql.contains("CAST(NULL AS NVARCHAR(128)) AS parent_schema"));
assert!(schema_sql.contains("CAST(NULL AS NVARCHAR(128)) AS parent_name"));
assert!(schema_sql.contains("CAST(NULL AS NVARCHAR(MAX)) AS object_comment"));
assert!(schema_sql.contains("CAST(NULL AS NVARCHAR(128)) AS data_type"));
let empty_request = CompletionAssistantRequest {
object_kinds: vec![CompletionAssistantObjectKind::Database],
..schema_request
};
let empty_sql = sqlserver_completion_assistant_sql(&empty_request, 100);
assert!(empty_sql.contains("CAST('' AS NVARCHAR(128)) AS name"));
assert!(empty_sql.contains("CAST(NULL AS NVARCHAR(128)) AS parent_schema"));
assert!(empty_sql.contains("CAST(NULL AS NVARCHAR(MAX)) AS object_comment"));
}
#[test]
fn sqlserver_tinyint_cells_are_json_numbers() {
assert_eq!(sqlserver_cell_to_json(&ColumnData::U8(Some(7))), serde_json::json!(7));
}
#[test]
fn sqlserver_numeric_cells_with_scale_over_28_do_not_panic() {
// NUMERIC(38, 29) with data used to abort the app: rust_decimal caps scale at 28 (issue #3648).
let cell = ColumnData::Numeric(Some(tiberius::numeric::Numeric::new_with_scale(5, 29)));
assert_eq!(sqlserver_cell_to_json(&cell), serde_json::json!("0.00000000000000000000000000005"));
}
#[test]
fn sqlserver_numeric_cells_beyond_96_bit_mantissa_do_not_panic() {
// Precision-38 values overflow rust_decimal's 96-bit mantissa even at low scale.
let value: i128 = 12_345_678_901_234_567_890_123_456_789_012_345_678;
let cell = ColumnData::Numeric(Some(tiberius::numeric::Numeric::new_with_scale(value, 10)));
assert_eq!(sqlserver_cell_to_json(&cell), serde_json::json!("1234567890123456789012345678.9012345678"));
}
#[test]
fn sqlserver_null_numeric_cells_are_json_null() {
assert_eq!(sqlserver_cell_to_json(&ColumnData::Numeric(None)), serde_json::Value::Null);
}
#[test]
fn format_sqlserver_numeric_covers_sign_scale_and_padding() {
assert_eq!(format_sqlserver_numeric(42, 0), "42");
assert_eq!(format_sqlserver_numeric(-42, 0), "-42");
assert_eq!(format_sqlserver_numeric(12345, 2), "123.45");
// Trailing zeros are kept, matching the previous rust_decimal display.
assert_eq!(format_sqlserver_numeric(1500, 3), "1.500");
assert_eq!(format_sqlserver_numeric(-15, 1), "-1.5");
assert_eq!(format_sqlserver_numeric(-5, 2), "-0.05");
assert_eq!(format_sqlserver_numeric(0, 2), "0.00");
// digits.len() == scale must keep the leading "0." (guards `>` vs `>=` in the split).
assert_eq!(format_sqlserver_numeric(123, 3), "0.123");
// Largest scale tiberius can deliver (its decoder asserts scale < 38).
assert_eq!(format_sqlserver_numeric(9, 37), format!("0.{}9", "0".repeat(36)));
}
#[test]
fn sqlserver_datetime2_cells_are_json_strings() {
let datetime = NaiveDate::from_ymd_opt(2026, 5, 13).unwrap().and_hms_milli_opt(9, 8, 7, 123).unwrap();
let cell: ColumnData<'static> = datetime.into_sql();
assert_eq!(sqlserver_cell_to_json(&cell), serde_json::json!("2026-05-13 09:08:07.123"));
}
#[test]
fn sqlserver_datetime_cells_display_millisecond_precision() {
let cell = ColumnData::DateTime(Some(tiberius::time::DateTime::new(46_200, 11_001_869)));
assert_eq!(sqlserver_cell_to_json(&cell), serde_json::json!("2026-06-29 10:11:12.897"));
}
#[test]
fn sqlserver_binary_cells_are_json_hex_strings() {
let cell =
ColumnData::Binary(Some(std::borrow::Cow::Owned(vec![0x00, 0x00, 0x00, 0x00, 0x00, 0x01, 0xCF, 0x53])));
assert_eq!(sqlserver_cell_to_json(&cell), serde_json::json!("0x000000000001cf53"));
}
#[test]
fn sqlserver_strips_generated_row_number_pagination_column() {
let sql = "SELECT * FROM (SELECT dbx_page_source.*, ROW_NUMBER() OVER (ORDER BY (SELECT NULL)) AS [__dbx_row_num] FROM (SELECT id FROM users) dbx_page_source) dbx_page WHERE [__dbx_row_num] > 100 AND [__dbx_row_num] <= 200 ORDER BY [__dbx_row_num];";
let mut result = QueryResult {
columns: vec!["id".to_string(), "__dbx_row_num".to_string()],
column_types: vec!["int".to_string(), "bigint".to_string()],
column_sortables: vec![],
rows: vec![vec![serde_json::json!(42), serde_json::json!(101)]],
affected_rows: 0,
execution_time_ms: 1,
truncated: false,
session_id: None,
has_more: false,
elasticsearch_raw_body: None,
};
strip_dbx_sqlserver_row_number_column(&mut result, sql);
assert_eq!(result.columns, vec!["id"]);
assert_eq!(result.column_types, vec!["int"]);
assert_eq!(result.rows, vec![vec![serde_json::json!(42)]]);
}
#[test]
fn sqlserver_detects_geometry_result_columns() {
assert!(is_sqlserver_spatial_column(&SqlServerDescribedColumn {
name: Some("polygon".to_string()),
system_type_name: Some("geometry".to_string()),
user_type_schema: Some("sys".to_string()),
user_type_name: Some("geometry".to_string()),
}));
assert!(is_sqlserver_spatial_column(&SqlServerDescribedColumn {
name: Some("shape".to_string()),
system_type_name: Some("geography".to_string()),
user_type_schema: Some("sys".to_string()),
user_type_name: Some("geography".to_string()),
}));
assert!(!is_sqlserver_spatial_column(&SqlServerDescribedColumn {
name: Some("name".to_string()),
system_type_name: Some("varchar(30)".to_string()),
user_type_schema: None,
user_type_name: None,
}));
}
#[test]
fn sqlserver_wraps_geometry_columns_as_text() {
let rewritten = build_sqlserver_unsafe_type_query(
"SELECT * FROM dbo.tLandPolygon;",
&[
SqlServerDescribedColumn {
name: Some("landId".to_string()),
system_type_name: Some("varchar(30)".to_string()),
user_type_schema: None,
user_type_name: None,
},
SqlServerDescribedColumn {
name: Some("polygon".to_string()),
system_type_name: Some("geometry".to_string()),
user_type_schema: Some("sys".to_string()),
user_type_name: Some("geometry".to_string()),
},
],
)
.unwrap();
assert_eq!(
rewritten,
"SELECT [landId] = [dbx_unsafe_source].[dbx_col_1], [polygon] = CASE WHEN [dbx_unsafe_source].[dbx_col_2] IS NULL THEN NULL ELSE [dbx_unsafe_source].[dbx_col_2].STAsText() END FROM (SELECT * FROM dbo.tLandPolygon) AS [dbx_unsafe_source]([dbx_col_1], [dbx_col_2])"
);
}
#[test]
fn sqlserver_does_not_wrap_non_spatial_columns() {
assert_eq!(
build_sqlserver_unsafe_type_query(
"SELECT landId FROM dbo.tLandPolygon",
&[SqlServerDescribedColumn {
name: Some("landId".to_string()),
system_type_name: Some("varchar(30)".to_string()),
user_type_schema: None,
user_type_name: None,
}]
),
None
);
}
#[test]
fn sqlserver_preserves_order_by_when_wrapping_geometry_columns() {
let rewritten = build_sqlserver_unsafe_type_query(
"SELECT landId, polygon FROM dbo.tLandPolygon ORDER BY landId DESC",
&[
SqlServerDescribedColumn {
name: Some("landId".to_string()),
system_type_name: Some("varchar(30)".to_string()),
user_type_schema: None,
user_type_name: None,
},
SqlServerDescribedColumn {
name: Some("polygon".to_string()),
system_type_name: Some("geometry".to_string()),
user_type_schema: Some("sys".to_string()),
user_type_name: Some("geometry".to_string()),
},
],
)
.unwrap();
// ORDER BY is stripped from the inner query so it can be used as a
// derived table subquery across all SQL Server versions (2008–2022).
assert!(!rewritten.contains("ORDER BY"));
assert!(rewritten.contains("FROM dbo.tLandPolygon"));
assert!(rewritten.contains(".STAsText()"));
}
#[test]
fn sqlserver_keeps_empty_result_sets_when_metadata_exists() {
let mut results = Vec::new();
super::push_sqlserver_result_set(
&mut results,
Some(SqlServerResultSet {
columns: vec!["id".to_string(), "name".to_string()],
column_types: vec![],
rows: vec![],
truncated: false,
}),
Instant::now(),
);
assert_eq!(results.len(), 1);
assert_eq!(results[0].columns, vec!["id".to_string(), "name".to_string()]);
assert!(results[0].rows.is_empty());
}
#[test]
fn sqlserver_drops_truly_empty_result_sets_without_metadata() {
let mut results = Vec::new();
super::push_sqlserver_result_set(
&mut results,
Some(SqlServerResultSet { columns: vec![], column_types: vec![], rows: vec![], truncated: false }),
Instant::now(),
);
assert!(results.is_empty());
}
#[test]
fn sqlserver_detects_sql_variant_columns() {
assert!(is_sqlserver_variant_column(&SqlServerDescribedColumn {
name: Some("value".to_string()),
system_type_name: Some("sql_variant".to_string()),
user_type_schema: None,
user_type_name: None,
}));
assert!(is_sqlserver_variant_column(&SqlServerDescribedColumn {
name: Some("value".to_string()),
system_type_name: None,
user_type_schema: None,
user_type_name: Some("sql_variant".to_string()),
}));
assert!(!is_sqlserver_variant_column(&SqlServerDescribedColumn {
name: Some("name".to_string()),
system_type_name: Some("nvarchar(128)".to_string()),
user_type_schema: None,
user_type_name: None,
}));
}
#[test]
fn sqlserver_result_type_probe_keeps_modern_and_legacy_paths_in_one_round_trip() {
assert!(SQLSERVER_RESULT_TYPE_PROBE_SQL.contains("sys.dm_exec_describe_first_result_set"));
assert!(SQLSERVER_RESULT_TYPE_PROBE_SQL.contains("##dbx_result_type_probe_"));
assert!(SQLSERVER_RESULT_TYPE_PROBE_SQL.contains("SELECT TOP (0) * INTO"));
assert!(SQLSERVER_RESULT_TYPE_PROBE_SQL.contains("N' FROM ' + @P3"));
assert!(SQLSERVER_RESULT_TYPE_PROBE_SQL.contains("FROM tempdb.sys.columns"));
assert!(!SQLSERVER_RESULT_TYPE_PROBE_SQL.contains("FMTONLY"));
assert_eq!(SQLSERVER_RESULT_TYPE_PROBE_SQL.matches("SELECT @dbx_use_describe_dmv").count(), 1);
}
#[test]
fn sqlserver_legacy_probe_uses_unique_internal_names_for_duplicate_outputs() {
let probe = sqlserver_legacy_probe(
"SELECT a.HJRQ, b.HJRQ, a.id + b.id AS total, GETDATE() FROM dbo.a a JOIN dbo.b b ON b.id = a.id",
)
.unwrap();
assert_eq!(
probe.source_sql,
"(SELECT a.HJRQ, b.HJRQ, a.id + b.id AS total, GETDATE() FROM dbo.a a JOIN dbo.b b ON b.id = a.id) AS [dbx_probe_source]([dbx_col_1], [dbx_col_2], [dbx_col_3], [dbx_col_4])"
);
assert_eq!(
probe.output_names,
Some(vec![Some("HJRQ".to_string()), Some("HJRQ".to_string()), Some("total".to_string()), None])
);
}
#[test]
fn sqlserver_legacy_probe_keeps_wildcards_on_existing_server_validation_path() {
let probe = sqlserver_legacy_probe("SELECT a.*, b.HJRQ FROM dbo.a a JOIN dbo.b b ON b.id = a.id").unwrap();
assert_eq!(
probe.source_sql,
"(SELECT a.*, b.HJRQ FROM dbo.a a JOIN dbo.b b ON b.id = a.id) AS [dbx_probe_source]"
);
assert_eq!(probe.output_names, None);
}
#[test]
fn sqlserver_legacy_probe_errors_are_blocking_only_when_marked_unsafe() {
assert!(is_blocking_sqlserver_unsafe_probe_error(
"SQL Server unsafe result type: legacy schema capture failed"
));
assert!(!is_blocking_sqlserver_unsafe_probe_error("Invalid object name"));
}
#[test]
fn sqlserver_wraps_sql_variant_columns_as_nvarchar() {
let rewritten = build_sqlserver_unsafe_type_query(
"SELECT name, value FROM sys.extended_properties;",
&[
SqlServerDescribedColumn {
name: Some("name".to_string()),
system_type_name: Some("sysname".to_string()),
user_type_schema: None,
user_type_name: None,
},
SqlServerDescribedColumn {
name: Some("value".to_string()),
system_type_name: Some("sql_variant".to_string()),
user_type_schema: None,
user_type_name: None,
},
],
)
.unwrap();
assert!(rewritten.contains("CAST("));
assert!(rewritten.contains("AS NVARCHAR(MAX))"));
assert!(rewritten.contains("FROM sys.extended_properties"));
// The name column should not be cast
assert_eq!(rewritten.matches("CAST(").count(), 1);
}
#[test]
fn sqlserver_does_not_wrap_non_variant_columns() {
assert_eq!(
build_sqlserver_unsafe_type_query(
"SELECT name FROM sys.extended_properties",
&[SqlServerDescribedColumn {
name: Some("name".to_string()),
system_type_name: Some("sysname".to_string()),
user_type_schema: None,
user_type_name: None,
}]
),
None
);
}
#[test]
fn sqlserver_wraps_both_spatial_and_variant_columns() {
let rewritten = build_sqlserver_unsafe_type_query(
"SELECT id, shape, metadata FROM dbo.t;",
&[
SqlServerDescribedColumn {
name: Some("id".to_string()),
system_type_name: Some("int".to_string()),
user_type_schema: None,
user_type_name: None,
},
SqlServerDescribedColumn {
name: Some("shape".to_string()),
system_type_name: Some("geometry".to_string()),
user_type_schema: Some("sys".to_string()),
user_type_name: Some("geometry".to_string()),
},
SqlServerDescribedColumn {
name: Some("metadata".to_string()),
system_type_name: Some("sql_variant".to_string()),
user_type_schema: None,
user_type_name: None,
},
],
)
.unwrap();
assert!(rewritten.contains(".STAsText()"));
assert!(rewritten.contains("CAST("));
assert!(rewritten.contains("AS NVARCHAR(MAX))"));
assert!(rewritten.contains("FROM dbo.t"));
}
#[tokio::test]
#[ignore = "requires DBX_LIVE_SQLSERVER_HOST/PORT/USER/PASSWORD"]
async fn live_sqlserver_legacy_probe_casts_variant_and_keeps_connection_usable() {
let host = std::env::var("DBX_LIVE_SQLSERVER_HOST").expect("DBX_LIVE_SQLSERVER_HOST");
let port = std::env::var("DBX_LIVE_SQLSERVER_PORT")
.expect("DBX_LIVE_SQLSERVER_PORT")
.parse::<u16>()
.expect("valid DBX_LIVE_SQLSERVER_PORT");
let user = std::env::var("DBX_LIVE_SQLSERVER_USER").expect("DBX_LIVE_SQLSERVER_USER");
let password = std::env::var("DBX_LIVE_SQLSERVER_PASSWORD").expect("DBX_LIVE_SQLSERVER_PASSWORD");
let mut client =
super::connect(&host, port, &user, &password, Some("tempdb"), None, std::time::Duration::from_secs(10))
.await
.unwrap();
let setup = "\
IF OBJECT_ID('tempdb..#dbx_issue_4002') IS NOT NULL DROP TABLE #dbx_issue_4002; \
CREATE TABLE #dbx_issue_4002 (id int NOT NULL, payload sql_variant NULL); \
INSERT INTO #dbx_issue_4002 (id, payload) VALUES (1, CAST(N'legacy' AS nvarchar(20)))";
client.simple_query(setup).await.unwrap().into_results().await.unwrap();
let sql = "SELECT id, payload FROM #dbx_issue_4002";
let ordinary_probe = super::sqlserver_legacy_probe("SELECT 42 AS answer").unwrap();
let ordinary_columns =
super::describe_sqlserver_result_set_with_mode(&mut client, "SELECT 42 AS answer", &ordinary_probe, true)
.await
.unwrap();
assert_eq!(ordinary_columns[0].system_type_name.as_deref(), Some("int"));
let legacy_probe = super::sqlserver_legacy_probe(sql).unwrap();
let legacy_columns =
super::describe_sqlserver_result_set_with_mode(&mut client, sql, &legacy_probe, true).await.unwrap();
assert_eq!(legacy_columns.len(), 2);
assert!(is_sqlserver_variant_column(&legacy_columns[1]));
let rewritten = build_sqlserver_unsafe_type_query(sql, &legacy_columns).unwrap();
let legacy_rows = client.query(rewritten, &[]).await.unwrap().into_first_result().await.unwrap();
assert_eq!(legacy_rows[0].get::<i32, _>(0), Some(1));
assert_eq!(legacy_rows[0].get::<&str, _>(1), Some("legacy"));
let ordinary = super::execute_query(&mut client, "SELECT CAST(42 AS int) AS answer").await.unwrap();
assert_eq!(ordinary.rows, vec![vec![serde_json::json!(42)]]);
let variant = super::execute_query(&mut client, sql).await.unwrap();
assert_eq!(variant.rows, vec![vec![serde_json::json!(1), serde_json::json!("legacy")]]);
let duplicate_sql = "SELECT id AS HJRQ, payload AS HJRQ FROM #dbx_issue_4002";
let duplicate_probe = super::sqlserver_legacy_probe(duplicate_sql).unwrap();
let duplicate_columns =
super::describe_sqlserver_result_set_with_mode(&mut client, duplicate_sql, &duplicate_probe, true)
.await
.unwrap();
assert_eq!(duplicate_columns.len(), 2);
assert_eq!(duplicate_columns[0].name.as_deref(), Some("HJRQ"));
assert_eq!(duplicate_columns[1].name.as_deref(), Some("HJRQ"));
assert!(is_sqlserver_variant_column(&duplicate_columns[1]));
let duplicate_rewritten = build_sqlserver_unsafe_type_query(duplicate_sql, &duplicate_columns).unwrap();
let duplicate_rows = client.query(duplicate_rewritten, &[]).await.unwrap().into_first_result().await.unwrap();
assert_eq!(duplicate_rows[0].columns()[0].name(), "HJRQ");
assert_eq!(duplicate_rows[0].columns()[1].name(), "HJRQ");
assert_eq!(duplicate_rows[0].get::<i32, _>(0), Some(1));
assert_eq!(duplicate_rows[0].get::<&str, _>(1), Some("legacy"));
let continued = super::execute_query(&mut client, "SELECT CAST(7 AS int) AS still_connected").await.unwrap();
assert_eq!(continued.rows, vec![vec![serde_json::json!(7)]]);
client.simple_query("DROP TABLE #dbx_issue_4002").await.unwrap().into_results().await.unwrap();
}
}
Reference in New Issue View Git Blame Copy Permalink