diff --git a/.github/workflows/python-psycopg2-integ-tests.yml b/.github/workflows/python-psycopg2-integ-tests.yml index b3bfc196..8fa9d167 100644 --- a/.github/workflows/python-psycopg2-integ-tests.yml +++ b/.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 diff --git a/python/psycopg2/README.md b/python/psycopg2/README.md index bb821a34..5d987334 100644 --- a/python/psycopg2/README.md +++ b/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="" -```sh -# Use the account credentials dedicated for python -export CLUSTER_ENDPOINT="" -export REGION="" -pytest test/test_example.py +# e.g. "foo0bar1baz2quux3quuux4.dsql.us-east-1.on.aws" +export CLUSTER_ENDPOINT="" + +# e.g. "us-east-1" +export 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 diff --git a/python/psycopg2/requirements.txt b/python/psycopg2/requirements.txt index 4b483932..a26cb7de 100644 --- a/python/psycopg2/requirements.txt +++ b/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 diff --git a/python/psycopg2/setup.sh b/python/psycopg2/setup.sh deleted file mode 100644 index 2a5ae65f..00000000 --- a/python/psycopg2/setup.sh +++ /dev/null @@ -1,13 +0,0 @@ -#!/bin/sh - -# NOTE: -# This is only there to help any one from internal dev group to configure their -# environment with a single command. We are not going to write the readme -# using this script for external users -if [ ! -z "$VIRTUAL_ENV" ]; then - echo "Deactivating existing virtual enviornment" - deactivate -fi -python3 -m venv box -source box/bin/activate -pip install --force-reinstall -r requirements.txt --no-cache-dir diff --git a/python/psycopg2/src/example.py b/python/psycopg2/src/example.py index f15d2558..2f455062 100644 --- a/python/psycopg2/src/example.py +++ b/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) - \ No newline at end of file + main() diff --git a/python/psycopg2/test/test_example.py b/python/psycopg2/test/test_example.py index 0cc78016..c1ff694d 100644 --- a/python/psycopg2/test/test_example.py +++ b/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}")