diff --git a/src/platform/notebooks/deepnote/sqlIntegrationEnvironmentVariablesProvider.ts b/src/platform/notebooks/deepnote/sqlIntegrationEnvironmentVariablesProvider.ts index f42598127..cfd9df1bd 100644 --- a/src/platform/notebooks/deepnote/sqlIntegrationEnvironmentVariablesProvider.ts +++ b/src/platform/notebooks/deepnote/sqlIntegrationEnvironmentVariablesProvider.ts @@ -199,6 +199,19 @@ export class SqlIntegrationEnvironmentVariablesProvider implements ISqlIntegrati return envVars; } + // Always add the internal DuckDB integration + const dataframeSqlIntegrationEnvVarName = convertToEnvironmentVariableName( + getSqlEnvVarName(DATAFRAME_SQL_INTEGRATION_ID) + ); + const dataframeSqlIntegrationCredentialsJson = JSON.stringify({ + url: 'deepnote+duckdb:///:memory:', + params: {}, + param_style: 'qmark' + }); + + envVars[dataframeSqlIntegrationEnvVarName] = dataframeSqlIntegrationCredentialsJson; + logger.debug(`SqlIntegrationEnvironmentVariablesProvider: Added env var for dataframe SQL integration`); + // Scan all cells for SQL integration IDs const integrationIds = this.scanNotebookForIntegrations(notebook); if (integrationIds.size === 0) { @@ -219,17 +232,7 @@ 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` - ); + // Internal DuckDB integration is handled above continue; } diff --git a/src/platform/notebooks/deepnote/sqlIntegrationEnvironmentVariablesProvider.unit.test.ts b/src/platform/notebooks/deepnote/sqlIntegrationEnvironmentVariablesProvider.unit.test.ts index 5a15e4f30..171d4a214 100644 --- a/src/platform/notebooks/deepnote/sqlIntegrationEnvironmentVariablesProvider.unit.test.ts +++ b/src/platform/notebooks/deepnote/sqlIntegrationEnvironmentVariablesProvider.unit.test.ts @@ -14,6 +14,10 @@ import { } from './integrationTypes'; import { mockedVSCodeNamespaces, resetVSCodeMocks } from '../../../test/vscode-mock'; +const EXPECTED_DATAFRAME_ONLY_ENV_VARS = { + SQL_DEEPNOTE_DATAFRAME_SQL: '{"url":"deepnote+duckdb:///:memory:","params":{},"param_style":"qmark"}' +}; + suite('SqlIntegrationEnvironmentVariablesProvider', () => { let provider: SqlIntegrationEnvironmentVariablesProvider; let integrationStorage: IntegrationStorage; @@ -55,7 +59,7 @@ suite('SqlIntegrationEnvironmentVariablesProvider', () => { when(mockedVSCodeNamespaces.workspace.notebookDocuments).thenReturn([notebook]); const envVars = await provider.getEnvironmentVariables(uri); - assert.deepStrictEqual(envVars, {}); + assert.deepStrictEqual(envVars, EXPECTED_DATAFRAME_ONLY_ENV_VARS); }); test('Returns empty object when SQL cells have no integration ID', async () => { @@ -67,7 +71,7 @@ suite('SqlIntegrationEnvironmentVariablesProvider', () => { when(mockedVSCodeNamespaces.workspace.notebookDocuments).thenReturn([notebook]); const envVars = await provider.getEnvironmentVariables(uri); - assert.deepStrictEqual(envVars, {}); + assert.deepStrictEqual(envVars, EXPECTED_DATAFRAME_ONLY_ENV_VARS); }); test('Returns environment variable for internal DuckDB integration', async () => { @@ -186,9 +190,9 @@ suite('SqlIntegrationEnvironmentVariablesProvider', () => { const envVars = await provider.getEnvironmentVariables(uri); - // Should only have one environment variable + // Should only have one environment variable apart from the internal DuckDB integration assert.property(envVars, 'SQL_MY_POSTGRES_DB'); - assert.strictEqual(Object.keys(envVars).length, 1); + assert.strictEqual(Object.keys(envVars).length, 2); }); test('Handles multiple SQL cells with different integrations', async () => { @@ -230,10 +234,10 @@ suite('SqlIntegrationEnvironmentVariablesProvider', () => { const envVars = await provider.getEnvironmentVariables(uri); - // Should have two environment variables + // Should have two environment variables apart from the internal DuckDB integration assert.property(envVars, 'SQL_MY_POSTGRES_DB'); assert.property(envVars, 'SQL_MY_BIGQUERY'); - assert.strictEqual(Object.keys(envVars).length, 2); + assert.strictEqual(Object.keys(envVars).length, 3); }); test('Handles missing integration configuration gracefully', async () => { @@ -251,8 +255,8 @@ suite('SqlIntegrationEnvironmentVariablesProvider', () => { const envVars = await provider.getEnvironmentVariables(uri); - // Should return empty object when integration config is missing - assert.deepStrictEqual(envVars, {}); + // Should return only dataframe integration when integration config is missing + assert.deepStrictEqual(envVars, EXPECTED_DATAFRAME_ONLY_ENV_VARS); }); test('Properly encodes special characters in PostgreSQL credentials', async () => { @@ -621,9 +625,9 @@ suite('SqlIntegrationEnvironmentVariablesProvider', () => { when(mockedVSCodeNamespaces.workspace.notebookDocuments).thenReturn([notebook]); when(integrationStorage.getIntegrationConfig(integrationId)).thenResolve(config); - // Should return empty object when unsupported auth method is encountered + // Should return only dataframe integration when unsupported auth method is encountered const envVars = await provider.getEnvironmentVariables(uri); - assert.deepStrictEqual(envVars, {}); + assert.deepStrictEqual(envVars, EXPECTED_DATAFRAME_ONLY_ENV_VARS); }); test('Skips unsupported Snowflake auth method (AZURE_AD)', async () => { @@ -647,7 +651,7 @@ suite('SqlIntegrationEnvironmentVariablesProvider', () => { when(integrationStorage.getIntegrationConfig(integrationId)).thenResolve(config); const envVars = await provider.getEnvironmentVariables(uri); - assert.deepStrictEqual(envVars, {}); + assert.deepStrictEqual(envVars, EXPECTED_DATAFRAME_ONLY_ENV_VARS); }); test('Skips unsupported Snowflake auth method (KEY_PAIR)', async () => { @@ -671,7 +675,7 @@ suite('SqlIntegrationEnvironmentVariablesProvider', () => { when(integrationStorage.getIntegrationConfig(integrationId)).thenResolve(config); const envVars = await provider.getEnvironmentVariables(uri); - assert.deepStrictEqual(envVars, {}); + assert.deepStrictEqual(envVars, EXPECTED_DATAFRAME_ONLY_ENV_VARS); }); }); });