fix: checks whether latestEngineSelected is in list of engines (#3685)

## 📝 Summary

<!--
Provide a concise summary of what this pull request is addressing.

If this PR fixes any issues, list them here by number (e.g., Fixes
#123).
-->
When you delete an engine and create a new SQL cell, the cell's code
still uses the deleted engine. This fixes that.

However, it's not an ideal fix, using the constant variable triggers a 
`Uncaught ReferenceError: can't access lexical declaration
'DEFAULT_ENGINE' before initialization`

makes me think I have designed the latestEngineSelected feature wrongly.

another smaller issue:
- delete an engine, refresh the page, undelete the engine, the sql
cells's dropdown do not update

## 🔍 Description of Changes

<!--
Detail the specific changes made in this pull request. Explain the
problem addressed and how it was resolved. If applicable, provide before
and after comparisons, screenshots, or any relevant details to help
reviewers understand the changes easily.
-->

## 📋 Checklist

- [x] I have read the [contributor
guidelines](https://github.com/marimo-team/marimo/blob/main/CONTRIBUTING.md).
- [ ] For large changes, or changes that affect the public API: this
change was discussed or approved through an issue, on
[Discord](https://marimo.io/discord?ref=pr), or the community
[discussions](https://github.com/marimo-team/marimo/discussions) (Please
provide a link if applicable).
- [ ] I have added tests for the changes made.
- [ ] I have run the code and verified that it works as expected.

## 📜 Reviewers

<!--
Tag potential reviewers from the community or maintainers who might be
interested in reviewing this pull request.

Your PR will be reviewed more quickly if you can figure out the right
person to tag with @ -->

@akshayka OR @mscolnick

---------

Co-authored-by: Myles Scolnick <[email protected]>

frontend/src/core/codemirror/language/__tests__/sql.test.ts

@@ -1,13 +1,13 @@
 /* Copyright 2024 Marimo. All rights reserved. */
 import { expect, describe, it, beforeAll, afterAll, afterEach } from "vitest";
-import {
-  DEFAULT_ENGINE,
-  latestEngineSelected,
-  SQLLanguageAdapter,
-} from "../sql";
+import { SQLLanguageAdapter } from "../sql";
 import { store } from "@/core/state/jotai";
 import { capabilitiesAtom } from "@/core/config/capabilities";
-import type { ConnectionName } from "@/core/cells/data-source-connections";
+import {
+  dataSourceConnectionsAtom,
+  DEFAULT_ENGINE,
+  type ConnectionName,
+} from "@/core/datasets/data-source-connections";
 
 const adapter = new SQLLanguageAdapter();
 
@@ -410,6 +410,25 @@ _df = mo.sql(
   describe("latestEngineSelected", () => {
     afterEach(() => {
       adapter.engine = DEFAULT_ENGINE;
+      const state = store.get(dataSourceConnectionsAtom);
+      const connections = new Map(state.connectionsMap);
+      connections
+        .set("postgres_engine" as ConnectionName, {
+          name: "postgres_engine" as ConnectionName,
+          source: "postgres",
+          display_name: "PostgreSQL",
+          dialect: "postgres",
+        })
+        .set("mysql_engine" as ConnectionName, {
+          name: "mysql_engine" as ConnectionName,
+          source: "mysql",
+          display_name: "MySQL",
+          dialect: "mysql",
+        });
+      store.set(dataSourceConnectionsAtom, {
+        ...state,
+        connectionsMap: connections,
+      });
     });
 
     it("should use default engine initially", () => {
@@ -420,7 +439,9 @@ _df = mo.sql(
       const engine = "postgres_engine" as ConnectionName;
       adapter.selectEngine(engine);
       expect(adapter.engine).toBe(engine);
-      expect(store.get(latestEngineSelected)).toBe(engine);
+      expect(store.get(dataSourceConnectionsAtom).latestEngineSelected).toBe(
+        engine,
+      );
     });
 
     it("should allow switching between engines", () => {
@@ -429,18 +450,40 @@ _df = mo.sql(
 
       adapter.selectEngine(engine1);
       expect(adapter.engine).toBe(engine1);
-      expect(store.get(latestEngineSelected)).toBe(engine1);
+      expect(store.get(dataSourceConnectionsAtom).latestEngineSelected).toBe(
+        engine1,
+      );
 
       adapter.selectEngine(engine2);
       expect(adapter.engine).toBe(engine2);
-      expect(store.get(latestEngineSelected)).toBe(engine2);
+      expect(store.get(dataSourceConnectionsAtom).latestEngineSelected).toBe(
+        engine2,
+      );
     });
 
     it("should update engine in transformIn when specified", () => {
       const pythonCode = '_df = mo.sql("""SELECT 1""", engine=postgres_engine)';
       adapter.transformIn(pythonCode);
       expect(adapter.engine).toBe("postgres_engine");
-      expect(store.get(latestEngineSelected)).toBe("postgres_engine");
+      expect(store.get(dataSourceConnectionsAtom).latestEngineSelected).toBe(
+        "postgres_engine",
+      );
+
+      // Don't update for unspecified engine
+      const pythonCode2 = '_df = mo.sql("""SELECT 1""")';
+      adapter.transformIn(pythonCode2);
+      expect(adapter.engine).toBe(DEFAULT_ENGINE);
+      expect(store.get(dataSourceConnectionsAtom).latestEngineSelected).toBe(
+        "postgres_engine",
+      );
+
+      // Don't update for unknown engine
+      const pythonCode3 = '_df = mo.sql("""SELECT 1""", engine=unknown_engine)';
+      adapter.transformIn(pythonCode3);
+      expect(adapter.engine).toBe("unknown_engine");
+      expect(store.get(dataSourceConnectionsAtom).latestEngineSelected).toBe(
+        "postgres_engine",
+      );
     });
 
     it("should maintain engine selection across transformIn/transformOut", () => {
@@ -473,7 +516,9 @@ _df = mo.sql(
       adapter.selectEngine(DEFAULT_ENGINE);
 
       expect(adapter.engine).toBe(DEFAULT_ENGINE);
-      expect(store.get(latestEngineSelected)).toBe(DEFAULT_ENGINE);
+      expect(store.get(dataSourceConnectionsAtom).latestEngineSelected).toBe(
+        DEFAULT_ENGINE,
+      );
     });
   });
 

frontend/src/core/codemirror/language/panel.tsx

@@ -8,7 +8,7 @@ import { getFeatureFlag } from "@/core/config/feature-flag";
 import {
   type ConnectionName,
   dataConnectionsMapAtom,
-} from "@/core/cells/data-source-connections";
+} from "@/core/datasets/data-source-connections";
 import { useAtomValue } from "jotai";
 import { AlertCircle, CircleHelpIcon } from "lucide-react";
 import { Tooltip, TooltipProvider } from "@/components/ui/tooltip";
@@ -118,7 +118,7 @@ const SQLEngineSelect: React.FC<SelectProps> = ({
   );
 
   const engineIsDisconnected =
-    selectedEngine && connectionsMap.get(selectedEngine.name) === undefined;
+    selectedEngine && !connectionsMap.has(selectedEngine.name);
 
   const handleSelectEngine = (value: string) => {
     const nextEngine = connectionsMap.get(value as ConnectionName);

frontend/src/core/codemirror/language/sql.ts

@@ -15,19 +15,17 @@ import { datasetsAtom } from "@/core/datasets/state";
 import { type QuotePrefixKind, upgradePrefixKind } from "./utils/quotes";
 import { capabilitiesAtom } from "@/core/config/capabilities";
 import { MarkdownLanguageAdapter } from "./markdown";
-import type { ConnectionName } from "@/core/cells/data-source-connections";
-import { atom } from "jotai";
+import {
+  dataSourceConnectionsAtom,
+  DEFAULT_ENGINE,
+  setLatestEngineSelected,
+  type ConnectionName,
+} from "@/core/datasets/data-source-connections";
 import { parser } from "@lezer/python";
 import type { SyntaxNode, TreeCursor } from "@lezer/common";
 import { parseArgsKwargs } from "./utils/ast";
 import { Logger } from "@/utils/Logger";
 
-// Default engine to use when not specified, shouldn't conflict with user_defined engines
-export const DEFAULT_ENGINE: ConnectionName =
-  "_marimo_duckdb" as ConnectionName;
-
-export const latestEngineSelected = atom<ConnectionName>(DEFAULT_ENGINE);
-
 /**
  * Language adapter for SQL.
  */
@@ -39,7 +37,7 @@ export class SQLLanguageAdapter implements LanguageAdapter {
   dataframeName = "_df";
   lastQuotePrefix: QuotePrefixKind = "f";
   showOutput = true;
-  engine: ConnectionName = store.get(latestEngineSelected);
+  engine = store.get(dataSourceConnectionsAtom).latestEngineSelected;
 
   getDefaultCode(): string {
     if (this.engine === DEFAULT_ENGINE) {
@@ -77,7 +75,7 @@ export class SQLLanguageAdapter implements LanguageAdapter {
       if (this.engine !== DEFAULT_ENGINE) {
         // User selected a new engine, set it as latest.
         // This makes new SQL statements use the new engine by default.
-        store.set(latestEngineSelected, this.engine);
+        setLatestEngineSelected(this.engine);
       }
 
       return [
@@ -130,7 +128,7 @@ export class SQLLanguageAdapter implements LanguageAdapter {
 
   selectEngine(connectionName: ConnectionName): void {
     this.engine = connectionName;
-    store.set(latestEngineSelected, connectionName);
+    setLatestEngineSelected(this.engine);
   }
 
   setShowOutput(showOutput: boolean): void {

frontend/src/core/cells/__tests__/data-source.test.ts → frontend/src/core/datasets/__tests__/data-source.test.ts

@@ -4,9 +4,9 @@ import {
   type ConnectionName,
   type DataSourceConnection,
   type DataSourceState,
+  DEFAULT_ENGINE,
   exportedForTesting,
 } from "../data-source-connections";
-import { DEFAULT_ENGINE } from "@/core/codemirror/language/sql";
 import type { VariableName } from "@/core/variables/types";
 
 const { reducer, initialState } = exportedForTesting;

frontend/src/core/cells/data-source-connections.ts → frontend/src/core/datasets/data-source-connections.ts

@@ -1,12 +1,14 @@
 /* Copyright 2024 Marimo. All rights reserved. */
 import { createReducerAndAtoms } from "@/utils/createReducer";
 import type { TypedString } from "@/utils/typed";
-import { DEFAULT_ENGINE } from "../codemirror/language/sql";
 import type { VariableName } from "../variables/types";
 import { atom } from "jotai";
+import { store } from "../state/jotai";
 
 export type ConnectionName = TypedString<"ConnectionName">;
 
+export const DEFAULT_ENGINE = "_marimo_duckdb" as ConnectionName;
+
 export interface DataSourceConnection {
   name: ConnectionName;
   source: string;
@@ -15,11 +17,13 @@ export interface DataSourceConnection {
 }
 
 export interface DataSourceState {
+  latestEngineSelected: ConnectionName;
   connectionsMap: ReadonlyMap<ConnectionName, DataSourceConnection>;
 }
 
 function initialState(): DataSourceState {
   return {
+    latestEngineSelected: DEFAULT_ENGINE,
     connectionsMap: new Map().set(DEFAULT_ENGINE, {
       name: DEFAULT_ENGINE,
       source: "duckdb",
@@ -43,45 +47,65 @@ const {
       return state;
     }
 
+    const { latestEngineSelected, connectionsMap } = state;
+
     // update existing connections with latest values
     // add new ones if they don't exist
     // Backend will dedupe by connection name & keep the latest, so we use this as the key
-    const newMap = new Map(state.connectionsMap);
+    const newMap = new Map(connectionsMap);
     for (const conn of opts.connections) {
       newMap.set(conn.name, conn);
     }
 
-    return { connectionsMap: newMap };
+    return {
+      latestEngineSelected,
+      connectionsMap: newMap,
+    };
   },
 
   // Keep default engine and any connections that are used by variables
   filterDataSourcesFromVariables: (
     state: DataSourceState,
     variableNames: VariableName[],
   ) => {
+    const { latestEngineSelected, connectionsMap } = state;
     const names = new Set(variableNames);
     const newMap = new Map(
-      [...state.connectionsMap].filter(([name]) => {
+      [...connectionsMap].filter(([name]) => {
         if (name === DEFAULT_ENGINE) {
           return true;
         }
         return names.has(name as unknown as VariableName);
       }),
     );
-    return { connectionsMap: newMap };
+    return {
+      // If the latest engine selected is not in the new map, use the default engine
+      latestEngineSelected: newMap.has(latestEngineSelected)
+        ? latestEngineSelected
+        : DEFAULT_ENGINE,
+      connectionsMap: newMap,
+    };
   },
 
   clearDataSourceConnections: (): DataSourceState => ({
+    latestEngineSelected: DEFAULT_ENGINE,
     connectionsMap: new Map(),
   }),
 
   removeDataSourceConnection: (
     state: DataSourceState,
     connectionName: ConnectionName,
   ): DataSourceState => {
-    const newMap = new Map(state.connectionsMap);
+    const { latestEngineSelected, connectionsMap } = state;
+
+    const newMap = new Map(connectionsMap);
     newMap.delete(connectionName);
-    return { connectionsMap: newMap };
+    return {
+      latestEngineSelected: newMap.has(latestEngineSelected)
+        ? latestEngineSelected
+        : DEFAULT_ENGINE,
+      connectionsMap: newMap,
+    };
   },
 });
 
@@ -91,6 +115,17 @@ export const dataConnectionsMapAtom = atom(
   (get) => get(dataSourceConnectionsAtom).connectionsMap,
 );
 
+export function setLatestEngineSelected(engine: ConnectionName) {
+  const existing = store.get(dataSourceConnectionsAtom);
+  // Don't update the map if the engine is not in the map
+  if (existing.connectionsMap.has(engine)) {
+    store.set(dataSourceConnectionsAtom, {
+      ...existing,
+      latestEngineSelected: engine,
+    });
+  }
+}
+
 export const exportedForTesting = {
   reducer,
   createActions,

frontend/src/core/websocket/useMarimoWebSocket.tsx

@@ -41,7 +41,7 @@ import { capabilitiesAtom } from "../config/capabilities";
 import { UI_ELEMENT_REGISTRY } from "../dom/uiregistry";
 import { reloadSafe } from "@/utils/reload-safe";
 import { useRunsActions } from "../cells/runs";
-import { useDataSourceActions } from "../cells/data-source-connections";
+import { useDataSourceActions } from "../datasets/data-source-connections";
 
 /**
  * WebSocket that connects to the Marimo kernel and handles incoming messages.