Update aurora-dsql-connectors to v0.2.0

CHANGELOG.md

@@ -7,7 +7,11 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 Try to keep listed changes to a concise bulleted list of simple explanations of changes. Aim for the amount of information needed so that readers can understand where they would look in the codebase to investigate the changes' implementation, or where they would look in the documentation to understand how to make use of the change in practice - better yet, link directly to the docs and provide detailed information there. Only elaborate if doing so is required to avoid breaking changes or experimental features from ruining someone's day.
 
 ## [Unreleased]
+### Added
+- Add Amazon Aurora DSQL as a supported storage backend.
+
 ### Changed
+- Update Aurora DSQL Go connector to v0.2.0. Pool settings are now passed via `pgxpool.Config` instead of `dsql.Config`.
 - HTTP gateway's internal gRPC client now uses dynamic TLS credentials that automatically update on certificate rotation via certwatcher, preventing connection failures when certificates are rotated (e.g., by cert-manager). [#2951](https://github.com/openfga/openfga/pull/2951)
 
 ### Fixed

assets/assets.go

@@ -7,6 +7,7 @@ const (
 	MySQLMigrationDir    = "migrations/mysql"
 	PostgresMigrationDir = "migrations/postgres"
 	SqliteMigrationDir   = "migrations/sqlite"
+	DSQLMigrationDir     = "migrations/dsql"
 )
 
 // EmbedMigrations within the openfga binary.

assets/migrations/dsql/001_initialize_schema.sql

@@ -0,0 +1,60 @@
+-- +goose Up
+-- +goose NO TRANSACTION
+-- DSQL schema: uses full indexes (no partial), ASYNC index creation, C collation by default
+
+CREATE TABLE tuple (
+    store TEXT NOT NULL,
+    object_type TEXT NOT NULL,
+    object_id TEXT NOT NULL,
+    relation TEXT NOT NULL,
+    _user TEXT NOT NULL,
+    user_type TEXT NOT NULL,
+    ulid TEXT NOT NULL,
+    inserted_at TIMESTAMPTZ NOT NULL,
+    PRIMARY KEY (store, object_type, object_id, relation, _user)
+);
+
+CREATE INDEX ASYNC idx_tuple_user ON tuple (store, object_type, object_id, relation, _user, user_type);
+CREATE UNIQUE INDEX ASYNC idx_tuple_ulid ON tuple (ulid);
+
+CREATE TABLE authorization_model (
+    store TEXT NOT NULL,
+    authorization_model_id TEXT NOT NULL,
+    type TEXT NOT NULL,
+    type_definition BYTEA,
+    PRIMARY KEY (store, authorization_model_id, type)
+);
+
+CREATE TABLE store (
+    id TEXT PRIMARY KEY,
+    name TEXT NOT NULL,
+    created_at TIMESTAMPTZ NOT NULL,
+    updated_at TIMESTAMPTZ,
+    deleted_at TIMESTAMPTZ
+);
+
+CREATE TABLE assertion (
+    store TEXT NOT NULL,
+    authorization_model_id TEXT NOT NULL,
+    assertions BYTEA,
+    PRIMARY KEY (store, authorization_model_id)
+);
+
+CREATE TABLE changelog (
+    store TEXT NOT NULL,
+    object_type TEXT NOT NULL,
+    object_id TEXT NOT NULL,
+    relation TEXT NOT NULL,
+    _user TEXT NOT NULL,
+    operation INTEGER NOT NULL,
+    ulid TEXT NOT NULL,
+    inserted_at TIMESTAMPTZ NOT NULL,
+    PRIMARY KEY (store, ulid, object_type)
+);
+
+-- +goose Down
+DROP TABLE IF EXISTS tuple;
+DROP TABLE IF EXISTS authorization_model;
+DROP TABLE IF EXISTS store;
+DROP TABLE IF EXISTS assertion;
+DROP TABLE IF EXISTS changelog;

assets/migrations/dsql/002_add_authorization_model_version.sql

@@ -0,0 +1,9 @@
+-- +goose Up
+-- +goose NO TRANSACTION
+-- DSQL: ADD COLUMN without DEFAULT, then UPDATE existing rows
+ALTER TABLE authorization_model ADD COLUMN schema_version TEXT;
+UPDATE authorization_model SET schema_version = '1.0' WHERE schema_version IS NULL;
+
+-- +goose Down
+-- +goose NO TRANSACTION
+ALTER TABLE authorization_model DROP COLUMN schema_version;

assets/migrations/dsql/003_add_reverse_lookup_index.sql

@@ -0,0 +1,7 @@
+-- +goose Up
+-- +goose NO TRANSACTION
+CREATE INDEX ASYNC idx_reverse_lookup_user ON tuple (store, object_type, relation, _user);
+
+-- +goose Down
+-- +goose NO TRANSACTION
+DROP INDEX IF EXISTS idx_reverse_lookup_user;

assets/migrations/dsql/004_add_authorization_model_serialized_protobuf.sql

@@ -0,0 +1,7 @@
+-- +goose Up
+-- +goose NO TRANSACTION
+ALTER TABLE authorization_model ADD COLUMN serialized_protobuf BYTEA;
+
+-- +goose Down
+-- +goose NO TRANSACTION
+ALTER TABLE authorization_model DROP COLUMN serialized_protobuf;

assets/migrations/dsql/005_add_conditions_to_tuples.sql

@@ -0,0 +1,14 @@
+-- +goose Up
+-- +goose NO TRANSACTION
+-- DSQL: separate ALTER statements (one DDL per transaction)
+ALTER TABLE tuple ADD COLUMN condition_name TEXT;
+ALTER TABLE tuple ADD COLUMN condition_context BYTEA;
+ALTER TABLE changelog ADD COLUMN condition_name TEXT;
+ALTER TABLE changelog ADD COLUMN condition_context BYTEA;
+
+-- +goose Down
+-- +goose NO TRANSACTION
+ALTER TABLE tuple DROP COLUMN condition_name;
+ALTER TABLE tuple DROP COLUMN condition_context;
+ALTER TABLE changelog DROP COLUMN condition_name;
+ALTER TABLE changelog DROP COLUMN condition_context;

assets/migrations/dsql/006_add_collate_index.sql

@@ -0,0 +1,17 @@
+-- +goose Up
+-- +goose NO TRANSACTION
+-- DSQL: uses C collation by default, ASYNC instead of CONCURRENTLY
+CREATE INDEX ASYNC idx_user_lookup ON tuple (
+    store,
+    _user,
+    relation,
+    object_type,
+    object_id
+);
+
+DROP INDEX IF EXISTS idx_reverse_lookup_user;
+
+-- +goose Down
+-- +goose NO TRANSACTION
+DROP INDEX IF EXISTS idx_user_lookup;
+CREATE INDEX ASYNC idx_reverse_lookup_user ON tuple (store, object_type, relation, _user);

assets/migrations/dsql/README.md

@@ -0,0 +1,39 @@
+# DSQL Migrations for OpenFGA
+
+These migrations are adapted from the PostgreSQL migrations for Aurora DSQL compatibility.
+
+## Key Differences from PostgreSQL Migrations
+
+1. **ASYNC Indexes**: All `CREATE INDEX` statements use `CREATE INDEX ASYNC` (DSQL requirement)
+2. **NO TRANSACTION**: All migrations use `-- +goose NO TRANSACTION` as DSQL runs one DDL statement per transaction
+3. **BIGINT IDENTITY**: The goose version table uses `BIGINT GENERATED BY DEFAULT AS IDENTITY`, as required by DSQL
+4. **No Partial Indexes**: DSQL doesn't support partial indexes, so full indexes are used
+
+## Migration Files
+
+| File | Description |
+|------|-------------|
+| 001_initialize_schema.sql | Creates core tables (tuple, authorization_model, store, assertion, changelog) |
+| 002_add_authorization_model_version.sql | Adds schema_version column |
+| 003_add_reverse_lookup_index.sql | Adds reverse lookup index |
+| 004_add_authorization_model_serialized_protobuf.sql | Adds serialized_protobuf column |
+| 005_add_conditions_to_tuples.sql | Adds condition columns to tuple and changelog |
+| 006_add_collate_index.sql | Adds user lookup index with C collation |
+
+## Future Consideration: Splitting Migrations
+
+Per DSQL best practices, each migration should ideally contain only one DDL statement. Currently, migrations 002 and 005 contain multiple statements:
+
+- **002**: ALTER TABLE + UPDATE (mixed DDL and DML)
+- **005**: Four ALTER TABLE statements
+
+While these work correctly with `-- +goose NO TRANSACTION` (goose executes each statement separately), consider splitting them into individual migration files if you encounter OCC errors during migration. For example:
+
+```
+005a_add_condition_name_to_tuple.sql
+005b_add_condition_context_to_tuple.sql
+005c_add_condition_name_to_changelog.sql
+005d_add_condition_context_to_changelog.sql
+```
+
+See [DSQL Development Guide](https://docs.aws.amazon.com/aurora-dsql/latest/userguide/) for more information on DDL best practices.

cmd/run/run.go

@@ -503,6 +503,13 @@ func (s *ServerContext) datastoreConfig(config *serverconfig.Config) (storage.Op
 		if err != nil {
 			return nil, nil, fmt.Errorf("initialize sqlite datastore: %w", err)
 		}
+	case "dsql":
+		// Aurora DSQL uses PostgreSQL wire protocol with IAM authentication
+		// The postgres.New function handles dsql:// URIs automatically
+		datastore, err = postgres.New(config.Datastore.URI, dsCfg)
+		if err != nil {
+			return nil, nil, fmt.Errorf("initialize dsql datastore: %w", err)
+		}
 	default:
 		return nil, nil, fmt.Errorf("storage engine '%s' is unsupported", config.Datastore.Engine)
 	}

go.mod

@@ -9,6 +9,7 @@ require (
 	github.com/Masterminds/squirrel v1.5.4
 	github.com/MicahParks/keyfunc/v2 v2.1.0
 	github.com/Yiling-J/theine-go v0.6.2
+	github.com/awslabs/aurora-dsql-connectors/go/pgx v0.2.0
 	github.com/cenkalti/backoff/v4 v4.3.0
 	github.com/cespare/xxhash/v2 v2.3.0
 	github.com/containerd/errdefs v1.0.0
@@ -68,6 +69,21 @@ require (
 	github.com/Azure/go-ansiterm v0.0.0-20230124172434-306776ec8161 // indirect
 	github.com/Microsoft/go-winio v0.6.2 // indirect
 	github.com/antlr4-go/antlr/v4 v4.13.1 // indirect
+	github.com/aws/aws-sdk-go-v2 v1.41.3 // indirect
+	github.com/aws/aws-sdk-go-v2/config v1.32.11 // indirect
+	github.com/aws/aws-sdk-go-v2/credentials v1.19.11 // indirect
+	github.com/aws/aws-sdk-go-v2/feature/dsql/auth v1.1.19 // indirect
+	github.com/aws/aws-sdk-go-v2/feature/ec2/imds v1.18.19 // indirect
+	github.com/aws/aws-sdk-go-v2/internal/configsources v1.4.19 // indirect
+	github.com/aws/aws-sdk-go-v2/internal/endpoints/v2 v2.7.19 // indirect
+	github.com/aws/aws-sdk-go-v2/internal/ini v1.8.5 // indirect
+	github.com/aws/aws-sdk-go-v2/service/internal/accept-encoding v1.13.6 // indirect
+	github.com/aws/aws-sdk-go-v2/service/internal/presigned-url v1.13.19 // indirect
+	github.com/aws/aws-sdk-go-v2/service/signin v1.0.7 // indirect
+	github.com/aws/aws-sdk-go-v2/service/sso v1.30.12 // indirect
+	github.com/aws/aws-sdk-go-v2/service/ssooidc v1.35.16 // indirect
+	github.com/aws/aws-sdk-go-v2/service/sts v1.41.8 // indirect
+	github.com/aws/smithy-go v1.24.2 // indirect
 	github.com/beorn7/perks v1.0.1 // indirect
 	github.com/cenkalti/backoff/v5 v5.0.3 // indirect
 	github.com/containerd/errdefs/pkg v0.3.0 // indirect

go.sum

@@ -22,6 +22,38 @@ github.com/Yiling-J/theine-go v0.6.2 h1:1GeoXeQ0O0AUkiwj2S9Jc0Mzx+hpqzmqsJ4kIC4M
 github.com/Yiling-J/theine-go v0.6.2/go.mod h1:08QpMa5JZ2pKN+UJCRrCasWYO1IKCdl54Xa836rpmDU=
 github.com/antlr4-go/antlr/v4 v4.13.1 h1:SqQKkuVZ+zWkMMNkjy5FZe5mr5WURWnlpmOuzYWrPrQ=
 github.com/antlr4-go/antlr/v4 v4.13.1/go.mod h1:GKmUxMtwp6ZgGwZSva4eWPC5mS6vUAmOABFgjdkM7Nw=
+github.com/aws/aws-sdk-go-v2 v1.41.3 h1:4kQ/fa22KjDt13QCy1+bYADvdgcxpfH18f0zP542kZA=
+github.com/aws/aws-sdk-go-v2 v1.41.3/go.mod h1:mwsPRE8ceUUpiTgF7QmQIJ7lgsKUPQOUl3o72QBrE1o=
+github.com/aws/aws-sdk-go-v2/config v1.32.11 h1:ftxI5sgz8jZkckuUHXfC/wMUc8u3fG1vQS0plr2F2Zs=
+github.com/aws/aws-sdk-go-v2/config v1.32.11/go.mod h1:twF11+6ps9aNRKEDimksp923o44w/Thk9+8YIlzWMmo=
+github.com/aws/aws-sdk-go-v2/credentials v1.19.11 h1:NdV8cwCcAXrCWyxArt58BrvZJ9pZ9Fhf9w6Uh5W3Uyc=
+github.com/aws/aws-sdk-go-v2/credentials v1.19.11/go.mod h1:30yY2zqkMPdrvxBqzI9xQCM+WrlrZKSOpSJEsylVU+8=
+github.com/aws/aws-sdk-go-v2/feature/dsql/auth v1.1.19 h1:jS9Xb83Lr9ilcmR8NhhIh2G52eQ2Dd2Ng2p8xv23uK8=
+github.com/aws/aws-sdk-go-v2/feature/dsql/auth v1.1.19/go.mod h1:XExMXXv/juteYG/NHNCMzsExbXzByiAJMmVMcdgPrOo=
+github.com/aws/aws-sdk-go-v2/feature/ec2/imds v1.18.19 h1:INUvJxmhdEbVulJYHI061k4TVuS3jzzthNvjqvVvTKM=
+github.com/aws/aws-sdk-go-v2/feature/ec2/imds v1.18.19/go.mod h1:FpZN2QISLdEBWkayloda+sZjVJL+e9Gl0k1SyTgcswU=
+github.com/aws/aws-sdk-go-v2/internal/configsources v1.4.19 h1:/sECfyq2JTifMI2JPyZ4bdRN77zJmr6SrS1eL3augIA=
+github.com/aws/aws-sdk-go-v2/internal/configsources v1.4.19/go.mod h1:dMf8A5oAqr9/oxOfLkC/c2LU/uMcALP0Rgn2BD5LWn0=
+github.com/aws/aws-sdk-go-v2/internal/endpoints/v2 v2.7.19 h1:AWeJMk33GTBf6J20XJe6qZoRSJo0WfUhsMdUKhoODXE=
+github.com/aws/aws-sdk-go-v2/internal/endpoints/v2 v2.7.19/go.mod h1:+GWrYoaAsV7/4pNHpwh1kiNLXkKaSoppxQq9lbH8Ejw=
+github.com/aws/aws-sdk-go-v2/internal/ini v1.8.5 h1:clHU5fm//kWS1C2HgtgWxfQbFbx4b6rx+5jzhgX9HrI=
+github.com/aws/aws-sdk-go-v2/internal/ini v1.8.5/go.mod h1:O3h0IK87yXci+kg6flUKzJnWeziQUKciKrLjcatSNcY=
+github.com/aws/aws-sdk-go-v2/service/internal/accept-encoding v1.13.6 h1:XAq62tBTJP/85lFD5oqOOe7YYgWxY9LvWq8plyDvDVg=
+github.com/aws/aws-sdk-go-v2/service/internal/accept-encoding v1.13.6/go.mod h1:x0nZssQ3qZSnIcePWLvcoFisRXJzcTVvYpAAdYX8+GI=
+github.com/aws/aws-sdk-go-v2/service/internal/presigned-url v1.13.19 h1:X1Tow7suZk9UCJHE1Iw9GMZJJl0dAnKXXP1NaSDHwmw=
+github.com/aws/aws-sdk-go-v2/service/internal/presigned-url v1.13.19/go.mod h1:/rARO8psX+4sfjUQXp5LLifjUt8DuATZ31WptNJTyQA=
+github.com/aws/aws-sdk-go-v2/service/signin v1.0.7 h1:Y2cAXlClHsXkkOvWZFXATr34b0hxxloeQu/pAZz2row=
+github.com/aws/aws-sdk-go-v2/service/signin v1.0.7/go.mod h1:idzZ7gmDeqeNrSPkdbtMp9qWMgcBwykA7P7Rzh5DXVU=
+github.com/aws/aws-sdk-go-v2/service/sso v1.30.12 h1:iSsvB9EtQ09YrsmIc44Heqlx5ByGErqhPK1ZQLppias=
+github.com/aws/aws-sdk-go-v2/service/sso v1.30.12/go.mod h1:fEWYKTRGoZNl8tZ77i61/ccwOMJdGxwOhWCkp6TXAr0=
+github.com/aws/aws-sdk-go-v2/service/ssooidc v1.35.16 h1:EnUdUqRP1CNzt2DkV67tJx6XDN4xlfBFm+bzeNOQVb0=
+github.com/aws/aws-sdk-go-v2/service/ssooidc v1.35.16/go.mod h1:Jic/xv0Rq/pFNCh3WwpH4BEqdbSAl+IyHro8LbibHD8=
+github.com/aws/aws-sdk-go-v2/service/sts v1.41.8 h1:XQTQTF75vnug2TXS8m7CVJfC2nniYPZnO1D4Np761Oo=
+github.com/aws/aws-sdk-go-v2/service/sts v1.41.8/go.mod h1:Xgx+PR1NUOjNmQY+tRMnouRp83JRM8pRMw/vCaVhPkI=
+github.com/aws/smithy-go v1.24.2 h1:FzA3bu/nt/vDvmnkg+R8Xl46gmzEDam6mZ1hzmwXFng=
+github.com/aws/smithy-go v1.24.2/go.mod h1:YE2RhdIuDbA5E5bTdciG9KrW3+TiEONeUWCqxX9i1Fc=
+github.com/awslabs/aurora-dsql-connectors/go/pgx v0.2.0 h1:irsxnXcNaqkVmkmp5lkzpJ+dNxTCojuuEWGWL+gKBcg=
+github.com/awslabs/aurora-dsql-connectors/go/pgx v0.2.0/go.mod h1:KmZA9VBz4ubBFxTN7/FWo+V6sbcwPr/n1Qup4KXBV34=
 github.com/benbjohnson/clock v1.1.0/go.mod h1:J11/hYXuz8f4ySSvYwY0FKfm+ezbsZBKZxNJlLklBHA=
 github.com/beorn7/perks v1.0.1 h1:VlbKKnNfV8bJzeqoa4cOKqO6bYr3WgKZxO8Z16+hsOM=
 github.com/beorn7/perks v1.0.1/go.mod h1:G2ZrVWU2WbWT9wwq4/hrbKbnv/1ERSJQ0ibhJ6rlkpw=

internal/dsql/dsql.go

@@ -0,0 +1,90 @@
+// Package dsql provides shared utilities for Aurora DSQL connections.
+package dsql
+
+import (
+	"context"
+	"database/sql"
+	"fmt"
+	"net/url"
+	"strings"
+
+	"github.com/awslabs/aurora-dsql-connectors/go/pgx/dsql"
+)
+
+// PreparePostgresURI converts a dsql:// URI to a postgres:// URI with IAM authentication.
+// If username is provided, it overrides any username in the URI for token generation.
+func PreparePostgresURI(uri string, username string) (string, error) {
+	// Inject username into URI for token generation if provided.
+	tokenURI := uri
+	if username != "" {
+		pgURI := "postgres" + strings.TrimPrefix(uri, "dsql")
+		parsed, err := url.Parse(pgURI)
+		if err != nil {
+			return "", fmt.Errorf("parse URI for username override: %w", err)
+		}
+		parsed.User = url.User(username)
+		tokenURI = "dsql" + strings.TrimPrefix(parsed.String(), "postgres")
+	}
+
+	token, err := dsql.GenerateTokenConnString(context.Background(), tokenURI)
+	if err != nil {
+		return "", fmt.Errorf("generate DSQL auth token: %w", err)
+	}
+
+	pgURI := "postgres" + strings.TrimPrefix(uri, "dsql")
+	dbURI, err := url.Parse(pgURI)
+	if err != nil {
+		return "", fmt.Errorf("parse database URI: %w", err)
+	}
+
+	// Determine final username: param > URI > default "admin"
+	finalUser := username
+	if finalUser == "" {
+		if dbURI.User != nil {
+			finalUser = dbURI.User.Username()
+		}
+		if finalUser == "" {
+			finalUser = "admin"
+		}
+	}
+	dbURI.User = url.UserPassword(finalUser, token)
+
+	q := dbURI.Query()
+	q.Set("sslmode", "require")
+	q.Del("region")
+	dbURI.RawQuery = q.Encode()
+
+	return dbURI.String(), nil
+}
+
+// GooseTableDDL is the DDL for creating a DSQL-compatible goose version table.
+// Uses BIGINT GENERATED BY DEFAULT AS IDENTITY, as required by DSQL.
+const GooseTableDDL = `
+	CREATE TABLE IF NOT EXISTS goose_db_version (
+		id BIGINT GENERATED BY DEFAULT AS IDENTITY (CACHE 1) PRIMARY KEY,
+		version_id BIGINT NOT NULL,
+		is_applied BOOLEAN NOT NULL,
+		tstamp TIMESTAMP DEFAULT now()
+	)
+`
+
+// EnsureGooseTable creates the goose_db_version table if it doesn't exist.
+func EnsureGooseTable(db *sql.DB) error {
+	if _, err := db.Exec(GooseTableDDL); err != nil {
+		return fmt.Errorf("create goose table: %w", err)
+	}
+
+	// Goose expects an initial row with version 0 to exist.
+	var hasRows bool
+	if err := db.QueryRow(`SELECT EXISTS (SELECT 1 FROM goose_db_version)`).Scan(&hasRows); err != nil {
+		return fmt.Errorf("check goose table: %w", err)
+	}
+
+	if !hasRows {
+		if _, err := db.Exec(`INSERT INTO goose_db_version (version_id, is_applied) VALUES (0, TRUE)`); err != nil {
+			return fmt.Errorf("insert initial goose row: %w", err)
+		}
+	}
+
+	return nil
+}