Using AWS Lambda with Amazon Aurora DSQL

Sign in to the AWS Management Console and open the AWS Lambda console at https://console.aws.amazon.com/lambda/.

Choose Create function.

Provide a name, such as dsql-sample.

Don't edit the default settings to make sure that Lambda creates a new role with basic Lambda permissions.

Choose Create function.

In your Lambda function, choose Configuration > Permissions.

Choose the role name to open the execution role in the IAM console.

Choose Add Permissions > Create inline policy, and use the JSON editor.

In Action paste in the following action to authorize your IAM identity to connect using the admin database role.

"Action": ["dsql:DbConnectAdmin"],

In Resource, add your cluster’s Amazon Resource Name (ARN). You can also use a wildcard.

"Resource": ["*"]

Choose Next.

Enter a name for the policy, such as dsql-sample-dbconnect.

Choose Create policy.

Create a folder named myfunction.

In the folder, create a new file named package.json with the following content.

{
  "dependencies": {
    "@aws-sdk/core": "^3.587.0",
    "@aws-sdk/credential-providers": "^3.587.0",
    "@smithy/protocol-http": "^4.0.0",
    "@smithy/signature-v4": "^3.0.0",
    "pg": "^8.11.5"
  }
}

In the folder, create a file named index.mjs in the directory with the following content.

import { formatUrl } from "@aws-sdk/util-format-url";
import { HttpRequest } from "@smithy/protocol-http";
import { SignatureV4 } from "@smithy/signature-v4";
import { fromNodeProviderChain } from "@aws-sdk/credential-providers";
import { NODE_REGION_CONFIG_FILE_OPTIONS, NODE_REGION_CONFIG_OPTIONS } from "@smithy/config-resolver";
import { Hash } from "@smithy/hash-node";
import { loadConfig } from "@smithy/node-config-provider";
import pg from "pg";
const { Client } = pg;

export const getRuntimeConfig = (config) => {
  return {
    runtime: "node",
    sha256: config?.sha256 ?? Hash.bind(null, "sha256"),
    credentials: config?.credentials ?? fromNodeProviderChain(),
    region: config?.region ?? loadConfig(NODE_REGION_CONFIG_OPTIONS, NODE_REGION_CONFIG_FILE_OPTIONS),
    ...config,
  };
};

// Aurora DSQL requires IAM authentication
// This class generates auth tokens signed using AWS Signature Version 4
export class Signer {
  constructor(hostname) {
    const runtimeConfiguration = getRuntimeConfig({});

    this.credentials = runtimeConfiguration.credentials;
    this.hostname = hostname;
    this.region = runtimeConfiguration.region;

    this.sha256 = runtimeConfiguration.sha256;
    this.service = "dsql";
    this.protocol = "https:";
  }
  
  async getAuthToken() {
    const signer = new SignatureV4({
      service: this.service,
      region: this.region,
      credentials: this.credentials,
      sha256: this.sha256,
    });

    // To connect with a custom database role, set Action as "DbConnect"
    const request = new HttpRequest({
      method: "GET",
      protocol: this.protocol,
      hostname: this.hostname,
      query: {
        Action: "DbConnectAdmin",
      },
      headers: {
        host: this.hostname,
      },
    });

    const presigned = await signer.presign(request, {
      expiresIn: 3600,
    });

    // RDS requires the scheme to be removed
    // https://docs.aws.amazon.com/AmazonRDS/latest/UserGuide/UsingWithRDS.IAMDBAuth.Connecting.html
    return formatUrl(presigned).replace(`${this.protocol}//`, "");
  }
}

// To connect with a custom database role, set user as the database role name
async function dsql_sample(token, endpoint) {
  const client = new Client({
    user: "admin",
    database: "postgres",
    host: endpoint,
    password: token,
    ssl: { 
      rejectUnauthorized: false
    },
  });

  await client.connect();
  console.log("[dsql_sample] connected to Aurora DSQL!");

  try {
    console.log("[dsql_sample] attempting transaction.");
    await client.query("BEGIN; SELECT txid_current_if_assigned(); COMMIT;");
    return 200;
  } catch (err) {
    console.log("[dsql_sample] transaction attempt failed!");
    console.error(err);
    return 500;
  } finally {
    await client.end();
  }
}

// https://docs.aws.amazon.com/lambda/latest/dg/nodejs-handler.html
export const handler = async (event) => {
  const endpoint = event.endpoint;
  const s = new Signer(endpoint);
  const token = await s.getAuthToken();
  const responseCode = await dsql_sample(token, endpoint);
  
  const response = {
    statusCode: responseCode,
    endpoint: endpoint,
  };
  return response;
};

Use the following commands to create a package.

npm install
zip -r pkg.zip .

In your Lambda function’s Code tab, choose Upload from > .zip file

Upload the pkg.zip you created. For more information, see Deploy Node.js Lambda functions with .zip file archives.

In your Lambda function’s Test tab, paste in the following JSON payload, and modify it to use your cluster ID.

In your Lambda function’s Test tab, use the following Event JSON modified to specify your cluster’s endpoint.

{"endpoint": "replace_with_your_cluster_endpoint"}

Enter an Event name, such as dsql-sample-test. Choose Save.

Choose Test.

Choose Details to expand the execution response and log output.

If it succeeded, the Lambda function execution response should return a 200 status code:

{statusCode": 200, "endpoint": "your_cluster_endpoint"}

If the database returns an error or if the connection to the database fails, the Lambda function execution response returns a 500 status code.

{"statusCode": 500,"endpoint": "your_cluster_endpoint"}