fix(query): preserve PG notices for statements without results

2026-10-21 04:43:49 +00:00 · 2026-09-15 14:30:59 +08:00
parent 9c22856285
commit 20e026eea9
3 changed files with 75 additions and 24 deletions
@@ -172,7 +172,7 @@ import ToolbarOverflowMenu from "@/components/ui/ToolbarOverflowMenu.vue";
 import { formatElapsedSeconds } from "@/lib/common/elapsedTime";
 import { copyToClipboard } from "@/lib/common/clipboard";
 import type { CustomSaveHandler } from "@/composables/useDataGridEditor";
-import type { QueryTab, TableInfoTab, TreeNode, VectorCollectionMeta } from "@/types/database";
+import type { QueryMessage, QueryTab, TableInfoTab, TreeNode, VectorCollectionMeta } from "@/types/database";
 import type { SqlObjectNavigationTarget } from "@/lib/sql/sqlNavigation";
 import { sqlFormatDialectForDbType, type SqlFormatDialect } from "@/lib/sql/sqlFormatter";
 import { productionContextForDatabase } from "@/lib/database/productionSafety";
@@ -558,7 +558,16 @@ const hasTabularResult = computed(() => {
 });
 const canShowResultOutput = computed(() => hasTabularResult.value || props.activeTab.isExecuting);
 const canShowExplainOutput = computed(() => !!props.activeTab.explainPlan || !!props.activeTab.explainError || !!props.activeTab.explainTableResult || !!props.activeTab.explainTableError || props.activeTab.isExplaining === true);
-const resultMessageCount = computed(() => props.activeTab.result?.messages?.length ?? 0);
+// A batch can attach server messages to more than one statement result (for
+// example a `DO $$ RAISE NOTICE $$` block followed by a SELECT). The messages
+// view is scoped to the whole run, so collect every result's messages instead
+// of only the active result's — otherwise a statement with no result set
+// silently loses its notices once a later statement owns the active result.
+const resultMessages = computed<QueryMessage[]>(() => {
+  const results = props.activeTab.results?.length ? props.activeTab.results : props.activeTab.result ? [props.activeTab.result] : [];
+  return results.flatMap((result) => result.messages ?? []);
+});
+const resultMessageCount = computed(() => resultMessages.value.length);
 const canShowMessagesOutput = computed(() => resultMessageCount.value > 0);
 const showStandaloneResultToolbar = computed(() => activeElasticsearchJsonResponse.value || props.activeOutputView !== "result" || !props.activeTab.result || !hasTabularResult.value);
 const standaloneResultToolbarCompact = computed(() => isDataGridToolbarCompact(standaloneResultToolbarWidth.value, standaloneResultToolbarViewportWidth.value));
@@ -1900,7 +1909,7 @@ defineExpose({
              </div>
            </div>

-            <QueryMessagesView v-else-if="activeOutputView === 'messages'" class="flex-1 min-h-0" :messages="activeTab.result?.messages ?? []" />
+            <QueryMessagesView v-else-if="activeOutputView === 'messages'" class="flex-1 min-h-0" :messages="resultMessages" />

            <template v-else>
              <ElasticsearchJsonResponsePanel v-if="activeElasticsearchJsonResponse" ref="elasticsearchJsonResponsePanelRef" class="flex-1 min-h-0" :status="activeElasticsearchJsonResponse.status" :body="activeElasticsearchJsonResponse.body" />
@@ -7551,29 +7551,46 @@ pub async fn execute_query_with_max_rows(
    let start = Instant::now();
    let row_limit = query_result_row_limit(max_rows);

-    if postgres_statement_returns_rows(sql) {
-        let client = checkout_postgres_client(pool, None, super::connection_timeout()).await?;
+    let client = checkout_postgres_client(pool, None, super::connection_timeout()).await?;
+    // Drop stale notices from infrastructure statements so only messages raised
+    // by this statement are attached to its result.
+    let _ = drain_postgres_notices(&client).await;
+
+    let result = if postgres_statement_returns_rows(sql) {
        execute_select_query(&client, sql, start, row_limit).await
    } else {
-        let client = checkout_postgres_client(pool, None, super::connection_timeout()).await?;
-        let affected = client.execute(sql, &[]).await.map_err(pg_error_to_string)?;
-        clear_postgres_caches_after_ddl(pool, Some(&client), sql);
+        client.execute(sql, &[]).await.map_err(pg_error_to_string).map(|affected| {
+            clear_postgres_caches_after_ddl(pool, Some(&client), sql);

-        Ok(QueryResult {
-            columns: vec![],
-            column_types: Vec::new(),
-            column_sortables: Vec::new(),
-            spatial_columns: vec![],
-            spatial_values: vec![],
-            rows: vec![],
-            affected_rows: affected,
-            execution_time_ms: start.elapsed().as_millis(),
-            truncated: false,
-            session_id: None,
-            has_more: false,
-            elasticsearch_raw_body: None,
-            messages: Vec::new(),
+            QueryResult {
+                columns: vec![],
+                column_types: Vec::new(),
+                column_sortables: Vec::new(),
+                spatial_columns: vec![],
+                spatial_values: vec![],
+                rows: vec![],
+                affected_rows: affected,
+                execution_time_ms: start.elapsed().as_millis(),
+                truncated: false,
+                session_id: None,
+                has_more: false,
+                elasticsearch_raw_body: None,
+                messages: Vec::new(),
+            }
        })
+    };
+
+    match result {
+        Ok(mut result) => {
+            result.messages = drain_postgres_notices(&client).await;
+            Ok(result)
+        }
+        Err(error) => {
+            // Drop notices so an errored statement's messages cannot leak into
+            // the next query on this pooled connection.
+            let _ = drain_postgres_notices(&client).await;
+            Err(error)
+        }
    }
 }

@@ -10048,6 +10065,25 @@ mod tests {
        assert!(result.messages.iter().any(|message| message.message == "dbx notice identity regression"));
    }

+    #[tokio::test]
+    #[ignore = "requires DBX_TEST_POSTGRES_URL pointing at a PostgreSQL database"]
+    async fn postgres_command_query_preserves_notice_capture() {
+        let url = std::env::var("DBX_TEST_POSTGRES_URL").expect("DBX_TEST_POSTGRES_URL");
+        let pool = connect_with_local_timezone(&url, Duration::from_secs(10), "UTC")
+            .await
+            .expect("connect PostgreSQL database");
+
+        // `execute_query_with_max_rows` is the public command helper (used by
+        // DROP DATABASE and the transfer/export fallback). A statement with no
+        // result set must still attach the notices it raised.
+        let result =
+            execute_query_with_max_rows(&pool, "DO $$ BEGIN RAISE NOTICE 'dbx public notice regression'; END $$", None)
+                .await
+                .expect("execute statement with notice");
+
+        assert!(result.messages.iter().any(|message| message.message == "dbx public notice regression"));
+    }
+
    #[test]
    fn take_notices_for_key_returns_buffered_notices_and_empties_buffer() {
        let key = ("test-host".to_string(), "9000001".to_string(), "9000001".to_string());
@@ -71,11 +71,17 @@ test("view switcher exposes a messages button with a count badge", () => {
 test("ContentArea wires server messages into the switcher and the messages view", () => {
  const contentArea = source(contentAreaPath);

-  assert.match(contentArea, /resultMessageCount = computed\(\(\) => props\.activeTab\.result\?\.messages\?\.length \?\? 0\)/);
+  // Messages are aggregated across every statement result of the run so a
+  // message-only statement (e.g. a PostgreSQL DO block) keeps its notices even
+  // when a later statement owns the active result.
+  assert.match(contentArea, /const resultMessages = computed<QueryMessage\[\]>\(\(\) => \{/);
+  assert.match(contentArea, /props\.activeTab\.results\?\.length \? props\.activeTab\.results : props\.activeTab\.result \? \[props\.activeTab\.result\] : \[\]/);
+  assert.match(contentArea, /results\.flatMap\(\(result\) => result\.messages \?\? \[\]\)/);
+  assert.match(contentArea, /resultMessageCount = computed\(\(\) => resultMessages\.value\.length\)/);
  assert.match(contentArea, /canShowMessagesOutput = computed\(\(\) => resultMessageCount\.value > 0\)/);
  assert.equal((contentArea.match(/:can-show-messages="canShowMessagesOutput"/g) ?? []).length, 2);
  assert.equal((contentArea.match(/:message-count="resultMessageCount"/g) ?? []).length, 2);
-  assert.match(contentArea, /<QueryMessagesView v-else-if="activeOutputView === 'messages'"[\s\S]*:messages="activeTab\.result\?\.messages \?\? \[\]"/);
+  assert.match(contentArea, /<QueryMessagesView v-else-if="activeOutputView === 'messages'"[\s\S]*:messages="resultMessages"/);
  // Results with no result set auto-switch via the shared default-view helper.
  assert.match(contentArea, /import \{ defaultViewForResult \} from "@\/lib\/query\/queryResultDefaultView"/);
  assert.match(contentArea, /emit\("update:activeOutputView", props\.activeTab\.id, result \? defaultViewForResult\(result\) : "summary"\)/);