diff --git a/superset-frontend/src/core/sqlLab.ts b/superset-frontend/src/core/sqlLab.ts
index 8f374dea1be..719d984fc49 100644
--- a/superset-frontend/src/core/sqlLab.ts
+++ b/superset-frontend/src/core/sqlLab.ts
@@ -102,8 +102,11 @@ export const onDidChangeEditorDatabase: typeof sqlLabType.onDidChangeEditorDatab
     createActionListener(
       predicate(QUERY_EDITOR_SETDB),
       listener,
-      (action: { type: string; queryEditor: { dbId: number } }) =>
-        action.queryEditor.dbId,
+      (action: {
+        type: string;
+        dbId?: number;
+        queryEditor: { dbId: number };
+      }) => action.dbId || action.queryEditor.dbId,
       thisArgs,
     );