Skip to content

mpg: route attach connection info through the public API #5191

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Draft
tripledoublev wants to merge 1 commit into
base: master
Choose a base branch
from
+254 −24
Draft

mpg: route attach connection info through the public API #5191

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
83 changes: 71 additions & 12 deletions internal/command/mpg/v2/run_attach.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -11,6 +11,7 @@ import (
"github.com/superfly/flyctl/internal/appsecrets"
"github.com/superfly/flyctl/internal/flag"
"github.com/superfly/flyctl/internal/flapsutil"
"github.com/superfly/flyctl/internal/mpgutil"
"github.com/superfly/flyctl/internal/prompt"
mpgv2 "github.com/superfly/flyctl/internal/uiex/mpg/v2"
"github.com/superfly/flyctl/iostreams"
Expand Down Expand Up @@ -131,17 +132,12 @@ func RunAttach(ctx context.Context, clusterID string) error {
}
}

// Get cluster details and credentials. The public Machines API cluster show does
// not expose credentials, so we always use the legacy client for the connection
// URI and default DB name. This preserves the existing behavior.
clusterResp, err := legacyClient.GetClusterById(ctx, clusterID)
info, err := getClusterConnectionInfoPublicFirst(ctx, flapsClient, legacyClient, clusterID, username == "")
if err != nil {
return fmt.Errorf("failed retrieving cluster %s: %w", clusterID, err)
}

baseUri := clusterResp.Credentials.ConnectionUri
if baseUri == "" {
return fmt.Errorf("connection URI is empty; cannot attach without valid credentials")
if info.BaseURI == "" {
return fmt.Errorf("cluster is not ready; cannot attach without valid connection information")
}

var connectionUri string
Expand All @@ -155,14 +151,14 @@ func RunAttach(ctx context.Context, clusterID string) error {
user = creds.User
password = creds.Password
} else {
user = clusterResp.Credentials.User
password = clusterResp.Credentials.Password
user = info.DefaultUser
password = info.DefaultPassword
}

if db == "" {
db = clusterResp.Credentials.DBName
db = info.DefaultDBName
}
connectionUri, err = buildConnectionUri(baseUri, user, password, db)
connectionUri, err = buildConnectionUri(info.BaseURI, user, password, db)
if err != nil {
return fmt.Errorf("failed to build connection URI: %w", err)
}
Expand Down Expand Up @@ -331,3 +327,66 @@ func createAttachmentPublicFirst(ctx context.Context, flapsClient flapsutil.Flap

return nil
}

// clusterConnectionInfo holds the base URI and default credentials for attach.
type clusterConnectionInfo struct {
BaseURI string
DefaultUser string
DefaultPassword string
DefaultDBName string
}

// getClusterConnectionInfoPublicFirst falls back on cluster or credential 404s.
func getClusterConnectionInfoPublicFirst(ctx context.Context, flapsClient flapsutil.FlapsClient, legacyClient mpgv2.ClientV2, clusterID string, needDefaultCredentials bool) (clusterConnectionInfo, error) {
cluster, err := flapsClient.GetManagedPostgresCluster(ctx, clusterID)
if errors.Is(err, flaps.ErrFlapsNotFound) {
return getClusterConnectionInfoLegacy(ctx, legacyClient, clusterID)
}
if err != nil {
return clusterConnectionInfo{}, err
}

if cluster.Status == flaps.ManagedPostgresStatusFailed || cluster.Status == flaps.ManagedPostgresStatusError {
return clusterConnectionInfo{}, fmt.Errorf("cluster is in a failed state (status: %s); cannot attach", cluster.Status)
}

// Check before credentials so a credentials 404 cannot bypass the host check.
pooler := cluster.Endpoints.Primary.Pooler
if pooler.Host == "" || pooler.Port == 0 {
return clusterConnectionInfo{}, nil
}

info := clusterConnectionInfo{
BaseURI: fmt.Sprintf("postgres://%s:%d/%s", pooler.Host, pooler.Port, mpgutil.DefaultDatabase),
DefaultDBName: mpgutil.DefaultDatabase,
}
if !needDefaultCredentials {
return info, nil
}

creds, err := flapsClient.GetManagedPostgresUserCredentials(ctx, clusterID, mpgutil.DefaultUsername)
if errors.Is(err, flaps.ErrFlapsNotFound) {
return getClusterConnectionInfoLegacy(ctx, legacyClient, clusterID)
}
if err != nil {
return clusterConnectionInfo{}, err
}
info.DefaultUser = creds.Username
info.DefaultPassword = creds.Password

return info, nil
}

func getClusterConnectionInfoLegacy(ctx context.Context, legacyClient mpgv2.ClientV2, clusterID string) (clusterConnectionInfo, error) {
clusterResp, err := legacyClient.GetClusterById(ctx, clusterID)
if err != nil {
return clusterConnectionInfo{}, err
}

return clusterConnectionInfo{
BaseURI: clusterResp.Credentials.ConnectionUri,
DefaultUser: clusterResp.Credentials.User,
DefaultPassword: clusterResp.Credentials.Password,
DefaultDBName: clusterResp.Credentials.DBName,
}, nil
}
195 changes: 183 additions & 12 deletions internal/command/mpg/v2/run_attach_test.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -52,6 +52,11 @@ func addAttachFlags(fs *pflag.FlagSet) {
// specific Func fields as needed.
func minimalAttachFlapsClient() *mock.FlapsClient {
return &mock.FlapsClient{
GetManagedPostgresClusterFunc: func(_ context.Context, id string) (flaps.ManagedPostgresCluster, error) {
cluster := flaps.ManagedPostgresCluster{ID: id}
cluster.Endpoints.Primary.Pooler = flaps.ManagedPostgresEndpoint{Host: "pooler.fly.dev", Port: 5432}
return cluster, nil
},
ListManagedPostgresUsersFunc: func(_ context.Context, id string) ([]flaps.ManagedPostgresUser, error) {
return []flaps.ManagedPostgresUser{{Username: "alice", Role: flaps.ManagedPostgresUserRoleWriter}}, nil
},
Expand Down Expand Up @@ -140,37 +145,50 @@ func TestRunAttach_publicSuccess(t *testing.T) {
require.Empty(t, stderr.String())
}

// TestRunAttach_usernameUsesClusterDatabase verifies that selecting a user without
// --database still preserves the cluster credential DB name
func TestRunAttach_usernameUsesClusterDatabase(t *testing.T) {
func TestRunAttach_explicitUsernameSkipsDefaultCredentials(t *testing.T) {
ctx, stdout, stderr, flags := attachTestContext(t)
addAttachFlags(flags)
require.NoError(t, flags.Set("username", "alice"))

ctx = flapsutil.NewContextWithClient(ctx, minimalAttachFlapsClient())
ctx = mpgv2.NewContextWithClient(ctx, minimalAttachLegacyClient())
flapsClient := minimalAttachFlapsClient()
var requestedUsers []string
flapsClient.GetManagedPostgresUserCredentialsFunc = func(_ context.Context, id, username string) (flaps.ManagedPostgresUserCredentials, error) {
requestedUsers = append(requestedUsers, username)
require.Equal(t, "mpg-123", id)
if username == "fly-user" {
return flaps.ManagedPostgresUserCredentials{}, errors.New("default credentials unavailable")
}

return flaps.ManagedPostgresUserCredentials{Username: username, Password: "alice-pass"}, nil
}
ctx = flapsutil.NewContextWithClient(ctx, flapsClient)
ctx = mpgv2.NewContextWithClient(ctx, &mock.MpgV2Client{})

require.NoError(t, RunAttach(ctx, "mpg-123"))
require.Equal(t, wantSecretOutput("my-app", "DATABASE_URL", "postgres://alice:alice-pass@pooler.fly.dev:5432/default_db"), stdout.String())
require.Equal(t, []string{"alice"}, requestedUsers)
require.Equal(t, wantSecretOutput("my-app", "DATABASE_URL", "postgres://alice:alice-pass@pooler.fly.dev:5432/fly-db"), stdout.String())
require.Empty(t, stderr.String())
}

// TestRunAttach_noUsernameDefaultCredentials verifies that when no username is provided,
// the legacy cluster credentials are used directly
func TestRunAttach_noUsernameDefaultCredentials(t *testing.T) {
ctx, stdout, stderr, flags := attachTestContext(t)
addAttachFlags(flags)
// No username or database flags set.

flapsClient := minimalAttachFlapsClient()
flapsClient.GetManagedPostgresUserCredentialsFunc = func(_ context.Context, id, username string) (flaps.ManagedPostgresUserCredentials, error) {
require.Equal(t, "fly-user", username, "default credentials path uses mpgutil.DefaultUsername")

return flaps.ManagedPostgresUserCredentials{Username: username, Password: "default-pass"}, nil
}
ctx = flapsutil.NewContextWithClient(ctx, flapsClient)
ctx = mpgv2.NewContextWithClient(ctx, minimalAttachLegacyClient())

err := RunAttach(ctx, "mpg-123")
require.NoError(t, err)

// Default credentials used; database defaults to default_db.
wantUri := "postgres://default_user:default-pass@pooler.fly.dev:5432/default_db"
// Default credentials used; database defaults to fly-db.
wantUri := "postgres://fly-user:default-pass@pooler.fly.dev:5432/fly-db"
require.Equal(t, wantSecretOutput("my-app", "DATABASE_URL", wantUri), stdout.String())
require.Empty(t, stderr.String())
}
Expand All @@ -194,7 +212,7 @@ func TestRunAttach_customVariableName(t *testing.T) {
err := RunAttach(ctx, "mpg-123")
require.NoError(t, err)

wantUri := "postgres://default_user:default-pass@pooler.fly.dev:5432/default_db"
wantUri := "postgres://fly-user:alice-pass@pooler.fly.dev:5432/fly-db"
require.Equal(t, wantSecretOutput("my-app", "MY_PG_URL", wantUri), stdout.String())
require.Empty(t, stderr.String())
}
Expand Down Expand Up @@ -224,6 +242,9 @@ func TestRunAttach_invalidConnectionUriError(t *testing.T) {
// No username.

flapsClient := minimalAttachFlapsClient()
flapsClient.GetManagedPostgresClusterFunc = func(_ context.Context, id string) (flaps.ManagedPostgresCluster, error) {
return flaps.ManagedPostgresCluster{}, flaps.ErrFlapsNotFound
}
ctx = flapsutil.NewContextWithClient(ctx, flapsClient)
ctx = mpgv2.NewContextWithClient(ctx, &mock.MpgV2Client{
GetClusterByIdFunc: func(_ context.Context, id string) (mpgv2.GetClusterResponse, error) {
Expand All @@ -240,7 +261,7 @@ func TestRunAttach_invalidConnectionUriError(t *testing.T) {

err := RunAttach(ctx, "mpg-123")
require.Error(t, err)
require.Contains(t, err.Error(), "connection URI is empty")
require.Contains(t, err.Error(), "cluster is not ready")
}

// TestRunAttach_attachmentWarningOnly verifies that a failed attachment creation produces
Expand Down Expand Up @@ -551,3 +572,153 @@ func TestCreateDatabasePublicFirst(t *testing.T) {
})
}
}

func TestGetClusterConnectionInfoPublicFirst(t *testing.T) {
publicCluster := flaps.ManagedPostgresCluster{ID: "mpg-123"}
publicCluster.Endpoints.Primary.Pooler = flaps.ManagedPostgresEndpoint{Host: "pooler.fly.dev", Port: 5432}
publicCreds := flaps.ManagedPostgresUserCredentials{Username: "fly-user", Password: "public-pass"}

legacyResp := mpgv2.GetClusterResponse{
Credentials: mpgv2.GetClusterCredentialsResponse{
User: "legacy_user",
Password: "legacy-pass",
DBName: "legacy_db",
ConnectionUri: "postgres://legacy_user:legacy-pass@legacy.host:5432/legacy_db",
},
}

legacyInfo := clusterConnectionInfo{
BaseURI: "postgres://legacy_user:legacy-pass@legacy.host:5432/legacy_db",
DefaultUser: "legacy_user",
DefaultPassword: "legacy-pass",
DefaultDBName: "legacy_db",
}

tests := []struct {
name string
mutate func(*flaps.ManagedPostgresCluster)
publicClusterErr error
publicCredsErr error
legacyErr error
wantCredsCalls int
wantLegacyCalls int
wantInfo clusterConnectionInfo
wantErr string
}{
{
name: "uses Machines API for cluster and default creds",
wantCredsCalls: 1,
wantInfo: clusterConnectionInfo{
BaseURI: "postgres://pooler.fly.dev:5432/fly-db",
DefaultUser: "fly-user",
DefaultPassword: "public-pass",
DefaultDBName: "fly-db",
},
},
{
name: "empty pooler host (not-ready cluster) yields empty BaseURI without legacy fallback",
mutate: func(c *flaps.ManagedPostgresCluster) { c.Endpoints.Primary.Pooler.Host = "" },
wantInfo: clusterConnectionInfo{},
},
{
name: "zero port is treated as not-ready (not substituted with DefaultPort)",
mutate: func(c *flaps.ManagedPostgresCluster) { c.Endpoints.Primary.Pooler.Port = 0 },
wantInfo: clusterConnectionInfo{},
},
{
name: "falls back to legacy bundle on cluster 404",
publicClusterErr: flaps.ErrFlapsNotFound,
wantLegacyCalls: 1,
wantInfo: legacyInfo,
},
{
name: "falls back to legacy bundle on default-creds 404",
publicCredsErr: flaps.ErrFlapsNotFound,
wantCredsCalls: 1,
wantLegacyCalls: 1,
wantInfo: legacyInfo,
},
{
name: "propagates non-404 public cluster error without fallback",
publicClusterErr: errors.New("internal server error"),
wantErr: "internal server error",
},
{
name: "propagates non-404 public creds error without fallback",
publicCredsErr: errors.New("conflict: default user unavailable"),
wantCredsCalls: 1,
wantErr: "conflict: default user unavailable",
},
{
name: "propagates legacy error after cluster 404 fallback",
publicClusterErr: flaps.ErrFlapsNotFound,
legacyErr: errors.New("legacy boom"),
wantLegacyCalls: 1,
wantErr: "legacy boom",
},
{
name: "propagates legacy error after default-creds 404 fallback",
publicCredsErr: flaps.ErrFlapsNotFound,
legacyErr: errors.New("legacy boom"),
wantCredsCalls: 1,
wantLegacyCalls: 1,
wantErr: "legacy boom",
},
{
name: "cluster in failed state returns distinct error",
mutate: func(c *flaps.ManagedPostgresCluster) { c.Status = flaps.ManagedPostgresStatusFailed },
wantErr: "is in a failed state",
},
{
name: "cluster in error state returns distinct error",
mutate: func(c *flaps.ManagedPostgresCluster) { c.Status = flaps.ManagedPostgresStatusError },
wantErr: "is in a failed state",
},
}

for _, tt := range tests {
t.Run(tt.name, func(t *testing.T) {
ctx := context.Background()
cluster := publicCluster
if tt.mutate != nil {
tt.mutate(&cluster)
}
clusterCalls, credsCalls, legacyCalls := 0, 0, 0
flapsClient := &mock.FlapsClient{
GetManagedPostgresClusterFunc: func(_ context.Context, id string) (flaps.ManagedPostgresCluster, error) {
clusterCalls++
require.Equal(t, "mpg-123", id)

return cluster, tt.publicClusterErr
},
GetManagedPostgresUserCredentialsFunc: func(_ context.Context, id, username string) (flaps.ManagedPostgresUserCredentials, error) {
credsCalls++
require.Equal(t, "mpg-123", id)
require.Equal(t, "fly-user", username)

return publicCreds, tt.publicCredsErr
},
}
legacyClient := &mock.MpgV2Client{
GetClusterByIdFunc: func(_ context.Context, id string) (mpgv2.GetClusterResponse, error) {
legacyCalls++
require.Equal(t, "mpg-123", id)

return legacyResp, tt.legacyErr
},
}

info, err := getClusterConnectionInfoPublicFirst(ctx, flapsClient, legacyClient, "mpg-123", true)
require.Equal(t, 1, clusterCalls)
require.Equal(t, tt.wantCredsCalls, credsCalls)
require.Equal(t, tt.wantLegacyCalls, legacyCalls)
if tt.wantErr != "" {
require.Error(t, err)
require.Contains(t, err.Error(), tt.wantErr)
} else {
require.NoError(t, err)
require.Equal(t, tt.wantInfo, info)
}
})
}
}
Loading