diff --git a/frontend/src/core/cells/__tests__/cells.test.ts b/frontend/src/core/cells/__tests__/cells.test.ts index d9d7bab4dd6..a2e47a5d033 100644 --- a/frontend/src/core/cells/__tests__/cells.test.ts +++ b/frontend/src/core/cells/__tests__/cells.test.ts @@ -2288,6 +2288,63 @@ describe("cell reducer", () => { expect(focusSpy).not.toHaveBeenCalled(); expect(state.cellIds.inOrderIds).toEqual(initialState.cellIds.inOrderIds); }); + + it("creates a new SQL cell when moving after a SQL cell", () => { + // Create a SQL cell + actions.createNewCell({ + cellId: "__end__", + before: false, + code: `_df = mo.sql(f"""SELECT * FROM table""")`, + }); + + const sqlCellId = + state.cellIds.inOrderIds[state.cellIds.inOrderIds.length - 1]; + const initialCellCount = state.cellIds.inOrderIds.length; + + // Move to next cell (should create a new SQL cell) + actions.moveToNextCell({ + cellId: sqlCellId, + before: false, + noCreate: false, + }); + + expect(state.cellIds.inOrderIds.length).toBe(initialCellCount + 1); + const newCellId = + state.cellIds.inOrderIds[state.cellIds.inOrderIds.length - 1]; + const newCell = state.cellData[newCellId]; + + // The new cell should have SQL code (mo.sql) + expect(newCell.code).toContain("mo.sql"); + expect(newCell.code).toContain("_df"); + }); + + it("creates a new Python cell when moving after a Python cell", () => { + // Create a Python cell + actions.createNewCell({ + cellId: "__end__", + before: false, + code: "x = 1", + }); + + const pythonCellId = + state.cellIds.inOrderIds[state.cellIds.inOrderIds.length - 1]; + const initialCellCount = state.cellIds.inOrderIds.length; + + // Move to next cell (should create a new Python cell) + actions.moveToNextCell({ + cellId: pythonCellId, + before: false, + noCreate: false, + }); + + expect(state.cellIds.inOrderIds.length).toBe(initialCellCount + 1); + const newCellId = + state.cellIds.inOrderIds[state.cellIds.inOrderIds.length - 1]; + const newCell = state.cellData[newCellId]; + + // The new cell should be empty (default Python cell) + expect(newCell.code).toBe(""); + }); }); describe("untouched cells functionality", () => { diff --git a/frontend/src/core/cells/cells.ts b/frontend/src/core/cells/cells.ts index 5fd8e220ba1..3562931b65b 100644 --- a/frontend/src/core/cells/cells.ts +++ b/frontend/src/core/cells/cells.ts @@ -18,6 +18,7 @@ import { Objects } from "@/utils/objects"; import { extractAllTracebackInfo, type TracebackInfo } from "@/utils/traceback"; import { createReducerAndAtoms } from "../../utils/createReducer"; import { foldAllBulk, unfoldAllBulk } from "../codemirror/editing/commands"; +import { LanguageAdapters } from "../codemirror/language/LanguageAdapters"; import { splitEditor, updateEditorCodeFromPython, @@ -955,15 +956,29 @@ const { const isPastLastCell = nextCellIndex === column.length; const isBeforeFirstCell = nextCellIndex === -1; + const getNewCellCode = (): string => { + const currentCell = state.cellData[cellId]; + if (!currentCell) { + return ""; + } + + const isSqlCell = LanguageAdapters.sql.isSupported(currentCell.code); + if (isSqlCell) { + return LanguageAdapters.sql.defaultCode; + } + + return ""; + }; // Create a new cell at the end and set scrollKey to focus it if (isPastLastCell && !noCreate) { const newCellId = CellId.create(); + const newCellCode = getNewCellCode(); return { ...state, cellIds: state.cellIds.insertId(newCellId, column.id, nextCellIndex), cellData: { ...state.cellData, - [newCellId]: createCell({ id: newCellId }), + [newCellId]: createCell({ id: newCellId, code: newCellCode }), }, cellRuntime: { ...state.cellRuntime, @@ -979,12 +994,13 @@ const { if (isBeforeFirstCell && !noCreate) { const newCellId = CellId.create(); + const newCellCode = getNewCellCode(); return { ...state, cellIds: state.cellIds.insertId(newCellId, column.id, 0), cellData: { ...state.cellData, - [newCellId]: createCell({ id: newCellId }), + [newCellId]: createCell({ id: newCellId, code: newCellCode }), }, cellRuntime: { ...state.cellRuntime,