Skip to content

Commit

Permalink
Poc lambda subhosting (#10192)
Browse files Browse the repository at this point in the history
Add `SERVERLESS_LAMBDA_SUBHOSTING_ROLE` to allow hosting lambdas in
separate aws account

Tested with old configuration and new configuration
  • Loading branch information
martmull authored Feb 14, 2025
1 parent 9cbcf62 commit f2da915
Show file tree
Hide file tree
Showing 7 changed files with 945 additions and 25 deletions.
1 change: 1 addition & 0 deletions package.json
Original file line number Diff line number Diff line change
Expand Up @@ -5,6 +5,7 @@
"@apollo/server": "^4.7.3",
"@aws-sdk/client-lambda": "^3.614.0",
"@aws-sdk/client-s3": "^3.363.0",
"@aws-sdk/client-sts": "^3.744.0",
"@aws-sdk/credential-providers": "^3.363.0",
"@blocknote/mantine": "^0.22.0",
"@blocknote/react": "^0.22.0",
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -409,9 +409,17 @@ export class EnvironmentVariables {
})
@ValidateIf((env) => env.SERVERLESS_TYPE === ServerlessDriverType.Lambda)
@IsString()
@IsOptional()
SERVERLESS_LAMBDA_ROLE: string;

@EnvironmentVariablesMetadata({
group: EnvironmentVariablesGroup.ServerlessConfig,
description: 'Role to assume when hosting lambdas in dedicated AWS account',
})
@ValidateIf((env) => env.SERVERLESS_TYPE === ServerlessDriverType.Lambda)
@IsString()
@IsOptional()
SERVERLESS_LAMBDA_SUBHOSTING_ROLE?: string;

@EnvironmentVariablesMetadata({
group: EnvironmentVariablesGroup.ServerlessConfig,
sensitive: true,
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -21,6 +21,7 @@ import {
UpdateFunctionConfigurationCommandInput,
waitUntilFunctionUpdatedV2,
} from '@aws-sdk/client-lambda';
import { AssumeRoleCommand, STSClient } from '@aws-sdk/client-sts';
import { CreateFunctionCommandInput } from '@aws-sdk/client-lambda/dist-types/commands/CreateFunctionCommand';
import { UpdateFunctionCodeCommandInput } from '@aws-sdk/client-lambda/dist-types/commands/UpdateFunctionCodeCommand';
import dotenv from 'dotenv';
Expand Down Expand Up @@ -55,34 +56,84 @@ import {
} from 'src/engine/metadata-modules/serverless-function/serverless-function.exception';

const UPDATE_FUNCTION_DURATION_TIMEOUT_IN_SECONDS = 60;
const CREDENTIALS_DURATION_IN_SECONDS = 10 * 60 * 60; // 10h

export interface LambdaDriverOptions extends LambdaClientConfig {
fileStorageService: FileStorageService;
region: string;
role: string;
lambdaRole: string;
subhostingRole?: string;
}

export class LambdaDriver implements ServerlessDriver {
private readonly lambdaClient: Lambda;
private readonly lambdaRole: string;
private lambdaClient: Lambda | undefined;
private credentialsExpiry: Date | null = null;
private readonly options: LambdaDriverOptions;
private readonly fileStorageService: FileStorageService;

constructor(options: LambdaDriverOptions) {
const { region, role, ...lambdaOptions } = options;

this.lambdaClient = new Lambda({ ...lambdaOptions, region });
this.lambdaRole = role;
this.options = options;
this.lambdaClient = undefined;
this.fileStorageService = options.fileStorageService;
}

private async getLambdaClient() {
if (
!isDefined(this.lambdaClient) ||
(isDefined(this.options.subhostingRole) &&
isDefined(this.credentialsExpiry) &&
new Date() >= this.credentialsExpiry)
) {
this.lambdaClient = new Lambda({
...this.options,
...(isDefined(this.options.subhostingRole) && {
credentials: await this.getAssumeRoleCredentials(),
}),
});
}

return this.lambdaClient;
}

private async getAssumeRoleCredentials() {
const stsClient = new STSClient({ region: this.options.region });

this.credentialsExpiry = new Date(
Date.now() + (CREDENTIALS_DURATION_IN_SECONDS - 60 * 5) * 1000,
);

const assumeRoleCommand = new AssumeRoleCommand({
RoleArn: 'arn:aws:iam::820242914089:role/LambdaDeploymentRole',
RoleSessionName: 'LambdaSession',
DurationSeconds: CREDENTIALS_DURATION_IN_SECONDS,
});

const { Credentials } = await stsClient.send(assumeRoleCommand);

if (
!isDefined(Credentials) ||
!isDefined(Credentials.AccessKeyId) ||
!isDefined(Credentials.SecretAccessKey) ||
!isDefined(Credentials.SessionToken)
) {
throw new Error('Failed to assume role');
}

return {
accessKeyId: Credentials.AccessKeyId,
secretAccessKey: Credentials.SecretAccessKey,
sessionToken: Credentials.SessionToken,
};
}

private async waitFunctionUpdates(
serverlessFunctionId: string,
maxWaitTime: number = UPDATE_FUNCTION_DURATION_TIMEOUT_IN_SECONDS,
) {
const waitParams = { FunctionName: serverlessFunctionId };

await waitUntilFunctionUpdatedV2(
{ client: this.lambdaClient, maxWaitTime },
{ client: await this.getLambdaClient(), maxWaitTime },
waitParams,
);
}
Expand All @@ -93,7 +144,9 @@ export class LambdaDriver implements ServerlessDriver {
MaxItems: 1,
};
const listLayerCommand = new ListLayerVersionsCommand(listLayerParams);
const listLayerResult = await this.lambdaClient.send(listLayerCommand);
const listLayerResult = await (
await this.getLambdaClient()
).send(listLayerCommand);

if (
isDefined(listLayerResult.LayerVersions) &&
Expand Down Expand Up @@ -128,7 +181,7 @@ export class LambdaDriver implements ServerlessDriver {

const command = new PublishLayerVersionCommand(params);

const result = await this.lambdaClient.send(command);
const result = await (await this.getLambdaClient()).send(command);

await lambdaBuildDirectoryManager.clean();

Expand All @@ -145,7 +198,7 @@ export class LambdaDriver implements ServerlessDriver {
FunctionName: functionName,
});

await this.lambdaClient.send(getFunctionCommand);
await (await this.getLambdaClient()).send(getFunctionCommand);

return true;
} catch (error) {
Expand All @@ -166,7 +219,7 @@ export class LambdaDriver implements ServerlessDriver {
FunctionName: serverlessFunction.id,
});

await this.lambdaClient.send(deleteFunctionCommand);
await (await this.getLambdaClient()).send(deleteFunctionCommand);
}
}

Expand Down Expand Up @@ -232,15 +285,15 @@ export class LambdaDriver implements ServerlessDriver {
Environment: {
Variables: envVariables,
},
Role: this.lambdaRole,
Role: this.options.lambdaRole,
Runtime: serverlessFunction.runtime,
Description: 'Lambda function to run user script',
Timeout: serverlessFunction.timeoutSeconds,
};

const command = new CreateFunctionCommand(params);

await this.lambdaClient.send(command);
await (await this.getLambdaClient()).send(command);
} else {
const updateCodeParams: UpdateFunctionCodeCommandInput = {
ZipFile: await fs.readFile(lambdaZipPath),
Expand All @@ -249,7 +302,7 @@ export class LambdaDriver implements ServerlessDriver {

const updateCodeCommand = new UpdateFunctionCodeCommand(updateCodeParams);

await this.lambdaClient.send(updateCodeCommand);
await (await this.getLambdaClient()).send(updateCodeCommand);

const updateConfigurationParams: UpdateFunctionConfigurationCommandInput =
{
Expand All @@ -266,7 +319,7 @@ export class LambdaDriver implements ServerlessDriver {

await this.waitFunctionUpdates(serverlessFunction.id);

await this.lambdaClient.send(updateConfigurationCommand);
await (await this.getLambdaClient()).send(updateConfigurationCommand);
}

await this.waitFunctionUpdates(serverlessFunction.id);
Expand All @@ -280,7 +333,7 @@ export class LambdaDriver implements ServerlessDriver {

const command = new PublishVersionCommand(params);

const result = await this.lambdaClient.send(command);
const result = await (await this.getLambdaClient()).send(command);
const newVersion = result.Version;

if (!newVersion) {
Expand Down Expand Up @@ -331,7 +384,7 @@ export class LambdaDriver implements ServerlessDriver {
const command = new InvokeCommand(params);

try {
const result = await this.lambdaClient.send(command);
const result = await (await this.getLambdaClient()).send(command);

const parsedResult = result.Payload
? JSON.parse(result.Payload.transformToString())
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -29,7 +29,11 @@ export const serverlessModuleFactory = async (
const secretAccessKey = environmentService.get(
'SERVERLESS_LAMBDA_SECRET_ACCESS_KEY',
);
const role = environmentService.get('SERVERLESS_LAMBDA_ROLE');
const lambdaRole = environmentService.get('SERVERLESS_LAMBDA_ROLE');

const subhostingRole = environmentService.get(
'SERVERLESS_LAMBDA_SUBHOSTING_ROLE',
);

return {
type: ServerlessDriverType.Lambda,
Expand All @@ -43,8 +47,9 @@ export const serverlessModuleFactory = async (
: fromNodeProviderChain({
clientConfig: { region },
}),
region: region ?? '',
role: role ?? '',
region,
lambdaRole,
subhostingRole,
},
};
}
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -197,8 +197,6 @@ export class WorkflowStatusesUpdateJob {
serverlessFunction.latestVersion;

newStep.settings = newStepSettings;

this.logger.log(`New step computed for code step -> ${newStep}`);
}
newSteps.push(newStep);
}
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -271,6 +271,7 @@ yarn command:prod cron:calendar:ongoing-stale
['SERVERLESS_TYPE', 'local', "Serverless driver type: 'local' or 'lambda'"],
['SERVERLESS_LAMBDA_REGION', '', 'Lambda Region'],
['SERVERLESS_LAMBDA_ROLE', '', 'Lambda Role'],
['SERVERLESS_LAMBDA_SUBHOSTING_ROLE', '', 'Role to assume when hosting lambdas in dedicated AWS account'],
['SERVERLESS_LAMBDA_ACCESS_KEY_ID', '', 'Optional depending on the authentication method'],
['SERVERLESS_LAMBDA_SECRET_ACCESS_KEY', '', 'Optional depending on the authentication method'],
]}></ArticleTable>
Expand Down Expand Up @@ -304,7 +305,8 @@ This feature is WIP and is not yet useful for most users.
<ArticleTable options={[
['SERVERLESS_TYPE', 'local', "Functions can either be executed through Lambda or directly by the main node process"],
['SERVERLESS_LAMBDA_REGION', 'us-east-1', 'If you use the Lambda driver, region of the Lambda function'],
['SERVERLESS_LAMBDA_ROLE', 'arn:aws:iam::121334:role/lambda-execution-role', "If you use the Lambda drive, name of the IAM role with the right permissions"],
['SERVERLESS_LAMBDA_ROLE', 'arn:aws:iam::121334:role/lambda-execution-role', "If you use the Lambda driver, name of the IAM role with the right permissions"],
['SERVERLESS_LAMBDA_SUBHOSTING_ROLE', 'arn:aws:iam::121334:role/lambda-deployment-role', "If you host lambdas in a dedicated AWS account, name of the IAM role to assume in the dedicated account"],
]}></ArticleTable>


Expand Down
Loading

0 comments on commit f2da915

Please sign in to comment.