feat: support Aurora DSQL connections (IAM token + SNI)#20646
feat: support Aurora DSQL connections (IAM token + SNI)#20646ouchi2501 wants to merge 3 commits into
Conversation
Aurora DSQL is PostgreSQL-compatible on the wire but differs from RDS/Aurora in two ways that previously blocked Bytebase from connecting: - It mandates SNI during the TLS handshake. GetTLSConfig built a tls.Config without ServerName, so the password+SSL path sent a ClientHello without SNI and DSQL rejected the connection. Set ServerName to the host, skipping literal IP addresses per RFC 6066. - It requires a DSQL-specific IAM auth token (dsql:DbConnectAdmin / dsql:DbConnect) generated by feature/dsql/auth, not the RDS token. Detect DSQL endpoints (*.dsql.*.on.aws) on the AWS_RDS_IAM path and sign the appropriate token, using the admin variant for the built-in "admin" user. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
|
We require contributors to sign our Contributor License Agreement, and we don't have yours on file. In order for us to review and merge your code, please sign CLA and add your name to contributors list. |
|
Review the following changes in direct dependencies. Learn more about Socket for GitHub.
|
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: 3b88b10ccc
ℹ️ About Codex in GitHub
Codex has been enabled to automatically review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
When you sign up for Codex through ChatGPT, Codex can also answer questions or update the PR, like "@codex address that feedback".
| // ClientHello carries no SNI and such servers reject the handshake. | ||
| // Per RFC 6066, SNI must not be a literal IP address. | ||
| if net.ParseIP(ds.GetHost()) == nil { | ||
| cfg.ServerName = ds.GetHost() |
There was a problem hiding this comment.
Avoid pinning SNI to the primary address
When SSL is enabled for a multi-address data source (for example MongoDB replica sets via additional_addresses, Redis clusters/sentinels, or Cassandra contact points), GetTLSConfig returns one tls.Config that those drivers reuse for every target address; setting ServerName here pins SNI to the primary host for all handshakes. Dialing a secondary that expects its own hostname/SNI can then route or fail incorrectly, whereas leaving ServerName empty lets the driver/tls dialer fill it from each actual address, or the caller can clone/set it per address.
Useful? React with 👍 / 👎.
There was a problem hiding this comment.
Good catch — fixed in 502d0bf. GetTLSConfig is shared across drivers (MongoDB/Redis/Cassandra multi-address), so I moved the ServerName assignment into the PostgreSQL-specific applyPGTLSConfigForHost, which runs per host (primary + fallbacks). SNI is now set per target address and other drivers are unaffected.
Signed the Bytebase CLA for PR bytebase/bytebase#20646. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
GetTLSConfig is shared by many drivers (MongoDB replica sets, Redis clusters, Cassandra contact points, etc.) that reuse one tls.Config for multiple target addresses. Setting ServerName there pinned SNI to the primary host for every handshake. Move the ServerName assignment into applyPGTLSConfigForHost, which is PostgreSQL-specific and runs per host (primary + fallbacks), so each target presents its own name and other drivers are unaffected. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: 502d0bf1a4
ℹ️ About Codex in GitHub
Codex has been enabled to automatically review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
When you sign up for Codex through ChatGPT, Codex can also answer questions or update the PR, like "@codex address that feedback".
| if util.IsAWSDSQLHost(conf.DataSource.Host) { | ||
| return getDSQLConnectionPassword(ctx, conf, cfg.Credentials) |
There was a problem hiding this comment.
Avoid unsupported read-only startup params for DSQL
When this branch treats a DSQL endpoint as AWS_RDS_IAM, it still goes through Open's generic PostgreSQL setup, which adds default_transaction_read_only=true to RuntimeParams whenever ConnectionContext.ReadOnly is set; the SQL service and data export paths set that flag for read-only data sources/exports. Aurora DSQL rejects unsupported startup/session parameters, and default_transaction_read_only is not in its accepted parameter list, so those DSQL read paths fail at connection time before any query can run. Special-case DSQL to avoid this startup parameter and enforce read-only another way.
Useful? React with 👍 / 👎.
There was a problem hiding this comment.
Confirmed against the AWS docs — DSQL's supported session parameter list does not include default_transaction_read_only, and it rejects unsupported startup params with "setting configuration parameter ... not supported". Fixed in cc4402d: the read-only startup parameter is now skipped for DSQL endpoints.
| if conf.DataSource.Username == "admin" { | ||
| token, err = dsqlauth.GenerateDBConnectAdminAuthToken(ctx, endpoint, region, creds) | ||
| } else { | ||
| token, err = dsqlauth.GenerateDbConnectAuthToken(ctx, endpoint, region, creds) | ||
| } |
There was a problem hiding this comment.
Regenerate DSQL IAM tokens per connection
This generates the DSQL IAM token once while building the driver config; getRDSConnectionConfig then stores that token as the static password in the registered pgx config that database/sql uses for any later physical connections. DSQL tokens are only valid for connection establishment and expire after their configured lifetime (15 minutes by default), so a long-running export/schema sync or an idle SQL session that opens or reopens a connection after expiry will start failing authentication until the driver is recreated. Generate the DSQL password per connection, or otherwise refresh the registered config before new connections use it.
Useful? React with 👍 / 👎.
There was a problem hiding this comment.
Valid point. Note this is the same behavior as the existing RDS IAM path (auth.BuildAuthToken is likewise baked in as a static password), so DSQL is consistent with how Bytebase already handles AWS IAM tokens here rather than introducing a new regression. A proper fix means regenerating the token per physical connection (e.g. switching this path to stdlib.OpenDB with OptionBeforeConnect), which also changes the connect/close lifecycle and should cover RDS and DSQL together. I'd prefer to keep this PR focused on enabling DSQL and address per-connection token refresh for both paths in a dedicated follow-up. Happy to file/track that separately.
Aurora DSQL only accepts a fixed set of session parameters and rejects anything outside it with "setting configuration parameter ... not supported". default_transaction_read_only is not in that set, so adding it as a startup parameter made every read-only DSQL connection (read-only data sources, data export) fail before any query ran. Skip the parameter when connecting to a DSQL endpoint. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
|
What this PR does
Adds support for connecting to Amazon Aurora DSQL from Bytebase.
Aurora DSQL speaks the PostgreSQL wire protocol, so Bytebase can treat it as a
POSTGRESengine. However two DSQL-specific requirements previously blocked theconnection:
1. SNI was not sent on the TLS handshake
util.GetTLSConfigbuilt atls.Config{}withoutServerName. On thepassword/SSL path (
getPGConnectionConfig) this config replaces the one pgxderives from the DSN, so the
ClientHellocarried no SNI. Aurora DSQL routes byhostname and rejects handshakes without SNI.
Fix: set
ServerNameto the data source host inGetTLSConfig, skippingliteral IP addresses per RFC 6066 (mirrors the pgx fork's own
sslsnibehavior).2. DSQL requires a different IAM auth token than RDS/Aurora
DSQL does not support password auth; it requires an IAM token generated for the
dsql:DbConnectAdmin/dsql:DbConnectactions viaaws-sdk-go-v2/feature/dsql/auth, which differs from the RDS token produced byfeature/rds/auth.Fix: on the existing
AWS_RDS_IAMpath, detect DSQL cluster endpoints(
*.dsql.*.on.aws) and sign a DSQL token. The admin variant is used for thebuilt-in
adminuser; otherwise the regularDbConnectvariant is used. Thetoken is signed against the cluster endpoint hostname (no port).
How to use
Configure the instance with authentication type = AWS RDS IAM, set the host
to the DSQL cluster endpoint (
<id>.dsql.<region>.on.aws) and the user toadmin(or a custom DB role). Bytebase auto-detects DSQL and signs the correcttoken. No new UI option is introduced.
Changes
backend/plugin/db/util/ssl.go— setServerName(SNI) inGetTLSConfig.backend/plugin/db/util/aws.go— addIsAWSDSQLHostendpoint detector.backend/plugin/db/util/aws_test.go— table-driven test forIsAWSDSQLHost.backend/plugin/db/pg/pg.go— branch to DSQL token generation on the RDS IAM path.go.mod/go.sum— addfeature/dsql/auth v1.1.23(requires core SDKv1.41.7, identical to the version already in use, so no core SDK bump).Testing
go build ./backend/plugin/db/...passes.go test ./backend/plugin/db/util/...passes (incl. the newIsAWSDSQLHosttest).gofmtapplied;go vet ./backend/plugin/db/...clean.Notes
password/SSL path, not only DSQL.
cluster are not auto-detected.
🤖 Generated with Claude Code