feat: Enhance Deepnote notebook functionality with SQL metadata support

- Added a new schema for Deepnote SQL metadata to facilitate better handling of SQL-related input blocks.
- Updated the Deepnote notebook command listener to utilize the new SQL metadata, improving variable name management and default metadata setup for new cells.
- Enhanced the cell insertion logic to ensure proper selection and editing of newly created cells.

Signed-off-by: Tomas Kislan <tomas@kislan.sk>

Author: tkislan

src/notebooks/deepnote/deepnoteNotebookCommandListener.ts

@@ -6,7 +6,8 @@ import {
     NotebookCellKind,
     NotebookEdit,
     NotebookRange,
-    NotebookCell
+    NotebookCell,
+    NotebookEditorRevealType
 } from 'vscode';
 import z from 'zod';
 
@@ -25,7 +26,8 @@ import {
     DeepnoteDateInputMetadataSchema,
     DeepnoteDateRangeInputMetadataSchema,
     DeepnoteFileInputMetadataSchema,
-    DeepnoteButtonMetadataSchema
+    DeepnoteButtonMetadataSchema,
+    DeepnoteSqlMetadata
 } from './deepnoteSchemas';
 
 type InputBlockType =
@@ -88,12 +90,17 @@ export function getNextDeepnoteVariableName(cells: NotebookCell[], prefix: 'df'
             acc.push(contentValue);
         }
 
-        const parsedMetadataValue = z.string().safeParse(cell.metadata.__deepnotePocket?.variableName);
-
+        const parsedMetadataValue = z.string().safeParse(cell.metadata['deepnote_variable_name']);
         if (parsedMetadataValue.success) {
             acc.push(parsedMetadataValue.data);
         }
 
+        const parsedPocketMetadataValue = z.string().safeParse(cell.metadata.__deepnotePocket?.deepnote_variable_name);
+
+        if (parsedPocketMetadataValue.success) {
+            acc.push(parsedPocketMetadataValue.data);
+        }
+
         return acc;
     }, []);
 
@@ -175,6 +182,14 @@ export class DeepnoteNotebookCommandListener implements IExtensionSyncActivation
         }
         const document = editor.notebook;
         const selection = editor.selection;
+        const cells = editor.notebook.getCells();
+        const deepnoteVariableName = getNextDeepnoteVariableName(cells, 'df');
+
+        const defaultMetadata: DeepnoteSqlMetadata = {
+            deepnote_variable_name: deepnoteVariableName,
+            deepnote_return_variable_type: 'dataframe',
+            sql_integration_id: ''
+        };
 
         // Determine the index where to insert the new cell (below current selection or at the end)
         const insertIndex = selection ? selection.end : document.cellCount;
@@ -185,14 +200,22 @@ export class DeepnoteNotebookCommandListener implements IExtensionSyncActivation
             const newCell = new NotebookCellData(NotebookCellKind.Code, '', 'sql');
             newCell.metadata = {
                 __deepnotePocket: {
-                    type: 'sql'
-                }
+                    type: 'sql',
+                    ...defaultMetadata
+                },
+                ...defaultMetadata
             };
             const nbEdit = NotebookEdit.insertCells(insertIndex, [newCell]);
             edit.set(document.uri, [nbEdit]);
         }).then(
             () => {
-                editor.selection = new NotebookRange(insertIndex, insertIndex + 1);
+                const notebookRange = new NotebookRange(insertIndex, insertIndex + 1);
+                editor.revealRange(notebookRange, NotebookEditorRevealType.Default);
+                editor.selection = notebookRange;
+                // Enter edit mode on the new cell
+                commands
+                    .executeCommand('notebook.cell.edit')
+                    .then(undefined, (error) => logger.error('Error entering edit mode', error));
             },
             (error) => {
                 logger.error('Error inserting SQL block', error);
@@ -231,7 +254,13 @@ export class DeepnoteNotebookCommandListener implements IExtensionSyncActivation
             edit.set(document.uri, [nbEdit]);
         }).then(
             () => {
-                editor.selection = new NotebookRange(insertIndex, insertIndex + 1);
+                const notebookRange = new NotebookRange(insertIndex, insertIndex + 1);
+                editor.revealRange(notebookRange, NotebookEditorRevealType.Default);
+                editor.selection = notebookRange;
+                // Enter edit mode on the new cell
+                commands
+                    .executeCommand('notebook.cell.edit')
+                    .then(undefined, (error) => logger.error('Error entering edit mode', error));
             },
             (error) => {
                 logger.error('Error inserting big number chart block', error);
@@ -273,7 +302,13 @@ export class DeepnoteNotebookCommandListener implements IExtensionSyncActivation
             edit.set(document.uri, [nbEdit]);
         }).then(
             () => {
-                editor.selection = new NotebookRange(insertIndex, insertIndex + 1);
+                const notebookRange = new NotebookRange(insertIndex, insertIndex + 1);
+                editor.revealRange(notebookRange, NotebookEditorRevealType.Default);
+                editor.selection = notebookRange;
+                // Enter edit mode on the new cell
+                commands
+                    .executeCommand('notebook.cell.edit')
+                    .then(undefined, (error) => logger.error('Error entering edit mode', error));
             },
             (error) => {
                 logger.error('Error inserting input block', error);

src/notebooks/deepnote/deepnoteSchemas.ts

@@ -43,6 +43,21 @@ export const DeepnoteBigNumberMetadataSchema = z.object({
         .transform((val) => val ?? false)
 });
 
+export const DeepnoteSqlMetadataSchema = z.object({
+    deepnote_variable_name: z
+        .string()
+        .nullish()
+        .transform((val) => val ?? ''),
+    deepnote_return_variable_type: z
+        .enum(['dataframe', 'query_preview'])
+        .nullish()
+        .transform((val) => val ?? 'dataframe'),
+    sql_integration_id: z
+        .string()
+        .nullish()
+        .transform((val) => val ?? '')
+});
+
 // Base schema with common fields for all input types
 const DeepnoteBaseInputMetadataSchema = z.object({
     deepnote_variable_name: z
@@ -213,3 +228,4 @@ export const DeepnoteButtonMetadataSchema = DeepnoteBaseInputMetadataSchema.exte
 
 export type DeepnoteChartBigNumberOutput = z.infer<typeof DeepnoteChartBigNumberOutputSchema>;
 export type DeepnoteBigNumberMetadata = z.infer<typeof DeepnoteBigNumberMetadataSchema>;
+export type DeepnoteSqlMetadata = z.infer<typeof DeepnoteSqlMetadataSchema>;

src/webviews/webview-side/chart-big-number-renderer/ChartBigNumberOutputRenderer.tsx

@@ -24,7 +24,7 @@ export function ChartBigNumberOutputRenderer({
             return 'NaN';
         }
 
-        return formatValue(parsedValue, metadata.deepnote_big_number_format ?? 'number');
+        return formatValue(parsedValue, metadata.deepnote_big_number_format);
     }, [output.value, metadata.deepnote_big_number_format]);
 
     const comparisonValue = useMemo(() => {