import * as cdk from 'aws-cdk-lib';
import * as lambda from 'aws-cdk-lib/aws-lambda';
import * as apigateway from 'aws-cdk-lib/aws-apigateway';
import * as iam from 'aws-cdk-lib/aws-iam';
import { Construct } from 'constructs';
import * as path from 'path';
export interface AthenaLambdaStackProps extends cdk.StackProps {
outputS3Path?: string;
athenaWorkgroup?: string;
}
export class AthenaLambdaStack extends cdk.Stack {
constructor(scope: Construct, id: string, props?: AthenaLambdaStackProps) {
super(scope, id, props);
const outputS3Path = props?.outputS3Path ||
this.node.tryGetContext('outputS3Path') ||
's3://your-bucket/athena-results/';
const athenaWorkgroup = props?.athenaWorkgroup ||
this.node.tryGetContext('athenaWorkgroup') ||
'primary';
// Lambda 函数
const athenaFunction = new lambda.Function(this, 'AthenaFunction', {
runtime: lambda.Runtime.NODEJS_20_X,
handler: 'lambda.handler',
code: lambda.Code.fromAsset(path.join(__dirname, '../build')),
timeout: cdk.Duration.seconds(300),
memorySize: 512,
architecture: lambda.Architecture.ARM_64,
environment: {
OUTPUT_S3_PATH: outputS3Path,
ATHENA_WORKGROUP: athenaWorkgroup,
AWS_REGION: this.region,
},
});
// IAM 权限
athenaFunction.addToRolePolicy(new iam.PolicyStatement({
effect: iam.Effect.ALLOW,
actions: [
'athena:StartQueryExecution',
'athena:GetQueryExecution',
'athena:GetQueryResults',
'athena:ListNamedQueries',
'athena:BatchGetNamedQuery',
'athena:GetNamedQuery',
'athena:GetWorkGroup',
],
resources: ['*'],
}));
// 从 outputS3Path 提取 bucket 名称
const bucketName = outputS3Path.replace('s3://', '').split('/')[0];
athenaFunction.addToRolePolicy(new iam.PolicyStatement({
effect: iam.Effect.ALLOW,
actions: [
's3:GetBucketLocation',
's3:GetObject',
's3:ListBucket',
's3:PutObject',
],
resources: [
`arn:aws:s3:::${bucketName}`,
`arn:aws:s3:::${bucketName}/*`,
],
}));
athenaFunction.addToRolePolicy(new iam.PolicyStatement({
effect: iam.Effect.ALLOW,
actions: [
'glue:GetDatabase',
'glue:GetTable',
'glue:GetPartitions',
],
resources: ['*'],
}));
// API Gateway
const api = new apigateway.RestApi(this, 'McpApi', {
restApiName: 'Athena MCP Service',
description: 'MCP server for AWS Athena queries',
defaultCorsPreflightOptions: {
allowOrigins: apigateway.Cors.ALL_ORIGINS,
allowMethods: apigateway.Cors.ALL_METHODS,
allowHeaders: [
'Content-Type',
'X-Amz-Date',
'Authorization',
'X-Api-Key',
'X-Amz-Security-Token',
],
},
});
const mcpResource = api.root.addResource('mcp');
const integration = new apigateway.LambdaIntegration(athenaFunction);
mcpResource.addMethod('POST', integration);
// 输出
new cdk.CfnOutput(this, 'ApiEndpoint', {
value: api.url + 'mcp',
description: 'API Gateway endpoint URL',
});
new cdk.CfnOutput(this, 'FunctionArn', {
value: athenaFunction.functionArn,
description: 'Lambda Function ARN',
});
}
}
