Skip to content

Commit 915fa56

Browse files
jankucajankuca
committed
feat: setting env vars for connected integrations
1 parent 953a32a commit 915fa56

File tree

8 files changed

+695
-8
lines changed

8 files changed

+695
-8
lines changed

‎src/kernels/deepnote/deepnoteServerStarter.node.ts‎

Lines changed: 35 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';
@@ -17,6 +17,7 @@ import * as fs from 'fs-extra';
1717
import * as os from 'os';
1818
import * as path from '../../platform/vscode-path/path';
1919
import { generateUuid } from '../../platform/common/uuid';
20+
import { SqlIntegrationEnvironmentVariablesProvider } from '../../notebooks/deepnote/integrations/sqlIntegrationEnvironmentVariablesProvider';
2021

2122
/**
2223
* Lock file data structure for tracking server ownership
@@ -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(SqlIntegrationEnvironmentVariablesProvider)
53+
@optional()
54+
private readonly sqlIntegrationEnvVars?: SqlIntegrationEnvironmentVariablesProvider
5155
) {
5256
// Register for disposal when the extension deactivates
5357
asyncRegistry.push(this);
@@ -149,10 +153,38 @@ 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.info(`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.info(
167+
`DeepnoteServerStarter: Injecting ${
168+
Object.keys(sqlEnvVars).length
169+
} SQL integration env vars: ${Object.keys(sqlEnvVars).join(', ')}`
170+
);
171+
Object.assign(env, sqlEnvVars);
172+
// Log the first 100 chars of each env var value for debugging
173+
for (const [key, value] of Object.entries(sqlEnvVars)) {
174+
if (value) {
175+
logger.info(`DeepnoteServerStarter: ${key} = ${value.substring(0, 100)}...`);
176+
}
177+
}
178+
} else {
179+
logger.info('DeepnoteServerStarter: No SQL integration env vars to inject');
180+
}
181+
} catch (error) {
182+
logger.error('DeepnoteServerStarter: Failed to get SQL integration env vars', error);
183+
}
184+
} else {
185+
logger.info('DeepnoteServerStarter: SqlIntegrationEnvironmentVariablesProvider not available');
186+
}
187+
156188
// Remove PYTHONHOME if it exists (can interfere with venv)
157189
delete env.PYTHONHOME;
158190

‎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.info(`Executing startup code with ${startupCode.length} lines`);
857+
const outputs = await this.executeSilently(session, startupCode, {
857858
traceErrors: true,
858859
traceErrorsMessage: 'Error executing jupyter extension internal startup code'
859860
});
861+
logger.info(`Startup code execution completed with ${outputs?.length || 0} outputs`);
862+
if (outputs && outputs.length > 0) {
863+
outputs.forEach((output, idx) => {
864+
logger.info(
865+
`Startup code output ${idx}: ${output.output_type} - ${JSON.stringify(output).substring(
866+
0,
867+
200
868+
)}`
869+
);
870+
});
871+
}
860872
// Run user specified startup commands
861873
await this.executeSilently(session, this.getUserStartupCommands(), { traceErrors: false });
862874
}

‎src/kernels/raw/launcher/kernelEnvVarsService.node.ts‎

Lines changed: 57 additions & 4 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 } from 'inversify';
4+
import { inject, injectable, optional } from 'inversify';
55
import { logger } from '../../../platform/logging';
66
import { getDisplayPath } from '../../../platform/common/platform/fs-paths';
77
import { IConfigurationService, Resource, type ReadWrite } from '../../../platform/common/types';
@@ -18,6 +18,7 @@ import { IJupyterKernelSpec } from '../../types';
1818
import { CancellationToken, Uri } from 'vscode';
1919
import { PYTHON_LANGUAGE } from '../../../platform/common/constants';
2020
import { trackKernelResourceInformation } from '../../telemetry/helper';
21+
import { SqlIntegrationEnvironmentVariablesProvider } from '../../../notebooks/deepnote/integrations/sqlIntegrationEnvironmentVariablesProvider';
2122

2223
/**
2324
* Class used to fetch environment variables for a kernel.
@@ -30,8 +31,17 @@ export class KernelEnvironmentVariablesService {
3031
@inject(IEnvironmentVariablesService) private readonly envVarsService: IEnvironmentVariablesService,
3132
@inject(ICustomEnvironmentVariablesProvider)
3233
private readonly customEnvVars: ICustomEnvironmentVariablesProvider,
33-
@inject(IConfigurationService) private readonly configService: IConfigurationService
34-
) {}
34+
@inject(IConfigurationService) private readonly configService: IConfigurationService,
35+
@inject(SqlIntegrationEnvironmentVariablesProvider)
36+
@optional()
37+
private readonly sqlIntegrationEnvVars?: SqlIntegrationEnvironmentVariablesProvider
38+
) {
39+
logger.info(
40+
`KernelEnvironmentVariablesService: Constructor called, sqlIntegrationEnvVars is ${
41+
sqlIntegrationEnvVars ? 'AVAILABLE' : 'UNDEFINED'
42+
}`
43+
);
44+
}
3545
/**
3646
* Generates the environment variables for the kernel.
3747
*
@@ -51,6 +61,11 @@ export class KernelEnvironmentVariablesService {
5161
kernelSpec: IJupyterKernelSpec,
5262
token?: CancellationToken
5363
) {
64+
logger.info(
65+
`KernelEnvVarsService.getEnvironmentVariables: Called for resource ${
66+
resource?.toString() || 'undefined'
67+
}, sqlIntegrationEnvVars is ${this.sqlIntegrationEnvVars ? 'AVAILABLE' : 'UNDEFINED'}`
68+
);
5469
let kernelEnv =
5570
kernelSpec.env && Object.keys(kernelSpec.env).length > 0
5671
? (Object.assign({}, kernelSpec.env) as ReadWrite<NodeJS.ProcessEnv>)
@@ -68,7 +83,7 @@ export class KernelEnvironmentVariablesService {
6883
if (token?.isCancellationRequested) {
6984
return;
7085
}
71-
let [customEnvVars, interpreterEnv] = await Promise.all([
86+
let [customEnvVars, interpreterEnv, sqlIntegrationEnvVars] = await Promise.all([
7287
this.customEnvVars
7388
.getCustomEnvironmentVariables(resource, isPythonKernel ? 'RunPythonCode' : 'RunNonPythonCode', token)
7489
.catch(noop),
@@ -82,6 +97,24 @@ export class KernelEnvironmentVariablesService {
8297
);
8398
return undefined;
8499
})
100+
: undefined,
101+
this.sqlIntegrationEnvVars
102+
? this.sqlIntegrationEnvVars
103+
.getEnvironmentVariables(resource, token)
104+
.then((vars) => {
105+
if (vars && Object.keys(vars).length > 0) {
106+
logger.info(
107+
`KernelEnvVarsService: Got ${
108+
Object.keys(vars).length
109+
} SQL integration env vars: ${Object.keys(vars).join(', ')}`
110+
);
111+
}
112+
return vars;
113+
})
114+
.catch<undefined>((ex) => {
115+
logger.error('Failed to get SQL integration env variables for Kernel', ex);
116+
return undefined;
117+
})
85118
: undefined
86119
]);
87120
if (token?.isCancellationRequested) {
@@ -123,6 +156,16 @@ export class KernelEnvironmentVariablesService {
123156

124157
Object.assign(mergedVars, interpreterEnv, kernelEnv); // kernels vars win over interpreter.
125158

159+
// Merge SQL integration environment variables
160+
if (sqlIntegrationEnvVars) {
161+
logger.info(
162+
`KernelEnvVarsService: Merging ${
163+
Object.keys(sqlIntegrationEnvVars).length
164+
} SQL integration env vars into kernel env`
165+
);
166+
Object.assign(mergedVars, sqlIntegrationEnvVars);
167+
}
168+
126169
// If user asks us to, set PYTHONNOUSERSITE
127170
// For more details see here https://github.com/microsoft/vscode-jupyter/issues/8553#issuecomment-997144591
128171
// https://docs.python.org/3/library/site.html#site.ENABLE_USER_SITE
@@ -138,6 +181,16 @@ export class KernelEnvironmentVariablesService {
138181
// We can support this, however since this has not been requested, lets not do it.'
139182
this.envVarsService.mergeVariables(kernelEnv, mergedVars); // kernels vars win over interpreter.
140183
this.envVarsService.mergeVariables(customEnvVars, mergedVars); // custom vars win over all.
184+
185+
// Merge SQL integration environment variables
186+
if (sqlIntegrationEnvVars) {
187+
logger.info(
188+
`KernelEnvVarsService: Merging ${
189+
Object.keys(sqlIntegrationEnvVars).length
190+
} SQL integration env vars into kernel env (non-python)`
191+
);
192+
this.envVarsService.mergeVariables(sqlIntegrationEnvVars, mergedVars);
193+
}
141194
}
142195

143196
// env variables in kernelSpecs can contain variables that need to be substituted

‎src/kernels/serviceRegistry.node.ts‎

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -48,6 +48,7 @@ import { LastCellExecutionTracker } from './execution/lastCellExecutionTracker';
4848
import { ClearJupyterServersCommand } from './jupyter/clearJupyterServersCommand';
4949
import { KernelChatStartupCodeProvider } from './chat/kernelStartupCodeProvider';
5050
import { KernelWorkingDirectory } from './raw/session/kernelWorkingDirectory.node';
51+
import { SqlIntegrationEnvironmentVariablesProvider } from '../notebooks/deepnote/integrations/sqlIntegrationEnvironmentVariablesProvider';
5152

5253
export function registerTypes(serviceManager: IServiceManager, isDevMode: boolean) {
5354
serviceManager.addSingleton<IExtensionSyncActivationService>(IExtensionSyncActivationService, Activation);
@@ -62,6 +63,10 @@ export function registerTypes(serviceManager: IServiceManager, isDevMode: boolea
6263
);
6364
serviceManager.addSingleton<IRawKernelSessionFactory>(IRawKernelSessionFactory, RawKernelSessionFactory);
6465
serviceManager.addSingleton<IKernelLauncher>(IKernelLauncher, KernelLauncher);
66+
serviceManager.addSingleton<SqlIntegrationEnvironmentVariablesProvider>(
67+
SqlIntegrationEnvironmentVariablesProvider,
68+
SqlIntegrationEnvironmentVariablesProvider
69+
);
6570
serviceManager.addSingleton<KernelEnvironmentVariablesService>(
6671
KernelEnvironmentVariablesService,
6772
KernelEnvironmentVariablesService

0 commit comments

Comments
 (0)