fix(sqlserver): honor USE context in completions

crates/dbx-core/src/db/sqlserver.rs

@@ -92,6 +92,56 @@ pub struct SqlServerColumnMetadata {
     pub generated_always_type: i32,
 }
 
+#[derive(Debug, Clone, serde::Serialize, serde::Deserialize, PartialEq, Eq)]
+pub struct SqlServerCompletionContext {
+    pub default_schema: String,
+    pub supports_session_database_switch: bool,
+}
+
+const SQLSERVER_COMPLETION_CONTEXT_SQL: &str = "\
+    SELECT COALESCE(\
+        (SELECT default_schema.name \
+         FROM sys.schemas default_schema \
+         WHERE default_schema.name = SCHEMA_NAME()), \
+        N'dbo'\
+    ) AS default_schema, \
+    CONVERT(int, SERVERPROPERTY(N'EngineEdition')) AS engine_edition";
+
+fn sqlserver_supports_session_database_switch(engine_edition: i32) -> bool {
+    // Only known boxed SQL Server, Managed Instance, and SQL Edge editions are
+    // allowed to switch databases. Cloud single-database endpoints and future
+    // editions default to opening a connection directly to the target database.
+    matches!(engine_edition, 1 | 2 | 3 | 4 | 8 | 9)
+}
+
+fn sqlserver_completion_context(
+    default_schema: Option<&str>,
+    engine_edition: Option<i32>,
+) -> Result<SqlServerCompletionContext, String> {
+    let default_schema = default_schema.map(str::trim).filter(|schema| !schema.is_empty()).unwrap_or("dbo");
+    let engine_edition = engine_edition.ok_or_else(|| "SQL Server EngineEdition is unavailable".to_string())?;
+    Ok(SqlServerCompletionContext {
+        default_schema: default_schema.to_string(),
+        supports_session_database_switch: sqlserver_supports_session_database_switch(engine_edition),
+    })
+}
+
+pub fn completion_context_sql() -> &'static str {
+    SQLSERVER_COMPLETION_CONTEXT_SQL
+}
+
+pub fn completion_context_from_query_result(result: QueryResult) -> Result<SqlServerCompletionContext, String> {
+    let row = result.rows.first().ok_or_else(|| "SQL Server completion context query returned no rows".to_string())?;
+    let default_schema = row.first().and_then(serde_json::Value::as_str);
+    let engine_edition = row.get(1).and_then(|value| {
+        value
+            .as_i64()
+            .and_then(|value| i32::try_from(value).ok())
+            .or_else(|| value.as_str()?.trim().parse::<i32>().ok())
+    });
+    sqlserver_completion_context(default_schema, engine_edition)
+}
+
 #[derive(Debug, PartialEq, Eq)]
 struct SqlServerEndpoint<'a> {
     host: &'a str,
@@ -1133,6 +1183,15 @@ pub async fn list_databases(client: &mut SqlServerClient) -> Result<Vec<Database
     Ok(rows.iter().map(|row| DatabaseInfo { name: row.get::<&str, _>(0).unwrap_or("").to_string() }).collect())
 }
 
+pub async fn get_completion_context(client: &mut SqlServerClient) -> Result<SqlServerCompletionContext, String> {
+    let stream = client.query(SQLSERVER_COMPLETION_CONTEXT_SQL, &[]).await.map_err(|e| e.to_string())?;
+    let rows = stream.into_first_result().await.map_err(|e| e.to_string())?;
+    let row = rows.first().ok_or_else(|| "SQL Server completion context query returned no rows".to_string())?;
+    let default_schema = row.try_get::<&str, _>(0).map_err(|e| e.to_string())?;
+    let engine_edition = row.try_get::<i32, _>(1).map_err(|e| e.to_string())?;
+    sqlserver_completion_context(default_schema, engine_edition)
+}
+
 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())?;
@@ -2503,17 +2562,18 @@ fn first_sql_tokens(sql: &str, limit: usize) -> Vec<String> {
 #[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, restore_sqlserver_legacy_probe_output_names,
-        sqlserver_batch_can_use_execute, sqlserver_bulk_token_row, sqlserver_cell_to_json, sqlserver_columns_sql,
-        sqlserver_completion_assistant_sql, sqlserver_dml_output_returns_rows, sqlserver_filter_definition_error,
-        sqlserver_hidden_schema_names, sqlserver_indexes_sql, sqlserver_legacy_indexes_sql, sqlserver_legacy_probe,
-        sqlserver_legacy_probe_with_nonce, sqlserver_list_objects_sql, sqlserver_list_schemas_sql,
-        sqlserver_list_tables_sql, sqlserver_probe_explicit_alias, sqlserver_schema_name_predicate,
+        build_sqlserver_unsafe_type_query, capture_sqlserver_messages, completion_context_from_query_result,
+        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,
+        restore_sqlserver_legacy_probe_output_names, sqlserver_batch_can_use_execute, sqlserver_bulk_token_row,
+        sqlserver_cell_to_json, sqlserver_columns_sql, sqlserver_completion_assistant_sql,
+        sqlserver_dml_output_returns_rows, sqlserver_filter_definition_error, sqlserver_hidden_schema_names,
+        sqlserver_indexes_sql, sqlserver_legacy_indexes_sql, sqlserver_legacy_probe, sqlserver_legacy_probe_with_nonce,
+        sqlserver_list_objects_sql, sqlserver_list_schemas_sql, sqlserver_list_tables_sql,
+        sqlserver_probe_explicit_alias, sqlserver_schema_name_predicate, sqlserver_supports_session_database_switch,
         sqlserver_table_comment_sql, sqlserver_triggers_sql, sqlserver_visible_object_predicate,
         strip_dbx_sqlserver_row_number_column, SqlServerDescribedColumn, SqlServerProbeOutputNameOverride,
-        SqlServerResultSet, SQLSERVER_RESULT_TYPE_PROBE_SQL,
+        SqlServerResultSet, SQLSERVER_COMPLETION_CONTEXT_SQL, SQLSERVER_RESULT_TYPE_PROBE_SQL,
     };
     use crate::types::{
         CompletionAssistantMatchMode, CompletionAssistantObjectKind, CompletionAssistantRequest, QueryResult,
@@ -2965,6 +3025,46 @@ mod tests {
         assert!(sqlserver_indexes_sql("", "orders").contains("OBJECT_ID(QUOTENAME(N'orders'))"));
     }
 
+    #[test]
+    fn sqlserver_completion_context_uses_the_database_user_default_schema() {
+        assert!(SQLSERVER_COMPLETION_CONTEXT_SQL.contains("SCHEMA_NAME()"));
+        assert!(SQLSERVER_COMPLETION_CONTEXT_SQL.contains("sys.schemas"));
+        assert!(SQLSERVER_COMPLETION_CONTEXT_SQL.contains("N'dbo'"));
+        assert!(SQLSERVER_COMPLETION_CONTEXT_SQL.contains("EngineEdition"));
+    }
+
+    #[test]
+    fn sqlserver_completion_context_disables_use_for_azure_database_endpoints() {
+        assert!(!sqlserver_supports_session_database_switch(5));
+        assert!(!sqlserver_supports_session_database_switch(6));
+        assert!(!sqlserver_supports_session_database_switch(11));
+        assert!(!sqlserver_supports_session_database_switch(12));
+        assert!(!sqlserver_supports_session_database_switch(99));
+        assert!(sqlserver_supports_session_database_switch(3));
+        assert!(sqlserver_supports_session_database_switch(8));
+        assert!(sqlserver_supports_session_database_switch(9));
+    }
+
+    #[test]
+    fn sqlserver_completion_context_parses_agent_query_results() {
+        let context = completion_context_from_query_result(QueryResult {
+            columns: vec!["default_schema".to_string(), "engine_edition".to_string()],
+            column_types: vec![],
+            column_sortables: vec![],
+            rows: vec![vec![serde_json::json!("app_user"), serde_json::json!("8")]],
+            affected_rows: 0,
+            execution_time_ms: 0,
+            truncated: false,
+            session_id: None,
+            has_more: false,
+            elasticsearch_raw_body: None,
+        })
+        .unwrap();
+
+        assert_eq!(context.default_schema, "app_user");
+        assert!(context.supports_session_database_switch);
+    }
+
     #[test]
     fn sqlserver_metadata_preserves_explicit_schema_matching() {
         let predicate = sqlserver_schema_name_predicate("Sales'Ops", "s.name");