fix: Fix dataframe output duplicate rendering (#136)

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

Author: tkislan

src/notebooks/deepnote/deepnoteNotebookCommandListener.ts

@@ -30,6 +30,7 @@ import {
     DeepnoteButtonMetadataSchema,
     DeepnoteSqlMetadata
 } from './deepnoteSchemas';
+import { DATAFRAME_SQL_INTEGRATION_ID } from '../../platform/notebooks/deepnote/integrationTypes';
 
 export type InputBlockType =
     | 'input-text'
@@ -189,7 +190,7 @@ export class DeepnoteNotebookCommandListener implements IExtensionSyncActivation
         const defaultMetadata: DeepnoteSqlMetadata = {
             deepnote_variable_name: deepnoteVariableName,
             deepnote_return_variable_type: 'dataframe',
-            sql_integration_id: 'deepnote-dataframe-sql'
+            sql_integration_id: DATAFRAME_SQL_INTEGRATION_ID
         };
 
         // Determine the index where to insert the new cell (below current selection or at the end)

src/notebooks/deepnote/deepnoteNotebookCommandListener.unit.test.ts

@@ -21,6 +21,7 @@ import {
 import { IDisposable } from '../../platform/common/types';
 import * as notebookUpdater from '../../kernels/execution/notebookUpdater';
 import { createMockedNotebookDocument } from '../../test/datascience/editor-integration/helpers';
+import { DATAFRAME_SQL_INTEGRATION_ID } from '../../platform/notebooks/deepnote/integrationTypes';
 
 suite('DeepnoteNotebookCommandListener', () => {
     let commandListener: DeepnoteNotebookCommandListener;
@@ -734,7 +735,7 @@ suite('DeepnoteNotebookCommandListener', () => {
                 );
                 assert.equal(
                     newCell.metadata.sql_integration_id,
-                    'deepnote-dataframe-sql',
+                    DATAFRAME_SQL_INTEGRATION_ID,
                     'Should have correct sql integration id'
                 );
 

src/platform/notebooks/deepnote/sqlIntegrationEnvironmentVariablesProvider.unit.test.ts

@@ -9,6 +9,7 @@ import {
     IntegrationType,
     PostgresIntegrationConfig,
     BigQueryIntegrationConfig,
+    DATAFRAME_SQL_INTEGRATION_ID,
     SnowflakeIntegrationConfig,
     SnowflakeAuthMethods
 } from './integrationTypes';
@@ -74,7 +75,7 @@ suite('SqlIntegrationEnvironmentVariablesProvider', () => {
         const uri = Uri.file('/test/notebook.deepnote');
         const notebook = createMockNotebook(uri, [
             createMockCell(0, NotebookCellKind.Code, 'sql', 'SELECT * FROM df', {
-                sql_integration_id: 'deepnote-dataframe-sql'
+                sql_integration_id: DATAFRAME_SQL_INTEGRATION_ID
             })
         ]);
 

src/webviews/webview-side/dataframe-renderer/DataframeRendererContainer.tsx

@@ -0,0 +1,34 @@
+import * as React from 'react';
+
+import { DataframeMetadata, DataframeRenderer } from './DataframeRenderer';
+import { RendererContext } from 'vscode-notebook-renderer';
+
+export interface Metadata {
+    cellId?: string;
+    cellIndex?: number;
+    executionCount: number;
+    metadata?: DataframeMetadata;
+    outputType: string;
+}
+
+export function DataframeRendererContainer({
+    context,
+    outputJson,
+    outputMetadata
+}: {
+    context: RendererContext<unknown>;
+    // eslint-disable-next-line @typescript-eslint/no-explicit-any
+    outputJson: any;
+    outputMetadata: Metadata;
+}) {
+    console.log(`Dataframe renderer - received data with ${Object.keys(outputJson).length} keys`);
+
+    console.log('[DataframeRenderer] Full metadata', outputMetadata);
+
+    const dataFrameMetadata = outputMetadata?.metadata as DataframeMetadata | undefined;
+    const cellId = outputMetadata?.cellId;
+
+    console.log(`[DataframeRenderer] Extracted cellId: ${cellId}`);
+
+    return <DataframeRenderer context={context} data={outputJson} metadata={dataFrameMetadata} cellId={cellId} />;
+}

src/webviews/webview-side/dataframe-renderer/ErrorBoundary.tsx

@@ -0,0 +1,37 @@
+import React from 'react';
+import type { FallbackProps } from 'react-error-boundary';
+
+export function ErrorFallback({ error }: FallbackProps) {
+    return (
+        <div
+            style={{
+                padding: '16px',
+                color: 'var(--vscode-errorForeground)',
+                backgroundColor: 'var(--vscode-inputValidation-errorBackground)',
+                border: '1px solid var(--vscode-inputValidation-errorBorder)',
+                borderRadius: '4px',
+                fontFamily: 'var(--vscode-font-family)',
+                fontSize: '13px'
+            }}
+        >
+            <div style={{ fontWeight: 'bold', marginBottom: '8px' }}>Error rendering dataframe</div>
+            <div style={{ marginBottom: '8px' }}>{error.message}</div>
+            <details style={{ marginTop: '8px', cursor: 'pointer' }}>
+                <summary>Stack trace</summary>
+                <pre
+                    style={{
+                        marginTop: '8px',
+                        padding: '8px',
+                        backgroundColor: 'var(--vscode-editor-background)',
+                        overflow: 'auto',
+                        fontSize: '11px',
+                        whiteSpace: 'pre-wrap',
+                        wordBreak: 'break-word'
+                    }}
+                >
+                    {error.stack}
+                </pre>
+            </details>
+        </div>
+    );
+}

src/webviews/webview-side/dataframe-renderer/index.ts

@@ -5,61 +5,50 @@ import * as ReactDOM from 'react-dom';
 
 import type { ActivationFunction, OutputItem, RendererContext } from 'vscode-notebook-renderer';
 
-import { DataframeMetadata, DataframeRenderer } from './DataframeRenderer';
-
-interface Metadata {
-    cellId?: string;
-    cellIndex?: number;
-    executionCount: number;
-    metadata?: DataframeMetadata;
-    outputType: string;
-}
+import { ErrorBoundary } from 'react-error-boundary';
+import { ErrorFallback } from './ErrorBoundary';
+import { DataframeRendererContainer, Metadata } from './DataframeRendererContainer';
 
 export const activate: ActivationFunction = (context: RendererContext<unknown>) => {
     return {
         renderOutputItem(outputItem: OutputItem, element: HTMLElement) {
-            console.log(`Dataframe renderer - rendering output item: ${outputItem.id}`);
-            try {
-                const data = outputItem.json();
-
-                console.log(`Dataframe renderer - received data with ${Object.keys(data).length} keys`);
-
-                const metadata = outputItem.metadata as Metadata | undefined;
-
-                console.log('[DataframeRenderer] Full metadata', metadata);
-
-                const dataFrameMetadata = metadata?.metadata as DataframeMetadata | undefined;
-                const cellId = metadata?.cellId;
-                const cellIndex = metadata?.cellIndex;
-
-                console.log(`[DataframeRenderer] Extracted cellId: ${cellId}, cellIndex: ${cellIndex}`);
-
-                const root = document.createElement('div');
-
-                element.appendChild(root);
-
-                ReactDOM.render(
-                    React.createElement(DataframeRenderer, {
+            ReactDOM.render(
+                React.createElement(
+                    ErrorBoundary,
+                    {
+                        FallbackComponent: ErrorFallback,
+                        onError: (error, info) => {
+                            console.error('Vega renderer error:', error, info);
+                        }
+                    },
+                    React.createElement(DataframeRendererContainer, {
                         context,
-                        data,
-                        metadata: dataFrameMetadata,
-                        cellId,
-                        cellIndex
-                    }),
-                    root
-                );
-            } catch (error) {
-                console.error(`Error rendering dataframe: ${error}`);
-                const errorDiv = document.createElement('div');
-                errorDiv.style.padding = '10px';
-                errorDiv.style.color = 'var(--vscode-errorForeground)';
-                errorDiv.textContent = `Error rendering dataframe: ${error}`;
-                element.appendChild(errorDiv);
-            }
+                        outputJson: outputItem.json(),
+                        outputMetadata: outputItem.metadata as Metadata
+                    })
+                ),
+                element
+            );
         },
 
-        disposeOutputItem(_id?: string) {
-            // Cleanup if needed
+        disposeOutputItem(id?: string) {
+            // If undefined, all cells are being removed.
+            if (id == null) {
+                for (let i = 0; i < document.children.length; i++) {
+                    const child = document.children.item(i);
+                    if (child == null) {
+                        continue;
+                    }
+                    ReactDOM.unmountComponentAtNode(child);
+                }
+                return;
+            }
+
+            const element = document.getElementById(id);
+            if (element == null) {
+                return;
+            }
+            ReactDOM.unmountComponentAtNode(element);
         }
     };
 };