Migrate Python Connector Examples

README.md

@@ -22,6 +22,7 @@ The subdirectories contain code examples for connecting and using Aurora DSQL in
 | JavaScript  | [node-postgres (standalone)](javascript/node-postgres/) |
 | JavaScript  |         [Postgres.js](javascript/postgres-js/)          |
 |   Python    |                [Jupyter](python/jupyter)                |
+|   Python    |               [asyncpg](python/asyncpg/)                |
 |   Python    |               [psycopg](python/psycopg/)                |
 |   Python    |              [psycopg2](python/psycopg2/)               |
 |   Python    |             [SQLAlchemy](python/sqlalchemy)             |

python/asyncpg/README.md

@@ -0,0 +1,151 @@
+# Aurora DSQL with asyncpg
+
+This example demonstrates how to use the Aurora DSQL Python Connector with asyncpg to connect to Amazon Aurora DSQL clusters and perform basic database operations.
+
+Aurora DSQL is a distributed SQL database service that provides high availability and scalability for
+your PostgreSQL-compatible applications.
+Asyncpg is a popular PostgreSQL database library 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).
+
+
+## TLS connection configuration
+
+This example uses direct TLS connections where supported, and verifies the server certificate is trusted. Verified SSL
+connections should be used where possible to ensure data security during transmission.
+
+* Driver versions following the release of PostgreSQL 17 support direct TLS connections, bypassing the traditional
+  PostgreSQL connection preamble
+* Direct TLS connections provide improved connection performance and enhanced security
+* Not all PostgreSQL drivers support direct TLS connections yet, or only in recent versions following PostgreSQL 17
+* Ensure your installed driver version supports direct TLS negotiation, or use a version that is at least as recent as
+  the one used in this sample
+* If your driver doesn't support direct TLS connections, you may need to use the traditional preamble connection instead
+
+
+### Prerequisites
+
+* 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.10.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:
+
+```
+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
+```
+
+2. Install the required packages for running the examples:
+
+```bash
+pip install -r requirements.txt
+```
+
+### Run the code
+
+#### What the Examples Do
+
+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.
+
+The connection pool examples demonstrate:
+- Creating a connection pool for Aurora DSQL
+- Using async context managers for connection management
+- Performing database operations through the pool
+- Running multiple concurrent database operations
+- Using asyncio.gather() for parallel execution
+- Proper resource management with connection pools
+
+#### Environment Cluster Details
+
+Set environment variables for your cluster details:
+
+```bash
+# e.g. "admin"
+export CLUSTER_USER="<your user>"
+
+# 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
+# Run example directly
+python src/example.py
+
+# Run example using pytest
+pytest ./test/test_example.py
+
+# Run all using pytest
+pytest ./test
+```
+
+The examples contain 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)
+* [Amazon Aurora DSQL Python Connector Documentation](https://docs.aws.amazon.com/aurora-dsql/latest/userguide/SECTION_program-with-dsql-connector-for-python.html)
+* [Asyncpg Documentation](https://magicstack.github.io/asyncpg/current/)
+* [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.
+
+SPDX-License-Identifier: MIT-0

python/asyncpg/pyproject.toml

@@ -0,0 +1,4 @@
+[tool.pytest.ini_options]
+pythonpath = [
+  "src"
+]

python/asyncpg/requirements.txt

@@ -0,0 +1,5 @@
+aurora-dsql-python-connector>=0.2.0
+boto3>=1.35.74
+asyncpg>=0.30.0
+pytest>=8.0
+pytest-asyncio>=1.2.0

python/asyncpg/src/example.py

@@ -0,0 +1,101 @@
+"""
+Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+SPDX-License-Identifier: MIT-0
+"""
+
+import asyncio
+import os
+
+import aurora_dsql_asyncpg as dsql
+
+
+async def create_connection(cluster_user, cluster_endpoint, region):
+    ssl_cert_path = "./root.pem"
+    if not os.path.isfile(ssl_cert_path):
+        raise FileNotFoundError(f"SSL certificate file not found: {ssl_cert_path}")
+
+    conn_params = {
+        "database": "postgres",
+        "user": cluster_user,
+        "host": cluster_endpoint,
+        "port": 5432,
+        "region": region,
+        "ssl": "verify-full",
+        "sslrootcert": ssl_cert_path,
+    }
+
+    # Make a connection to the cluster
+    conn = await dsql.connect(**conn_params)
+
+    if cluster_user == "admin":
+        schema = "public"
+    else:
+        schema = "myschema"
+
+    try:
+        await conn.execute(f"SET search_path = {schema};")
+    except Exception as e:
+        await conn.close()
+        raise e
+
+    return conn
+
+
+async def exercise_connection(conn):
+    await conn.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))
+            """
+    )
+
+    # Insert some rows
+    await conn.execute(
+        "INSERT INTO owner(name, city, telephone) VALUES($1, $2, $3)",
+        "John Doe",
+        "Anytown",
+        "555-555-1999",
+    )
+
+    row = await conn.fetchrow("SELECT * FROM owner WHERE name=$1", "John Doe")
+
+    # Verify the result we got is what we inserted before
+    assert row[0] is not None
+    assert row[1] == "John Doe"
+    assert row[2] == "Anytown"
+    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
+    await conn.execute("DELETE FROM owner WHERE name = $1", "John Doe")
+
+
+async 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 = await create_connection(cluster_user, cluster_endpoint, region)
+        await exercise_connection(conn)
+    finally:
+        if conn is not None:
+            await conn.close()
+
+    print("Connection exercised successfully")
+
+
+if __name__ == "__main__":
+    asyncio.run(main())

python/asyncpg/src/example_with_connection_pool.py

@@ -0,0 +1,64 @@
+"""
+Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+SPDX-License-Identifier: MIT-0
+"""
+
+import asyncio
+import os
+
+import aurora_dsql_asyncpg as dsql
+
+
+async def connect_with_pool(cluster_user, cluster_endpoint, region):
+    ssl_cert_path = "./root.pem"
+    if not os.path.isfile(ssl_cert_path):
+        raise FileNotFoundError(f"SSL certificate file not found: {ssl_cert_path}")
+
+    pool_params = {
+        "database": "postgres",
+        "user": cluster_user,
+        "host": cluster_endpoint,
+        "port": 5432,
+        "region": region,
+        "ssl": "verify-full",
+        "sslrootcert": ssl_cert_path,
+        "min_size": 2,
+        "max_size": 5,
+    }
+
+    pool = await dsql.create_pool(**pool_params)
+    try:
+        async with pool.acquire() as conn:
+            result = await conn.fetchval("SELECT 1")
+            assert result == 1
+    finally:
+        await pool.close()
+
+
+async def main():
+    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"
+
+        ssl_cert_path = "./root.pem"
+        if not os.path.isfile(ssl_cert_path):
+            raise FileNotFoundError(f"SSL certificate file not found: {ssl_cert_path}")
+
+        await connect_with_pool(cluster_user, cluster_endpoint, region)
+
+    finally:
+        pass
+
+    print("Pool exercised successfully")
+
+
+if __name__ == "__main__":
+    asyncio.run(main())