use Aurora DSQL Connectors for Python

Author: jojo786

apigw-lambda-dsql/src/app.py

@@ -1,36 +1,22 @@
-#Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+    #Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
 #SPDX-License-Identifier: MIT-0
 
 import json
-import boto3
-import psycopg2
-import psycopg2.extensions
+import aurora_dsql_psycopg2 as dsql
 import os
 
 cluster_endpoint = os.environ['cluster_endpoint']
 region = os.environ['AWS_REGION']
 
-client = boto3.client("dsql", region_name=region)
-
 def lambda_handler(event, context):
-    # Generate a fresh password token for each connection, to ensure the token is not expired when the connection is established
-    password_token = client.generate_db_connect_admin_auth_token(cluster_endpoint, region)
-
-    conn_params = {
-        "dbname": "postgres",
-        "user": "admin",
-        "host": cluster_endpoint,
-        "port": "5432",
-        "sslmode": "require",
-        "password": password_token
+    config = {
+        'host': cluster_endpoint,
+        'region': region,
+        'user': "admin",
     }
 
-     # Use the more efficient connection method if it's supported.
-    if psycopg2.extensions.libpq_version() >= 170000:
-        conn_params["sslnegotiation"] = "direct"
-
-     # Make a connection to the cluster
-    conn = psycopg2.connect(**conn_params)
+    # Make a connection to the cluster
+    conn = dsql.connect(**config)
 
     try:
         with conn.cursor() as cur:

apigw-lambda-dsql/src/requirements.txt

@@ -1,3 +1,3 @@
 psycopg2-binary>=2.9
 botocore>=1.35.74
-boto3>=1.35.74
+aurora-dsql-python-connector

apigw-lambda-dsql/template.yaml

@@ -27,7 +27,7 @@ Resources:
         - Statement:
             - Effect: Allow
               Action:
-                - dsql:DbConnectAdmin
+                - dsql:DbConnectAdmin #IAM action dsql:DbConnectAdmin required to connect to the cluster.
               Resource:
                 - !Sub arn:${AWS::Partition}:dsql:${AWS::Region}:${AWS::AccountId}:cluster/${DSQL}
       Environment: # Function environment variables
@@ -42,7 +42,7 @@ Resources:
             Method: get
 
   DSQL:
-    Type: AWS::DSQL::Cluster
+    Type: AWS::DSQL::Cluster #Creates a single-Region DSQL cluster
     Properties:
       DeletionProtectionEnabled: false
       Tags: