Skip to content

Commit 4382e25

Browse files
authored
fix: Dataframe SQL block support (#118)
1 parent e2fa9b5 commit 4382e25

File tree

2 files changed

+25
-8
lines changed

2 files changed

+25
-8
lines changed

src/platform/notebooks/deepnote/sqlIntegrationEnvironmentVariablesProvider.ts

Lines changed: 17 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -149,6 +149,22 @@ export class SqlIntegrationEnvironmentVariablesProvider implements ISqlIntegrati
149149
}
150150

151151
try {
152+
// Handle internal DuckDB integration specially
153+
if (integrationId === DATAFRAME_SQL_INTEGRATION_ID) {
154+
const envVarName = convertToEnvironmentVariableName(getSqlEnvVarName(integrationId));
155+
const credentialsJson = JSON.stringify({
156+
url: 'deepnote+duckdb:///:memory:',
157+
params: {},
158+
param_style: 'qmark'
159+
});
160+
161+
envVars[envVarName] = credentialsJson;
162+
logger.debug(
163+
`SqlIntegrationEnvironmentVariablesProvider: Added env var for dataframe SQL integration`
164+
);
165+
continue;
166+
}
167+
152168
const config = await this.integrationStorage.getIntegrationConfig(integrationId);
153169
if (!config) {
154170
logger.warn(
@@ -162,7 +178,7 @@ export class SqlIntegrationEnvironmentVariablesProvider implements ISqlIntegrati
162178
const credentialsJson = convertIntegrationConfigToJson(config);
163179

164180
envVars[envVarName] = credentialsJson;
165-
logger.info(
181+
logger.debug(
166182
`SqlIntegrationEnvironmentVariablesProvider: Added env var ${envVarName} for integration ${integrationId}`
167183
);
168184
} catch (error) {
@@ -194,11 +210,6 @@ export class SqlIntegrationEnvironmentVariablesProvider implements ISqlIntegrati
194210
if (metadata && typeof metadata === 'object') {
195211
const integrationId = (metadata as Record<string, unknown>).sql_integration_id;
196212
if (typeof integrationId === 'string') {
197-
// Skip the internal DuckDB integration
198-
if (integrationId === DATAFRAME_SQL_INTEGRATION_ID) {
199-
continue;
200-
}
201-
202213
integrationIds.add(integrationId);
203214
logger.trace(
204215
`SqlIntegrationEnvironmentVariablesProvider: Found integration ${integrationId} in cell ${cell.index}`

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

Lines changed: 8 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -64,7 +64,7 @@ suite('SqlIntegrationEnvironmentVariablesProvider', () => {
6464
assert.deepStrictEqual(envVars, {});
6565
});
6666

67-
test('Skips internal DuckDB integration', async () => {
67+
test('Returns environment variable for internal DuckDB integration', async () => {
6868
const uri = Uri.file('/test/notebook.deepnote');
6969
const notebook = createMockNotebook(uri, [
7070
createMockCell(0, NotebookCellKind.Code, 'sql', 'SELECT * FROM df', {
@@ -75,7 +75,13 @@ suite('SqlIntegrationEnvironmentVariablesProvider', () => {
7575
when(mockedVSCodeNamespaces.workspace.notebookDocuments).thenReturn([notebook]);
7676

7777
const envVars = await provider.getEnvironmentVariables(uri);
78-
assert.deepStrictEqual(envVars, {});
78+
79+
// Check that the environment variable is set for dataframe SQL
80+
assert.property(envVars, 'SQL_DEEPNOTE_DATAFRAME_SQL');
81+
const credentialsJson = JSON.parse(envVars['SQL_DEEPNOTE_DATAFRAME_SQL']!);
82+
assert.strictEqual(credentialsJson.url, 'deepnote+duckdb:///:memory:');
83+
assert.deepStrictEqual(credentialsJson.params, {});
84+
assert.strictEqual(credentialsJson.param_style, 'qmark');
7985
});
8086

8187
test('Returns environment variable for PostgreSQL integration', async () => {

0 commit comments

Comments
 (0)