Merge pull request #3117 from sundersc/dsql-lambda

feat(graphql-model-transformer): add dsql support for sql lambda

packages/amplify-graphql-model-transformer/rds-lambda/dsql-helpers.ts

@@ -0,0 +1,17 @@
+import { DsqlSigner } from "@aws-sdk/dsql-signer";
+
+export const generateDSQLAuthToken = async (endpoint: string): Promise<string> => {
+  const signer = new DsqlSigner({
+    hostname: endpoint,
+  });
+  try {
+    const token = await signer.getDbConnectAdminAuthToken();
+    return token;
+  } catch (error) {
+    throw error;
+  }
+}
+
+export const isDSQLHostname = (endpoint: string): boolean => {
+  return endpoint.includes(".dsql.") && endpoint.endsWith(".on.aws");
+}

packages/amplify-graphql-model-transformer/rds-lambda/handler.ts

@@ -2,6 +2,7 @@ import { SSMClient, GetParameterCommand, GetParameterCommandOutput } from '@aws-
 import { GetSecretValueCommand, SecretsManagerClient, GetSecretValueCommandOutput } from '@aws-sdk/client-secrets-manager';
 // @ts-ignore
 import { DBAdapter, DBConfig, getDBAdapter } from 'rds-query-processor';
+import { generateDSQLAuthToken, isDSQLHostname } from './dsql-helpers';
 
 let adapter: DBAdapter;
 let ssmClient: SSMClient;
@@ -52,7 +53,9 @@ const isRetryableError = (error: Error & {code?: string, errno?: string}): boole
   // https://dev.mysql.com/doc/mysql-errors/8.0/en/server-error-reference.html
   const mysqlRetryableError = error.errno === '1045';
 
-  return postgresRetryableError || mysqlRetryableError;
+  const dsqlRetryableError = error.code === '08006' && error.message?.includes('access denied');
+
+  return postgresRetryableError || mysqlRetryableError || dsqlRetryableError;
 }
 
 const createSSMClient = (): void => {
@@ -180,7 +183,7 @@ const retrieveSsmValueFromEnvPaths = async (path: string): Promise<string> => {
 };
 
 const getDBConfig = async (): DBConfig => {
-  const config: DBConfig = {};
+  let config: DBConfig = {};
 
   const sslCertificate = await getCustomSslCert();
   if (sslCertificate) {
@@ -195,7 +198,26 @@ const getDBConfig = async (): DBConfig => {
 
     const jsonConnectionString = process.env.connectionString;
     if (jsonConnectionString) {
-      config.connectionString = await retrieveSsmValueFromEnvPaths(jsonConnectionString);
+      const connectionString = await retrieveSsmValueFromEnvPaths(jsonConnectionString);
+
+      // If the host is a DSQL hostname, generate an auth token
+      const { hostname } = new URL(connectionString);
+      config.host = decodeURIComponent(hostname);
+      const defaultDSQLConfig = {
+        username: 'admin',
+        engine: 'postgres',
+        port: 5432,
+        database: 'postgres',
+      }
+      if (isDSQLHostname(config.host)) {
+        config = { ...defaultDSQLConfig, ...config };
+        config.password = await generateDSQLAuthToken(config.host);
+        return config;
+      }
+
+      // Fall back to the connection string if the host is not a DSQL hostname
+      delete config.host;
+      config.connectionString = connectionString;
       return config;
     }
 

packages/amplify-graphql-model-transformer/rds-lambda/package.json

@@ -7,6 +7,7 @@
   "dependencies": {
     "@aws-sdk/client-secrets-manager": "3.462.0",
     "@aws-sdk/client-ssm": "3.624.0",
+    "@aws-sdk/dsql-signer": "^3.726.1",
     "babel-jest": "^29.1.2",
     "bestzip": "^2.1.5",
     "jest": "^29.1.2",

packages/amplify-graphql-model-transformer/src/resolvers/rds/resolver.ts

@@ -469,7 +469,7 @@ export const createRdsLambdaRole = (
     new PolicyStatement({
       effect: Effect.ALLOW,
       resources: ['*'],
-      actions: ['ec2:CreateNetworkInterface', 'ec2:DescribeNetworkInterfaces', 'ec2:DeleteNetworkInterface'],
+      actions: ['ec2:CreateNetworkInterface', 'ec2:DescribeNetworkInterfaces', 'ec2:DeleteNetworkInterface', 'dsql:DbConnectAdmin'],
     }),
   );