From d94fc7b7feb79efc7d14f802f6b867754295546a Mon Sep 17 00:00:00 2001
From: jankuca <jan@jankuca.com>
Date: Mon, 27 Oct 2025 11:05:41 +0100
Subject: [PATCH] fix: Dataframe SQL block support

---
 ...IntegrationEnvironmentVariablesProvider.ts | 23 ++++++++++++++-----
 ...nEnvironmentVariablesProvider.unit.test.ts | 10 ++++++--
 2 files changed, 25 insertions(+), 8 deletions(-)

diff --git a/src/platform/notebooks/deepnote/sqlIntegrationEnvironmentVariablesProvider.ts b/src/platform/notebooks/deepnote/sqlIntegrationEnvironmentVariablesProvider.ts
index 13da2ab4a3..9a5aa93b93 100644
--- a/src/platform/notebooks/deepnote/sqlIntegrationEnvironmentVariablesProvider.ts
+++ b/src/platform/notebooks/deepnote/sqlIntegrationEnvironmentVariablesProvider.ts
@@ -149,6 +149,22 @@ export class SqlIntegrationEnvironmentVariablesProvider implements ISqlIntegrati
             }
 
             try {
+                // Handle internal DuckDB integration specially
+                if (integrationId === DATAFRAME_SQL_INTEGRATION_ID) {
+                    const envVarName = convertToEnvironmentVariableName(getSqlEnvVarName(integrationId));
+                    const credentialsJson = JSON.stringify({
+                        url: 'deepnote+duckdb:///:memory:',
+                        params: {},
+                        param_style: 'qmark'
+                    });
+
+                    envVars[envVarName] = credentialsJson;
+                    logger.debug(
+                        `SqlIntegrationEnvironmentVariablesProvider: Added env var for dataframe SQL integration`
+                    );
+                    continue;
+                }
+
                 const config = await this.integrationStorage.getIntegrationConfig(integrationId);
                 if (!config) {
                     logger.warn(
@@ -162,7 +178,7 @@ export class SqlIntegrationEnvironmentVariablesProvider implements ISqlIntegrati
                 const credentialsJson = convertIntegrationConfigToJson(config);
 
                 envVars[envVarName] = credentialsJson;
-                logger.info(
+                logger.debug(
                     `SqlIntegrationEnvironmentVariablesProvider: Added env var ${envVarName} for integration ${integrationId}`
                 );
             } catch (error) {
@@ -194,11 +210,6 @@ export class SqlIntegrationEnvironmentVariablesProvider implements ISqlIntegrati
             if (metadata && typeof metadata === 'object') {
                 const integrationId = (metadata as Record<string, unknown>).sql_integration_id;
                 if (typeof integrationId === 'string') {
-                    // Skip the internal DuckDB integration
-                    if (integrationId === DATAFRAME_SQL_INTEGRATION_ID) {
-                        continue;
-                    }
-
                     integrationIds.add(integrationId);
                     logger.trace(
                         `SqlIntegrationEnvironmentVariablesProvider: Found integration ${integrationId} in cell ${cell.index}`
diff --git a/src/platform/notebooks/deepnote/sqlIntegrationEnvironmentVariablesProvider.unit.test.ts b/src/platform/notebooks/deepnote/sqlIntegrationEnvironmentVariablesProvider.unit.test.ts
index f7348d124a..82cf82c20e 100644
--- a/src/platform/notebooks/deepnote/sqlIntegrationEnvironmentVariablesProvider.unit.test.ts
+++ b/src/platform/notebooks/deepnote/sqlIntegrationEnvironmentVariablesProvider.unit.test.ts
@@ -64,7 +64,7 @@ suite('SqlIntegrationEnvironmentVariablesProvider', () => {
         assert.deepStrictEqual(envVars, {});
     });
 
-    test('Skips internal DuckDB integration', async () => {
+    test('Returns environment variable for internal DuckDB integration', async () => {
         const uri = Uri.file('/test/notebook.deepnote');
         const notebook = createMockNotebook(uri, [
             createMockCell(0, NotebookCellKind.Code, 'sql', 'SELECT * FROM df', {
@@ -75,7 +75,13 @@ suite('SqlIntegrationEnvironmentVariablesProvider', () => {
         when(mockedVSCodeNamespaces.workspace.notebookDocuments).thenReturn([notebook]);
 
         const envVars = await provider.getEnvironmentVariables(uri);
-        assert.deepStrictEqual(envVars, {});
+
+        // Check that the environment variable is set for dataframe SQL
+        assert.property(envVars, 'SQL_DEEPNOTE_DATAFRAME_SQL');
+        const credentialsJson = JSON.parse(envVars['SQL_DEEPNOTE_DATAFRAME_SQL']!);
+        assert.strictEqual(credentialsJson.url, 'deepnote+duckdb:///:memory:');
+        assert.deepStrictEqual(credentialsJson.params, {});
+        assert.strictEqual(credentialsJson.param_style, 'qmark');
     });
 
     test('Returns environment variable for PostgreSQL integration', async () => {