Skip to content

Commit 80392e7

Browse files
tkislantkislan
authored
fix: Fix SQL dataframe integration (#139)
* fix: Fix SQL dataframe integration Signed-off-by: Tomas Kislan <tomas@kislan.sk> * test: Update test expectations Signed-off-by: Tomas Kislan <tomas@kislan.sk> * test: Update comments Signed-off-by: Tomas Kislan <tomas@kislan.sk> * test: Update comments Signed-off-by: Tomas Kislan <tomas@kislan.sk> * test: Rename empty env vars test variable Signed-off-by: Tomas Kislan <tomas@kislan.sk> --------- Signed-off-by: Tomas Kislan <tomas@kislan.sk>
1 parent c1b5526 commit 80392e7

File tree

2 files changed

+30
-23
lines changed

2 files changed

+30
-23
lines changed

‎src/platform/notebooks/deepnote/sqlIntegrationEnvironmentVariablesProvider.ts‎

Lines changed: 14 additions & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -199,6 +199,19 @@ export class SqlIntegrationEnvironmentVariablesProvider implements ISqlIntegrati
199199
return envVars;
200200
}
201201

202+
// Always add the internal DuckDB integration
203+
const dataframeSqlIntegrationEnvVarName = convertToEnvironmentVariableName(
204+
getSqlEnvVarName(DATAFRAME_SQL_INTEGRATION_ID)
205+
);
206+
const dataframeSqlIntegrationCredentialsJson = JSON.stringify({
207+
url: 'deepnote+duckdb:///:memory:',
208+
params: {},
209+
param_style: 'qmark'
210+
});
211+
212+
envVars[dataframeSqlIntegrationEnvVarName] = dataframeSqlIntegrationCredentialsJson;
213+
logger.debug(`SqlIntegrationEnvironmentVariablesProvider: Added env var for dataframe SQL integration`);
214+
202215
// Scan all cells for SQL integration IDs
203216
const integrationIds = this.scanNotebookForIntegrations(notebook);
204217
if (integrationIds.size === 0) {
@@ -219,17 +232,7 @@ export class SqlIntegrationEnvironmentVariablesProvider implements ISqlIntegrati
219232
try {
220233
// Handle internal DuckDB integration specially
221234
if (integrationId === DATAFRAME_SQL_INTEGRATION_ID) {
222-
const envVarName = convertToEnvironmentVariableName(getSqlEnvVarName(integrationId));
223-
const credentialsJson = JSON.stringify({
224-
url: 'deepnote+duckdb:///:memory:',
225-
params: {},
226-
param_style: 'qmark'
227-
});
228-
229-
envVars[envVarName] = credentialsJson;
230-
logger.debug(
231-
`SqlIntegrationEnvironmentVariablesProvider: Added env var for dataframe SQL integration`
232-
);
235+
// Internal DuckDB integration is handled above
233236
continue;
234237
}
235238

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

Lines changed: 16 additions & 12 deletions
Original file line numberDiff line numberDiff line change
@@ -15,6 +15,10 @@ import {
1515
} from './integrationTypes';
1616
import { mockedVSCodeNamespaces, resetVSCodeMocks } from '../../../test/vscode-mock';
1717

18+
const EXPECTED_DATAFRAME_ONLY_ENV_VARS = {
19+
SQL_DEEPNOTE_DATAFRAME_SQL: '{"url":"deepnote+duckdb:///:memory:","params":{},"param_style":"qmark"}'
20+
};
21+
1822
suite('SqlIntegrationEnvironmentVariablesProvider', () => {
1923
let provider: SqlIntegrationEnvironmentVariablesProvider;
2024
let integrationStorage: IntegrationStorage;
@@ -56,7 +60,7 @@ suite('SqlIntegrationEnvironmentVariablesProvider', () => {
5660
when(mockedVSCodeNamespaces.workspace.notebookDocuments).thenReturn([notebook]);
5761

5862
const envVars = await provider.getEnvironmentVariables(uri);
59-
assert.deepStrictEqual(envVars, {});
63+
assert.deepStrictEqual(envVars, EXPECTED_DATAFRAME_ONLY_ENV_VARS);
6064
});
6165

6266
test('Returns empty object when SQL cells have no integration ID', async () => {
@@ -68,7 +72,7 @@ suite('SqlIntegrationEnvironmentVariablesProvider', () => {
6872
when(mockedVSCodeNamespaces.workspace.notebookDocuments).thenReturn([notebook]);
6973

7074
const envVars = await provider.getEnvironmentVariables(uri);
71-
assert.deepStrictEqual(envVars, {});
75+
assert.deepStrictEqual(envVars, EXPECTED_DATAFRAME_ONLY_ENV_VARS);
7276
});
7377

7478
test('Returns environment variable for internal DuckDB integration', async () => {
@@ -187,9 +191,9 @@ suite('SqlIntegrationEnvironmentVariablesProvider', () => {
187191

188192
const envVars = await provider.getEnvironmentVariables(uri);
189193

190-
// Should only have one environment variable
194+
// Should only have one environment variable apart from the internal DuckDB integration
191195
assert.property(envVars, 'SQL_MY_POSTGRES_DB');
192-
assert.strictEqual(Object.keys(envVars).length, 1);
196+
assert.strictEqual(Object.keys(envVars).length, 2);
193197
});
194198

195199
test('Handles multiple SQL cells with different integrations', async () => {
@@ -231,10 +235,10 @@ suite('SqlIntegrationEnvironmentVariablesProvider', () => {
231235

232236
const envVars = await provider.getEnvironmentVariables(uri);
233237

234-
// Should have two environment variables
238+
// Should have two environment variables apart from the internal DuckDB integration
235239
assert.property(envVars, 'SQL_MY_POSTGRES_DB');
236240
assert.property(envVars, 'SQL_MY_BIGQUERY');
237-
assert.strictEqual(Object.keys(envVars).length, 2);
241+
assert.strictEqual(Object.keys(envVars).length, 3);
238242
});
239243

240244
test('Handles missing integration configuration gracefully', async () => {
@@ -252,8 +256,8 @@ suite('SqlIntegrationEnvironmentVariablesProvider', () => {
252256

253257
const envVars = await provider.getEnvironmentVariables(uri);
254258

255-
// Should return empty object when integration config is missing
256-
assert.deepStrictEqual(envVars, {});
259+
// Should return only dataframe integration when integration config is missing
260+
assert.deepStrictEqual(envVars, EXPECTED_DATAFRAME_ONLY_ENV_VARS);
257261
});
258262

259263
test('Properly encodes special characters in PostgreSQL credentials', async () => {
@@ -622,9 +626,9 @@ suite('SqlIntegrationEnvironmentVariablesProvider', () => {
622626
when(mockedVSCodeNamespaces.workspace.notebookDocuments).thenReturn([notebook]);
623627
when(integrationStorage.getIntegrationConfig(integrationId)).thenResolve(config);
624628

625-
// Should return empty object when unsupported auth method is encountered
629+
// Should return only dataframe integration when unsupported auth method is encountered
626630
const envVars = await provider.getEnvironmentVariables(uri);
627-
assert.deepStrictEqual(envVars, {});
631+
assert.deepStrictEqual(envVars, EXPECTED_DATAFRAME_ONLY_ENV_VARS);
628632
});
629633

630634
test('Skips unsupported Snowflake auth method (AZURE_AD)', async () => {
@@ -648,7 +652,7 @@ suite('SqlIntegrationEnvironmentVariablesProvider', () => {
648652
when(integrationStorage.getIntegrationConfig(integrationId)).thenResolve(config);
649653

650654
const envVars = await provider.getEnvironmentVariables(uri);
651-
assert.deepStrictEqual(envVars, {});
655+
assert.deepStrictEqual(envVars, EXPECTED_DATAFRAME_ONLY_ENV_VARS);
652656
});
653657

654658
test('Skips unsupported Snowflake auth method (KEY_PAIR)', async () => {
@@ -672,7 +676,7 @@ suite('SqlIntegrationEnvironmentVariablesProvider', () => {
672676
when(integrationStorage.getIntegrationConfig(integrationId)).thenResolve(config);
673677

674678
const envVars = await provider.getEnvironmentVariables(uri);
675-
assert.deepStrictEqual(envVars, {});
679+
assert.deepStrictEqual(envVars, EXPECTED_DATAFRAME_ONLY_ENV_VARS);
676680
});
677681
});
678682
});

0 commit comments

Comments
 (0)