feat: add Aurora DSQL storage backend

Add support for Amazon Aurora DSQL as a storage engine for OpenFGA.

- Add DSQL-specific migrations adapted for DSQL constraints (ASYNC
  index creation, NO TRANSACTION DDL, epoch-based goose table IDs,
  full indexes instead of partial)
- Add OCC retry logic with exponential backoff for write operations
- Skip FOR UPDATE clause (DSQL uses OCC at commit time)
- Use aurora-dsql-connectors/go/pgx for IAM authentication
- Optimize query plans to avoid table scans on DSQL
- Respect OPENFGA_DATASTORE_USERNAME for DSQL connections
- Add DSQL migration tooling and test fixtures

Author: amaksimo

CHANGELOG.md

@@ -45,11 +45,13 @@ Try to keep listed changes to a concise bulleted list of simple explanations of
 
 ### Changed
 - Datastore throttling separated from dispatch throttling in BatchCheck, ListUsers metadata. Also, `throttling_type` label added to `throttledRequestCounter` metric to differentiate between dispatch/datastore throttling. [#2839](https://github.com/openfga/openfga/pull/2839)
+- Update Aurora DSQL connector to use the new official monorepo location (`github.com/awslabs/aurora-dsql-connectors/go/pgx`). [#15](https://github.com/amaksimo/openfga-dsql-alemaksi/pull/15)
 
 ### Removed
 - Removed custom grpc_prometheus fork, replace with go-grpc-middleware's provider. Removes the custom `grpc_code` label on this metric. [#2855](https://github.com/openfga/openfga/pull/2855)
 
 ### Fixed
+- DSQL connections now respect `OPENFGA_DATASTORE_USERNAME` for IAM token generation. [#12](https://github.com/amaksimo/openfga-dsql-alemaksi/pull/12)
 - ListUsers will now properly get datastore throttled if enabled. [#2846](https://github.com/openfga/openfga/pull/2846)
 - Cache controller now uses the logger provided to the server instead of always using a no-op logger. [#2847](https://github.com/openfga/openfga/pull/2847)
 - Typesystem invalidate model with empty intersection and union. [#2865](https://github.com/openfga/openfga/pull/2865)

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` since DSQL DDL is always non-transactional
+3. **No SERIAL/IDENTITY**: The goose version table uses epoch microseconds instead of auto-increment
+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)
 	}