sqlalchemy

Aurora DSQL with SQLAlchemy

Overview

This code example demonstrates how to use SQLAlchemy with Amazon Aurora DSQL. The example shows you how to connect to an Aurora DSQL cluster with SQLAlchemy using the Aurora DSQL Python Connector, define ORM models, and perform CRUD operations.

Aurora DSQL is a distributed SQL database service that provides high availability and scalability for your PostgreSQL-compatible applications. SQLAlchemy is a popular Python SQL toolkit and ORM that provides a full suite of well-known enterprise-level persistence patterns.

About the code example

This example uses the Aurora DSQL Python Connector which automatically handles IAM token generation for authentication.

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.
• This code is not tested in every AWS Region. For more information, see AWS Regional Services.

Run the example

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 guide.
• Python 3.10 or later.
• You must have an Aurora DSQL cluster. For information about creating an Aurora DSQL cluster, see the Getting started with Aurora DSQL 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 guide.

Set up environment for examples

1. Create and activate a Python virtual environment:

python3 -m venv .venv
source .venv/bin/activate

2. Install the required packages:

pip install -r requirements.txt

Run the code

The example demonstrates the following operations:

• Connecting to Aurora DSQL using SQLAlchemy with IAM authentication
• Creating tables using SQLAlchemy's Base.metadata.create_all()
• Defining ORM models with UUID primary keys and application-level relationships
• Inserting and querying data with relationship loading (joinedload)

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:

# e.g. "admin"
export CLUSTER_USER="<your user>"

# e.g. "foo0bar1baz2quux3quuux4.dsql.us-east-1.on.aws"
export CLUSTER_ENDPOINT="<your endpoint>"

Run the example:

python -m src.example

Run the tests:

pytest

Usage notes

Connecting to DSQL

Aurora DSQL is PostgreSQL-compatible, so use the postgresql+psycopg dialect. The Aurora DSQL Python Connector handles IAM token generation automatically. Pass a connection creator to SQLAlchemy's create_engine.

Set isolation_level="AUTOCOMMIT" and autocommit=True on the connection to avoid SQLAlchemy's psycopg dialect issuing SAVEPOINT during initialization:

import aurora_dsql_psycopg as dsql
from sqlalchemy import create_engine

engine = create_engine(
    "postgresql+psycopg://",
    creator=lambda: dsql.connect(host=endpoint, user=user, autocommit=True),
    isolation_level="AUTOCOMMIT",
    pool_recycle=3300,  # Recycle before DSQL's 1-hour connection limit
)

For non-admin users, set the search path using a connection event:

from sqlalchemy import event

@event.listens_for(engine, "connect")
def set_search_path(dbapi_connection, connection_record):
    cursor = dbapi_connection.cursor()
    cursor.execute("SET search_path TO myschema")
    cursor.close()

Primary keys

The SERIAL pseudo-type is not available in Aurora DSQL. Aurora DSQL supports sequences and identity columns (with CACHE specified), but UUIDs with gen_random_uuid() are the recommended default for primary keys:

from uuid import UUID
from sqlalchemy import Uuid, text
from sqlalchemy.orm import Mapped, mapped_column

id: Mapped[UUID] = mapped_column(
    Uuid, primary_key=True, server_default=text("gen_random_uuid()")
)

Relationships with application-layer referential integrity

This sample uses application-layer referential integrity. Define relationships using relationship() with explicit primaryjoin and foreign() annotations instead of ForeignKey():

from sqlalchemy.orm import relationship, foreign

class Owner(Base):
    __tablename__ = "owner"
    id: Mapped[UUID] = mapped_column(Uuid, primary_key=True, server_default=text("gen_random_uuid()"))

class Pet(Base):
    __tablename__ = "pet"
    owner_id: Mapped[Optional[UUID]] = mapped_column(Uuid, nullable=True)

# Define after both classes exist
Owner.pets = relationship(Pet, primaryjoin=Owner.id == foreign(Pet.owner_id), back_populates="owner")
Pet.owner = relationship(Owner, primaryjoin=foreign(Pet.owner_id) == Owner.id, back_populates="pets")

The foreign() annotation tells SQLAlchemy which column is the "foreign" side of the join, replacing the role that ForeignKey() normally plays.

For the full list of Aurora DSQL SQL compatibility details, see the PostgreSQL compatibility reference.

Additional resources

• Amazon Aurora DSQL Documentation
• Sequences and identity columns in Aurora DSQL
• Migrating from PostgreSQL to Aurora DSQL
• Aurora DSQL Python Connector
• SQLAlchemy Documentation

Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.

SPDX-License-Identifier: MIT-0