Update Psycopg2 example to include non-admin and connection refresh logic

.github/workflows/python-psycopg2-integ-tests.yml

@@ -53,19 +53,19 @@ jobs:
         role-to-assume: ${{ secrets.PYTHON_IAM_ROLE }}
         aws-region: us-east-1
 
-    - name: Configure and run integration for psycopg2
+    - name: Configure and run integration for psycopg2 - admin
       working-directory: ./python/psycopg2
       env:
+        CLUSTER_USER: "admin"
         CLUSTER_ENDPOINT: ${{ secrets.PYTHON_PSYCOPG2_CLUSTER_ENDPOINT }}
         REGION: ${{ secrets.PYTHON_PSYCOPG2_CLUSTER_REGION }}
       run: |
-        python3 -m venv psycopg2-integ
-        source psycopg2-integ/bin/activate
+        python3 -m venv .venv
+        source .venv/bin/activate
         pip install --upgrade pip
         pip install --force-reinstall -r requirements.txt
         python3 -c "import boto3; print(boto3.__version__)"
-        pip install pytest pytest-cov
         pip list
         echo "$GITHUB_WORKSPACE" >> $GITHUB_PATH
-        pytest -v test/
-
+        wget https://www.amazontrust.com/repository/AmazonRootCA1.pem -O root.pem
+        python src/example.py

python/psycopg2/README.md

@@ -1,32 +1,122 @@
-# Amazon Distributed SQL Psycopg2 code examples
+# Psycopg2 with Aurora DSQL
 
 ## Overview
 
-The code examples in this topic show you how to use psycopg2 with Amazon Distributed SQL. 
+This code example demonstrates how to use Psycopg (version 2) with Amazon Aurora DSQL. The example shows you how to
+connect to an Aurora DSQL cluster and perform basic database operations.
 
-## Run the examples
+Aurora DSQL is a distributed SQL database service that provides high availability and scalability for
+your PostgreSQL-compatible applications. Psycopg2 is a popular PostgreSQL adapter for Python that allows
+you to interact with PostgreSQL databases using Python code.
+
+## About the code example
+
+The example demonstrates a flexible connection approach that works for both admin and non-admin users:
+
+* When connecting as an **admin user**, the example uses the `public` schema and generates an admin authentication
+  token.
+* When connecting as a **non-admin user**, the example uses a custom `myschema` schema and generates a standard
+  authentication token.
+
+The code automatically detects the user type and adjusts its behavior accordingly.
+
+## ⚠️ Important
+
+* Running this code might result in charges to your AWS account.
+* We recommend that you grant your code least privilege. At most, grant only the
+  minimum permissions required to perform the task. For more information, see
+  [Grant least privilege](https://docs.aws.amazon.com/IAM/latest/UserGuide/best-practices.html#grant-least-privilege).
+* This code is not tested in every AWS Region. For more information, see
+  [AWS Regional Services](https://aws.amazon.com/about-aws/global-infrastructure/regional-product-services).
+
+## Run the example
 
 ### Prerequisites
 
-* python version >=3.8.0 is needed
+* You must have an AWS account, and have your default credentials and AWS Region
+  configured as described in the
+  [Globally configuring AWS SDKs and tools](https://docs.aws.amazon.com/credref/latest/refdocs/creds-config-files.html)
+  guide.
+* [Python 3.8.0](https://www.python.org/) or later.
+* You must have an Aurora DSQL cluster. For information about creating an Aurora DSQL cluster, see the
+  [Getting started with Aurora DSQL](https://docs.aws.amazon.com/aurora-dsql/latest/userguide/getting-started.html)
+  guide.
+* If connecting as a non-admin user, ensure the user is linked to an IAM role and is granted access to the `myschema`
+  schema. See the
+  [Using database roles with IAM roles](https://docs.aws.amazon.com/aurora-dsql/latest/userguide/using-database-and-iam-roles.html)
+  guide.
+
+### Download the Amazon root certificate from the official trust store
+
+Download the Amazon root certificate from the official trust store. This example shows one of the available certs that
+can be used by the client. Other certs such as AmazonRootCA2.pem, AmazonRootCA3.pem, etc. can also be used.
+
+```
+wget https://www.amazontrust.com/repository/AmazonRootCA1.pem -O root.pem
+```
+
+### Set up environment for examples
+
+1. Create and activate a Python virtual environment:
+
+```bash
+python3 -m venv .venv
+source .venv/bin/activate  # Linux, macOS
+# or
+.venv\Scripts\activate     # Windows
+```
 
-### Setup test running environment 
+2. Install the required packages for running the examples:
 
-```sh
-source setup.sh
+```bash
+pip install "boto3>=1.35.74"
+pip install --prefer-binary "psycopg2-binary>=2.9"
 ```
 
-### Run the example tests
+### Run the code
+
+The example demonstrates the following operations:
+
+- Opening a connection to an Aurora DSQL cluster
+- Creating a table
+- Inserting and querying data
+
+The example is designed to work with both admin and non-admin users:
+
+- When run as an admin user, it uses the `public` schema
+- When run as a non-admin user, it uses the `myschema` schema
+
+**Note:** running the example will use actual resources in your AWS account and may incur charges.
+
+Set environment variables for your cluster details:
+
+```bash
+# e.g. "admin"
+export CLUSTER_USER="<your user>"
 
-```sh
-# Use the account credentials dedicated for python
-export CLUSTER_ENDPOINT="<your cluster endpoint>"
-export REGION="<your cluster region>"
-pytest test/test_example.py
+# e.g. "foo0bar1baz2quux3quuux4.dsql.us-east-1.on.aws"
+export CLUSTER_ENDPOINT="<your endpoint>"
+
+# e.g. "us-east-1"
+export REGION="<your region>"
+```
+
+Run the example:
+
+```bash
+python src/example.py
 ```
 
+The example contains comments explaining the code and the operations being performed.
+
+## Additional resources
+
+* [Amazon Aurora DSQL Documentation](https://docs.aws.amazon.com/aurora-dsql/latest/userguide/what-is-aurora-dsql.html)
+* [Psycopg2 Documentation](https://www.psycopg.org/docs/)
+* [AWS SDK for Python (Boto3) Documentation](https://boto3.amazonaws.com/v1/documentation/api/latest/index.html)
+
 ---
 
-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

python/psycopg2/requirements.txt

@@ -1,4 +1,3 @@
+boto3>=1.35.74
 psycopg2-binary>=2.9
 pytest>=8
-botocore>=1.35.74
-boto3>=1.35.74

python/psycopg2/src/example.py

@@ -1,62 +1,99 @@
-import psycopg2
 import boto3
-import os, sys
+import psycopg2
+import os
+import sys
+
 
-def main(cluster_endpoint, region):
-    # Generate a password token
+def create_connection(cluster_user, cluster_endpoint, region):
+    # Generate a fresh password token for each connection, to ensure the token is not expired
+    # when the connection is established
     client = boto3.client("dsql", region_name=region)
-    # The token expiration time is optional, and the default value 900 seconds
-    password_token = client.generate_db_connect_admin_auth_token(cluster_endpoint, region)
 
-    # connection parameters
-    dbname = "dbname=postgres"
-    user = "user=admin"
-    host = f'host={cluster_endpoint}'
-    sslmode = "sslmode=require"
-    password = f'password={password_token}'
+    if cluster_user == "admin":
+        password_token = client.generate_db_connect_admin_auth_token(cluster_endpoint, region)
+    else:
+        password_token = client.generate_db_connect_auth_token(cluster_endpoint, region)
+
+    conn_params = {
+        "dbname": "postgres",
+        "user": cluster_user,
+        "host": cluster_endpoint,
+        "port": "5432",
+        "sslmode": "verify-full",
+        "sslrootcert": "./root.pem",
+        "password": password_token
+    }
 
     # Make a connection to the cluster
-    conn = psycopg2.connect('%s %s %s %s %s' % (dbname, user, host, sslmode, password))
+    conn = psycopg2.connect(**conn_params)
+
+    if cluster_user == "admin":
+        schema = "public"
+    else:
+        schema = "myschema"
 
+    try:
+        with conn.cursor() as cur:
+            cur.execute(f"SET search_path = {schema};")
+            conn.commit()
+    except Exception as e:
+        conn.close()
+        raise e
+
+    return conn
+
+
+def exercise_connection(conn):
     conn.set_session(autocommit=True)
 
     cur = conn.cursor()
-
-    cur.execute("DROP TABLE IF EXISTS owner")
-
-    cur.execute(b"""
+
+    cur.execute("""
         CREATE TABLE IF NOT EXISTS owner(
             id uuid NOT NULL DEFAULT gen_random_uuid(),
             name varchar(30) NOT NULL,
             city varchar(80) NOT NULL,
             telephone varchar(20) DEFAULT NULL,
-            PRIMARY KEY (id))"""
-        )
+            PRIMARY KEY (id))
+            """)
 
     # Insert some rows
-    cur.execute("INSERT INTO owner(name, city, telephone) VALUES('John Doe', 'Anytown', '555-555-0150')")
+    cur.execute("INSERT INTO owner(name, city, telephone) VALUES('John Doe', 'Anytown', '555-555-1999')")
 
-    # Read back what we have inserted
     cur.execute("SELECT * FROM owner WHERE name='John Doe'")
     row = cur.fetchone()
-    
-    # Verify that the result we got is what we inserted before
+
+    # Verify the result we got is what we inserted before
     assert row[0] != None
     assert row[1] == "John Doe"
     assert row[2] == "Anytown"
-    assert row[3] == "555-555-0150"
-
-    # Insert some rows
-    # Placing this cleanup the table after the example. If we run the example
-    # again we do not have to worry about data inserted by previous runs
+    assert row[3] == "555-555-1999"
+
+    # Clean up the table after the example. If we run the example again
+    # we do not have to worry about data inserted by previous runs
     cur.execute("DELETE FROM owner where name = 'John Doe'")
 
+
+def main():
+    conn = None
+    try:
+        cluster_user = os.environ.get("CLUSTER_USER", None)
+        assert cluster_user is not None, "CLUSTER_USER environment variable is not set"
+
+        cluster_endpoint = os.environ.get("CLUSTER_ENDPOINT", None)
+        assert cluster_endpoint is not None, "CLUSTER_ENDPOINT environment variable is not set"
+
+        region = os.environ.get("REGION", None)
+        assert region is not None, "REGION environment variable is not set"
+
+        conn = create_connection(cluster_user, cluster_endpoint, region)
+        exercise_connection(conn)
+    finally:
+        if conn is not None:
+            conn.close()
+
+    print("Connection exercised successfully")
+
+
 if __name__ == "__main__":
-    cluster_endpoint = os.environ.get("CLUSTER_ENDPOINT", None)
-    region = os.environ.get("REGION", None)
-    if cluster_endpoint is None:
-        sys.exit("CLUSTER_ENDPOINT environment variable is not set")
-    if region is None:
-        sys.exit("REGION environment variable is not set")
-    main(cluster_endpoint, region)
-
+    main()

python/psycopg2/test/test_example.py

@@ -1,16 +1,11 @@
 from example import main
+
 import pytest
-import os, sys
+
 
 # Smoke tests that our example works fine
 def test_example():
     try:
-        cluster_endpoint = os.environ.get("CLUSTER_ENDPOINT", None)
-        region = os.environ.get("REGION", None)
-        if cluster_endpoint is None:
-            sys.exit("CLUSTER_ENDPOINT environment variable is not set")
-        if region is None:
-            sys.exit("REGION environment variable is not set")
-        main(cluster_endpoint, region)
+        main()
     except Exception as e:
-        pytest.fail("Unexpected exception: " + str(e))
+        pytest.fail(f"Unexpected exception: {e}")