diff --git a/internal/command/mpg/v2/run_attach.go b/internal/command/mpg/v2/run_attach.go index 9d4b7b76df..2d48cbf701 100644 --- a/internal/command/mpg/v2/run_attach.go +++ b/internal/command/mpg/v2/run_attach.go @@ -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" @@ -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 @@ -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) } @@ -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 +} diff --git a/internal/command/mpg/v2/run_attach_test.go b/internal/command/mpg/v2/run_attach_test.go index 1d7c751783..393ce5ac4a 100644 --- a/internal/command/mpg/v2/run_attach_test.go +++ b/internal/command/mpg/v2/run_attach_test.go @@ -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 }, @@ -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()) } @@ -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()) } @@ -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) { @@ -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 @@ -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) + } + }) + } +}