Skip to content

Commit 55d8454

Browse files
authored
feat: SQL block credentials propagation to Jupyter kernel (#78)
1 parent 614334b commit 55d8454

22 files changed

+1811
-61
lines changed

INTEGRATIONS_CREDENTIALS.md

Lines changed: 414 additions & 0 deletions
Large diffs are not rendered by default.

src/kernels/deepnote/deepnoteServerStarter.node.ts

Lines changed: 25 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -1,7 +1,7 @@
11
// Copyright (c) Microsoft Corporation.
22
// Licensed under the MIT License.
33

4-
import { inject, injectable, named } from 'inversify';
4+
import { inject, injectable, named, optional } from 'inversify';
55
import { CancellationToken, Uri } from 'vscode';
66
import { PythonEnvironment } from '../../platform/pythonEnvironments/info';
77
import { IDeepnoteServerStarter, IDeepnoteToolkitInstaller, DeepnoteServerInfo, DEEPNOTE_DEFAULT_PORT } from './types';
@@ -12,6 +12,7 @@ import { STANDARD_OUTPUT_CHANNEL } from '../../platform/common/constants';
1212
import { sleep } from '../../platform/common/utils/async';
1313
import { Cancellation, raceCancellationError } from '../../platform/common/cancellation';
1414
import { IExtensionSyncActivationService } from '../../platform/activation/types';
15+
import { ISqlIntegrationEnvVarsProvider } from '../../platform/notebooks/deepnote/types';
1516
import getPort from 'get-port';
1617
import * as fs from 'fs-extra';
1718
import * as os from 'os';
@@ -47,7 +48,10 @@ export class DeepnoteServerStarter implements IDeepnoteServerStarter, IExtension
4748
@inject(IDeepnoteToolkitInstaller) private readonly toolkitInstaller: IDeepnoteToolkitInstaller,
4849
@inject(IOutputChannel) @named(STANDARD_OUTPUT_CHANNEL) private readonly outputChannel: IOutputChannel,
4950
@inject(IHttpClient) private readonly httpClient: IHttpClient,
50-
@inject(IAsyncDisposableRegistry) asyncRegistry: IAsyncDisposableRegistry
51+
@inject(IAsyncDisposableRegistry) asyncRegistry: IAsyncDisposableRegistry,
52+
@inject(ISqlIntegrationEnvVarsProvider)
53+
@optional()
54+
private readonly sqlIntegrationEnvVars?: ISqlIntegrationEnvVarsProvider
5155
) {
5256
// Register for disposal when the extension deactivates
5357
asyncRegistry.push(this);
@@ -149,10 +153,28 @@ export class DeepnoteServerStarter implements IDeepnoteServerStarter, IExtension
149153

150154
// Detached mode ensures no requests are made to the backend (directly, or via proxy)
151155
// as there is no backend running in the extension, therefore:
152-
// 1. integration environment variables won't work / be injected
156+
// 1. integration environment variables are injected here instead
153157
// 2. post start hooks won't work / are not executed
154158
env.DEEPNOTE_RUNTIME__RUNNING_IN_DETACHED_MODE = 'true';
155159

160+
// Inject SQL integration environment variables
161+
if (this.sqlIntegrationEnvVars) {
162+
logger.debug(`DeepnoteServerStarter: Injecting SQL integration env vars for ${deepnoteFileUri.toString()}`);
163+
try {
164+
const sqlEnvVars = await this.sqlIntegrationEnvVars.getEnvironmentVariables(deepnoteFileUri, token);
165+
if (sqlEnvVars && Object.keys(sqlEnvVars).length > 0) {
166+
logger.debug(`DeepnoteServerStarter: Injecting ${Object.keys(sqlEnvVars).length} SQL env vars`);
167+
Object.assign(env, sqlEnvVars);
168+
} else {
169+
logger.debug('DeepnoteServerStarter: No SQL integration env vars to inject');
170+
}
171+
} catch (error) {
172+
logger.error('DeepnoteServerStarter: Failed to get SQL integration env vars', error.message);
173+
}
174+
} else {
175+
logger.debug('DeepnoteServerStarter: SqlIntegrationEnvironmentVariablesProvider not available');
176+
}
177+
156178
// Remove PYTHONHOME if it exists (can interfere with venv)
157179
delete env.PYTHONHOME;
158180

src/kernels/helpers.unit.test.ts

Lines changed: 311 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -646,4 +646,315 @@ suite('Kernel Connection Helpers', () => {
646646
assert.strictEqual(name, '.env (Python 9.8.7)');
647647
});
648648
});
649+
650+
suite('executeSilently', () => {
651+
interface MockKernelOptions {
652+
status: 'ok' | 'error';
653+
messages?: Array<{
654+
msg_type: 'stream' | 'error' | 'display_data' | 'execute_result';
655+
content: any;
656+
}>;
657+
errorContent?: {
658+
ename: string;
659+
evalue: string;
660+
traceback: string[];
661+
};
662+
}
663+
664+
function createMockKernel(options: MockKernelOptions) {
665+
return {
666+
requestExecute: () => {
667+
let resolvePromise: (value: any) => void;
668+
669+
// Create a promise that will be resolved after IOPub messages are dispatched
670+
const donePromise = new Promise<any>((resolve) => {
671+
resolvePromise = resolve;
672+
});
673+
674+
return {
675+
done: donePromise,
676+
set onIOPub(cb: (msg: any) => void) {
677+
// Invoke IOPub callback synchronously with all messages
678+
if (options.messages && options.messages.length > 0) {
679+
options.messages.forEach((msg) => {
680+
cb({
681+
header: { msg_type: msg.msg_type },
682+
content: msg.content
683+
});
684+
});
685+
}
686+
// Resolve the done promise after messages are dispatched
687+
resolvePromise({
688+
content:
689+
options.status === 'ok'
690+
? { status: 'ok' as const }
691+
: {
692+
status: 'error' as const,
693+
...options.errorContent
694+
}
695+
});
696+
}
697+
};
698+
}
699+
};
700+
}
701+
702+
test('Returns outputs from kernel execution', async () => {
703+
const mockKernel = createMockKernel({
704+
status: 'ok',
705+
messages: [
706+
{
707+
msg_type: 'stream',
708+
content: {
709+
name: 'stdout',
710+
text: 'hello\n'
711+
}
712+
}
713+
]
714+
});
715+
716+
const code = 'print("hello")';
717+
const { executeSilently } = await import('./helpers');
718+
const result = await executeSilently(mockKernel as any, code);
719+
720+
// executeSilently should return outputs array with collected stream output
721+
assert.isArray(result);
722+
assert.equal(result.length, 1);
723+
assert.equal(result[0].output_type, 'stream');
724+
assert.equal((result[0] as any).name, 'stdout');
725+
assert.equal((result[0] as any).text, 'hello\n');
726+
});
727+
728+
test('Collects error outputs', async () => {
729+
const mockKernel = createMockKernel({
730+
status: 'error',
731+
errorContent: {
732+
ename: 'NameError',
733+
evalue: 'name not defined',
734+
traceback: ['Traceback...']
735+
},
736+
messages: [
737+
{
738+
msg_type: 'error',
739+
content: {
740+
ename: 'NameError',
741+
evalue: 'name not defined',
742+
traceback: ['Traceback...']
743+
}
744+
}
745+
]
746+
});
747+
748+
const code = 'undefined_variable';
749+
const { executeSilently } = await import('./helpers');
750+
const result = await executeSilently(mockKernel as any, code);
751+
752+
assert.isArray(result);
753+
assert.equal(result.length, 1);
754+
assert.equal(result[0].output_type, 'error');
755+
assert.equal((result[0] as any).ename, 'NameError');
756+
assert.equal((result[0] as any).evalue, 'name not defined');
757+
assert.deepStrictEqual((result[0] as any).traceback, ['Traceback...']);
758+
});
759+
760+
test('Collects display_data outputs', async () => {
761+
const mockKernel = createMockKernel({
762+
status: 'ok',
763+
messages: [
764+
{
765+
msg_type: 'display_data',
766+
content: {
767+
data: {
768+
'text/plain': 'some data'
769+
},
770+
metadata: {}
771+
}
772+
}
773+
]
774+
});
775+
776+
const code = 'display("data")';
777+
const { executeSilently } = await import('./helpers');
778+
const result = await executeSilently(mockKernel as any, code);
779+
780+
assert.isArray(result);
781+
assert.equal(result.length, 1);
782+
assert.equal(result[0].output_type, 'display_data');
783+
assert.deepStrictEqual((result[0] as any).data, { 'text/plain': 'some data' });
784+
assert.deepStrictEqual((result[0] as any).metadata, {});
785+
});
786+
787+
test('Handles multiple outputs', async () => {
788+
const mockKernel = createMockKernel({
789+
status: 'ok',
790+
messages: [
791+
{
792+
msg_type: 'stream',
793+
content: {
794+
name: 'stdout',
795+
text: 'output 1'
796+
}
797+
},
798+
{
799+
msg_type: 'stream',
800+
content: {
801+
name: 'stdout',
802+
text: 'output 2'
803+
}
804+
}
805+
]
806+
});
807+
808+
const code = 'print("1"); print("2")';
809+
const { executeSilently } = await import('./helpers');
810+
const result = await executeSilently(mockKernel as any, code);
811+
812+
assert.isArray(result);
813+
// Consecutive stream messages with the same name are concatenated
814+
assert.equal(result.length, 1);
815+
assert.equal(result[0].output_type, 'stream');
816+
assert.equal((result[0] as any).name, 'stdout');
817+
assert.equal((result[0] as any).text, 'output 1output 2');
818+
});
819+
820+
test('Collects execute_result outputs', async () => {
821+
const mockKernel = createMockKernel({
822+
status: 'ok',
823+
messages: [
824+
{
825+
msg_type: 'execute_result',
826+
content: {
827+
data: {
828+
'text/plain': '42'
829+
},
830+
metadata: {},
831+
execution_count: 1
832+
}
833+
}
834+
]
835+
});
836+
837+
const code = '42';
838+
const { executeSilently } = await import('./helpers');
839+
const result = await executeSilently(mockKernel as any, code);
840+
841+
assert.isArray(result);
842+
assert.equal(result.length, 1);
843+
assert.equal(result[0].output_type, 'execute_result');
844+
assert.deepStrictEqual((result[0] as any).data, { 'text/plain': '42' });
845+
assert.deepStrictEqual((result[0] as any).metadata, {});
846+
assert.equal((result[0] as any).execution_count, 1);
847+
});
848+
849+
test('Stream messages with different names produce separate outputs', async () => {
850+
const mockKernel = createMockKernel({
851+
status: 'ok',
852+
messages: [
853+
{
854+
msg_type: 'stream',
855+
content: {
856+
name: 'stdout',
857+
text: 'standard output'
858+
}
859+
},
860+
{
861+
msg_type: 'stream',
862+
content: {
863+
name: 'stderr',
864+
text: 'error output'
865+
}
866+
},
867+
{
868+
msg_type: 'stream',
869+
content: {
870+
name: 'stdout',
871+
text: ' more stdout'
872+
}
873+
}
874+
]
875+
});
876+
877+
const code = 'print("test")';
878+
const { executeSilently } = await import('./helpers');
879+
const result = await executeSilently(mockKernel as any, code);
880+
881+
assert.isArray(result);
882+
// Should have 3 outputs: stdout, stderr, stdout (not concatenated because stderr is in between)
883+
assert.equal(result.length, 3);
884+
assert.equal(result[0].output_type, 'stream');
885+
assert.equal((result[0] as any).name, 'stdout');
886+
assert.equal((result[0] as any).text, 'standard output');
887+
assert.equal(result[1].output_type, 'stream');
888+
assert.equal((result[1] as any).name, 'stderr');
889+
assert.equal((result[1] as any).text, 'error output');
890+
assert.equal(result[2].output_type, 'stream');
891+
assert.equal((result[2] as any).name, 'stdout');
892+
assert.equal((result[2] as any).text, ' more stdout');
893+
});
894+
895+
test('errorOptions with traceErrors logs errors', async () => {
896+
const mockKernel = createMockKernel({
897+
status: 'error',
898+
errorContent: {
899+
ename: 'ValueError',
900+
evalue: 'invalid value',
901+
traceback: ['Traceback (most recent call last):', ' File "<stdin>", line 1']
902+
},
903+
messages: [
904+
{
905+
msg_type: 'error',
906+
content: {
907+
ename: 'ValueError',
908+
evalue: 'invalid value',
909+
traceback: ['Traceback (most recent call last):', ' File "<stdin>", line 1']
910+
}
911+
}
912+
]
913+
});
914+
915+
const code = 'raise ValueError("invalid value")';
916+
const { executeSilently } = await import('./helpers');
917+
const result = await executeSilently(mockKernel as any, code, {
918+
traceErrors: true,
919+
traceErrorsMessage: 'Custom error message'
920+
});
921+
922+
assert.isArray(result);
923+
assert.equal(result.length, 1);
924+
assert.equal(result[0].output_type, 'error');
925+
assert.equal((result[0] as any).ename, 'ValueError');
926+
});
927+
928+
test('errorOptions without traceErrors still collects errors', async () => {
929+
const mockKernel = createMockKernel({
930+
status: 'error',
931+
errorContent: {
932+
ename: 'RuntimeError',
933+
evalue: 'runtime issue',
934+
traceback: ['Traceback...']
935+
},
936+
messages: [
937+
{
938+
msg_type: 'error',
939+
content: {
940+
ename: 'RuntimeError',
941+
evalue: 'runtime issue',
942+
traceback: ['Traceback...']
943+
}
944+
}
945+
]
946+
});
947+
948+
const code = 'raise RuntimeError("runtime issue")';
949+
const { executeSilently } = await import('./helpers');
950+
const result = await executeSilently(mockKernel as any, code, {
951+
traceErrors: false
952+
});
953+
954+
assert.isArray(result);
955+
assert.equal(result.length, 1);
956+
assert.equal(result[0].output_type, 'error');
957+
assert.equal((result[0] as any).ename, 'RuntimeError');
958+
});
959+
});
649960
});

src/kernels/kernel.ts

Lines changed: 13 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -853,10 +853,22 @@ abstract class BaseKernel implements IBaseKernel {
853853

854854
// Gather all of the startup code at one time and execute as one cell
855855
const startupCode = await this.gatherInternalStartupCode();
856-
await this.executeSilently(session, startupCode, {
856+
logger.trace(`Executing startup code with ${startupCode.length} lines`);
857+
858+
const outputs = await this.executeSilently(session, startupCode, {
857859
traceErrors: true,
858860
traceErrorsMessage: 'Error executing jupyter extension internal startup code'
859861
});
862+
logger.trace(`Startup code execution completed with ${outputs?.length || 0} outputs`);
863+
if (outputs && outputs.length > 0) {
864+
// Avoid logging content; output types only.
865+
logger.trace(
866+
`Startup code produced ${outputs.length} output(s): ${outputs
867+
.map((o) => o.output_type)
868+
.join(', ')}`
869+
);
870+
}
871+
860872
// Run user specified startup commands
861873
await this.executeSilently(session, this.getUserStartupCommands(), { traceErrors: false });
862874
}

0 commit comments

Comments
 (0)