From 33dcefb76b1354c2c48609f72384c867fc195566 Mon Sep 17 00:00:00 2001 From: Shahmir Varqha Date: Wed, 5 Feb 2025 02:32:46 +0800 Subject: [PATCH] fix: checks whether latestEngineSelected is in list of engines (#3685) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit ## 📝 Summary 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 ## 📋 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 @akshayka OR @mscolnick --------- Co-authored-by: Myles Scolnick --- .../codemirror/language/__tests__/sql.test.ts | 67 ++++++++++++++++--- .../src/core/codemirror/language/panel.tsx | 4 +- frontend/src/core/codemirror/language/sql.ts | 20 +++--- .../__tests__/data-source.test.ts | 2 +- .../data-source-connections.ts | 49 ++++++++++++-- .../src/core/websocket/useMarimoWebSocket.tsx | 2 +- 6 files changed, 111 insertions(+), 33 deletions(-) rename frontend/src/core/{cells => datasets}/__tests__/data-source.test.ts (98%) rename frontend/src/core/{cells => datasets}/data-source-connections.ts (62%) diff --git a/frontend/src/core/codemirror/language/__tests__/sql.test.ts b/frontend/src/core/codemirror/language/__tests__/sql.test.ts index e65266adbee..82562282f36 100644 --- a/frontend/src/core/codemirror/language/__tests__/sql.test.ts +++ b/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, + ); }); }); diff --git a/frontend/src/core/codemirror/language/panel.tsx b/frontend/src/core/codemirror/language/panel.tsx index 5990e2e543f..036bfa6c4a5 100644 --- a/frontend/src/core/codemirror/language/panel.tsx +++ b/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 = ({ ); const engineIsDisconnected = - selectedEngine && connectionsMap.get(selectedEngine.name) === undefined; + selectedEngine && !connectionsMap.has(selectedEngine.name); const handleSelectEngine = (value: string) => { const nextEngine = connectionsMap.get(value as ConnectionName); diff --git a/frontend/src/core/codemirror/language/sql.ts b/frontend/src/core/codemirror/language/sql.ts index 270fef733a8..469fc20d79a 100644 --- a/frontend/src/core/codemirror/language/sql.ts +++ b/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(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 { diff --git a/frontend/src/core/cells/__tests__/data-source.test.ts b/frontend/src/core/datasets/__tests__/data-source.test.ts similarity index 98% rename from frontend/src/core/cells/__tests__/data-source.test.ts rename to frontend/src/core/datasets/__tests__/data-source.test.ts index 492c1d2b3ad..126886ab490 100644 --- a/frontend/src/core/cells/__tests__/data-source.test.ts +++ b/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; diff --git a/frontend/src/core/cells/data-source-connections.ts b/frontend/src/core/datasets/data-source-connections.ts similarity index 62% rename from frontend/src/core/cells/data-source-connections.ts rename to frontend/src/core/datasets/data-source-connections.ts index 6f3a5454702..3cb79751d3d 100644 --- a/frontend/src/core/cells/data-source-connections.ts +++ b/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; } function initialState(): DataSourceState { return { + latestEngineSelected: DEFAULT_ENGINE, connectionsMap: new Map().set(DEFAULT_ENGINE, { name: DEFAULT_ENGINE, source: "duckdb", @@ -43,15 +47,20 @@ 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 @@ -59,19 +68,27 @@ const { 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(), }), @@ -79,9 +96,16 @@ const { 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, diff --git a/frontend/src/core/websocket/useMarimoWebSocket.tsx b/frontend/src/core/websocket/useMarimoWebSocket.tsx index 09f502f4f9e..1ab3315b9e5 100644 --- a/frontend/src/core/websocket/useMarimoWebSocket.tsx +++ b/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.