diff --git a/internal/command/mpg/v1/run_connect.go b/internal/command/mpg/v1/run_connect.go index 8f10e81262..66f599b5d5 100644 --- a/internal/command/mpg/v1/run_connect.go +++ b/internal/command/mpg/v1/run_connect.go @@ -3,6 +3,7 @@ package cmdv1 import ( "context" "fmt" + "io" "os" "os/exec" "os/signal" @@ -82,9 +83,7 @@ func RunConnect(ctx context.Context, clusterID string, resolvedOrgSlug string) ( return err } - if cluster.Status != "ready" { - fmt.Fprintf(io.ErrOut, "%s Cluster is not in ready state, currently: %s\n", aurora.Yellow("WARN"), cluster.Status) - } + maybeWarnNotReady(io.ErrOut, cluster) psqlPath, err := exec.LookPath("psql") if err != nil { @@ -103,15 +102,7 @@ func RunConnect(ctx context.Context, clusterID string, resolvedOrgSlug string) ( return err } - user := credentials.User - password := credentials.Password - - // Use selected database or fall back to default from credentials - if db == "" { - db = credentials.DBName - } - - connectUrl := fmt.Sprintf("postgresql://%s:%s@localhost:%s/%s", user, password, localProxyPort, db) + connectUrl := buildConnectURL(credentials, db, localProxyPort) // Allow Ctrl+C signals to hit psql psqlCtx, psqlCancel := context.WithCancel(context.WithoutCancel(ctx)) @@ -174,3 +165,21 @@ func RunConnect(ctx context.Context, clusterID string, resolvedOrgSlug string) ( return err } + +// buildConnectURL prefers the selected database over the credential default. +func buildConnectURL(credentials *mpgv1.GetManagedClusterCredentialsResponse, db string, localProxyPort string) string { + if db == "" { + db = credentials.DBName + } + + return fmt.Sprintf("postgresql://%s:%s@localhost:%s/%s", credentials.User, credentials.Password, localProxyPort, db) +} + +// maybeWarnNotReady warns when a cluster is not in ready state. +func maybeWarnNotReady(errOut io.Writer, cluster *mpgv1.ManagedCluster) { + if cluster == nil || cluster.Status == "ready" { + return + } + + fmt.Fprintf(errOut, "%s Cluster is not in ready state, currently: %s\n", aurora.Yellow("WARN"), cluster.Status) +} diff --git a/internal/command/mpg/v1/run_connect_test.go b/internal/command/mpg/v1/run_connect_test.go new file mode 100644 index 0000000000..9fda1fdf01 --- /dev/null +++ b/internal/command/mpg/v1/run_connect_test.go @@ -0,0 +1,40 @@ +package cmdv1 + +import ( + "bytes" + "testing" + + "github.com/stretchr/testify/require" + mpgv1 "github.com/superfly/flyctl/internal/uiex/mpg/v1" +) + +func TestMaybeWarnNotReady(t *testing.T) { + const name = "test-cluster" + tests := []struct { + name string + status string + wantWarn bool + }{ + {name: "ready silent", status: "ready", wantWarn: false}, + {name: "creating warns", status: "creating", wantWarn: true}, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + var buf bytes.Buffer + cluster := &mpgv1.ManagedCluster{Name: name, Status: tt.status} + maybeWarnNotReady(&buf, cluster) + + got := buf.String() + if !tt.wantWarn { + require.Empty(t, got, "no warning expected for status=%q", tt.status) + + return + } + // Check the warning text independently of ANSI color codes. + require.Contains(t, got, "WARN", "warning must contain the literal 'WARN' marker") + require.Contains(t, got, "Cluster is not in ready state, currently: "+tt.status) + require.True(t, bytes.HasSuffix(buf.Bytes(), []byte("\n")), "warning must end with a newline (pre-migration format)") + }) + } +} diff --git a/internal/command/mpg/v1/run_proxy.go b/internal/command/mpg/v1/run_proxy.go index a82b04a159..d93ab65eae 100644 --- a/internal/command/mpg/v1/run_proxy.go +++ b/internal/command/mpg/v1/run_proxy.go @@ -2,11 +2,18 @@ package cmdv1 import ( "context" + "errors" "fmt" + "strconv" + fly "github.com/superfly/fly-go" + "github.com/superfly/fly-go/flaps" "github.com/superfly/flyctl/agent" "github.com/superfly/flyctl/internal/flag" + "github.com/superfly/flyctl/internal/flapsutil" "github.com/superfly/flyctl/internal/flyutil" + "github.com/superfly/flyctl/internal/mpgutil" + "github.com/superfly/flyctl/internal/uiex/mpg" mpgv1 "github.com/superfly/flyctl/internal/uiex/mpg/v1" "github.com/superfly/flyctl/proxy" ) @@ -29,12 +36,12 @@ func GetMpgProxyParams( clusterID string, resolvedOrgSlug string, ) (*mpgv1.ManagedCluster, *proxy.ConnectParams, error) { - response, err := getCluster(ctx, clusterID) + response, _, port, err := getCluster(ctx, clusterID) if err != nil { return nil, nil, err } - cluster, params, err := buildProxyParams(ctx, response, localProxyPort, resolvedOrgSlug) + cluster, params, err := buildProxyParams(ctx, response, port, localProxyPort, resolvedOrgSlug) if err != nil { return nil, nil, err } @@ -42,8 +49,7 @@ func GetMpgProxyParams( return cluster, params, nil } -// GetMpgConnectParams builds proxy connection parameters and resolves the -// database credentials needed by fly mpg connect. +// GetMpgConnectParams resolves credentials and proxy parameters. func GetMpgConnectParams( ctx context.Context, localProxyPort string, @@ -51,17 +57,17 @@ func GetMpgConnectParams( clusterID string, resolvedOrgSlug string, ) (*mpgv1.ManagedCluster, *proxy.ConnectParams, *mpgv1.GetManagedClusterCredentialsResponse, error) { - response, err := getCluster(ctx, clusterID) + response, useLegacy, port, err := getCluster(ctx, clusterID) if err != nil { return nil, nil, nil, err } - credentials, err := resolveConnectCredentials(ctx, response, username) + credentials, err := resolveConnectCredentials(ctx, response, useLegacy, username) if err != nil { return nil, nil, nil, err } - cluster, params, err := buildProxyParams(ctx, response, localProxyPort, resolvedOrgSlug) + cluster, params, err := buildProxyParams(ctx, response, port, localProxyPort, resolvedOrgSlug) if err != nil { return nil, nil, nil, err } @@ -69,23 +75,62 @@ func GetMpgConnectParams( return cluster, params, credentials, nil } -func getCluster(ctx context.Context, clusterID string) (*mpgv1.GetManagedClusterResponse, error) { - mpgClient := mpgv1.ClientFromContext(ctx) - response, err := mpgClient.GetManagedClusterById(ctx, clusterID) +// getCluster tries the public API, falling back to the legacy client only on 404. +// It returns the credential source and direct endpoint port (5432 for legacy). +func getCluster(ctx context.Context, clusterID string) (*mpgv1.GetManagedClusterResponse, bool, int, error) { + flapsClient := flapsutil.ClientFromContext(ctx) + publicCluster, err := flapsClient.GetManagedPostgresCluster(ctx, clusterID) + if err == nil { + response, port := publicToLegacyClusterResponse(publicCluster) + + return &response, false, port, nil + } + + if !errors.Is(err, flaps.ErrFlapsNotFound) { + return nil, false, 0, fmt.Errorf("failed retrieving cluster %s: %w", clusterID, err) + } + + legacyClient := mpgv1.ClientFromContext(ctx) + response, err := legacyClient.GetManagedClusterById(ctx, clusterID) if err != nil { - return nil, fmt.Errorf("failed retrieving cluster %s: %w", clusterID, err) + return nil, true, 0, fmt.Errorf("failed retrieving cluster %s: %w", clusterID, err) } - return &response, nil + return &response, true, mpgutil.DefaultPort, nil +} + +// publicToLegacyClusterResponse adapts the public cluster to the legacy shape. +// The advertised direct port is returned unchanged for validation. +func publicToLegacyClusterResponse(c flaps.ManagedPostgresCluster) (mpgv1.GetManagedClusterResponse, int) { + port := c.Endpoints.Primary.Direct.Port + + return mpgv1.GetManagedClusterResponse{ + Data: mpgv1.ManagedCluster{ + Id: c.ID, + Name: c.Name, + Status: c.Status, + Region: c.Region, + Plan: c.Plan, + Disk: c.DiskSizeGB, + Replicas: c.Replicas, + Organization: fly.Organization{Name: c.Organization.Name, Slug: c.Organization.Slug}, + IpAssignments: mpg.ManagedClusterIpAssignments{Direct: c.Endpoints.Primary.Direct.Host}, + }, + }, port } +// resolveConnectCredentials uses the same API as the cluster lookup. +// Public credentials default to fly-user and fly-db. func resolveConnectCredentials( ctx context.Context, response *mpgv1.GetManagedClusterResponse, + useLegacy bool, username string, ) (*mpgv1.GetManagedClusterCredentialsResponse, error) { var credentials mpgv1.GetManagedClusterCredentialsResponse - if username != "" { + + switch { + case username != "" && useLegacy: mpgClient := mpgv1.ClientFromContext(ctx) userCreds, err := mpgClient.GetUserCredentials(ctx, response.Data.Id, username) if err != nil { @@ -97,19 +142,56 @@ func resolveConnectCredentials( Password: userCreds.Data.Password, DBName: response.Credentials.DBName, } - } else { + case username != "": + flapsClient := flapsutil.ClientFromContext(ctx) + userCreds, err := flapsClient.GetManagedPostgresUserCredentials(ctx, response.Data.Id, username) + if err != nil { + return nil, fmt.Errorf("failed retrieving credentials for user %s: %w", username, err) + } + + credentials = mpgv1.GetManagedClusterCredentialsResponse{ + User: userCreds.Username, + Password: userCreds.Password, + DBName: mpgutil.DefaultDatabase, + } + case useLegacy: credentials = response.Credentials - } + default: + flapsClient := flapsutil.ClientFromContext(ctx) + userCreds, err := flapsClient.GetManagedPostgresUserCredentials(ctx, response.Data.Id, mpgutil.DefaultUsername) + if err != nil { + if errors.Is(err, flaps.ErrFlapsNotFound) { + return nil, fmt.Errorf("cluster is still initializing, wait a bit more") + } - if username == "" { - if credentials.Status == "initializing" { - return nil, fmt.Errorf("cluster is still initializing, wait a bit more") + return nil, fmt.Errorf("failed retrieving credentials for user %s: %w", mpgutil.DefaultUsername, err) } - if credentials.Status == "error" || credentials.Password == "" { - return nil, fmt.Errorf("error getting cluster password") + credentials = mpgv1.GetManagedClusterCredentialsResponse{ + User: userCreds.Username, + Password: userCreds.Password, + DBName: mpgutil.DefaultDatabase, + } + } + + if useLegacy { + // Only legacy default-user credentials include a status. + if username == "" { + if credentials.Status == "initializing" { + return nil, fmt.Errorf("cluster is still initializing, wait a bit more") + } + + if credentials.Status == "error" || credentials.Password == "" { + return nil, fmt.Errorf("error getting cluster password") + } + } else if credentials.Password == "" { + return nil, fmt.Errorf("error getting user password") } } else if credentials.Password == "" { + if username == "" { + return nil, fmt.Errorf("error getting cluster password") + } + return nil, fmt.Errorf("error getting user password") } @@ -119,10 +201,11 @@ func resolveConnectCredentials( func buildProxyParams( ctx context.Context, response *mpgv1.GetManagedClusterResponse, + port int, localProxyPort string, resolvedOrgSlug string, ) (*mpgv1.ManagedCluster, *proxy.ConnectParams, error) { - cluster, params, err := proxyParams(response, localProxyPort, resolvedOrgSlug, flag.GetBindAddr(ctx), nil) + cluster, params, err := proxyParams(response, port, localProxyPort, resolvedOrgSlug, flag.GetBindAddr(ctx), nil) if err != nil { return nil, nil, err } @@ -147,6 +230,7 @@ func buildProxyParams( func proxyParams( response *mpgv1.GetManagedClusterResponse, + port int, localProxyPort string, resolvedOrgSlug string, bindAddr string, @@ -157,8 +241,12 @@ func proxyParams( return nil, nil, fmt.Errorf("error getting cluster IP") } + if port < 1 || port > 65535 { + return nil, nil, fmt.Errorf("invalid cluster port %d: must be between 1 and 65535", port) + } + return cluster, &proxy.ConnectParams{ - Ports: []string{localProxyPort, "5432"}, + Ports: []string{localProxyPort, strconv.Itoa(port)}, OrganizationSlug: resolvedOrgSlug, Dialer: dialer, BindAddr: bindAddr, diff --git a/internal/command/mpg/v1/run_proxy_test.go b/internal/command/mpg/v1/run_proxy_test.go index 884d0daf9d..a37f155323 100644 --- a/internal/command/mpg/v1/run_proxy_test.go +++ b/internal/command/mpg/v1/run_proxy_test.go @@ -3,12 +3,20 @@ package cmdv1 import ( "context" "errors" + "fmt" "net" + "strconv" "testing" + "github.com/spf13/pflag" "github.com/stretchr/testify/assert" "github.com/stretchr/testify/require" + fly "github.com/superfly/fly-go" + "github.com/superfly/fly-go/flaps" + "github.com/superfly/flyctl/internal/flag" + "github.com/superfly/flyctl/internal/flapsutil" "github.com/superfly/flyctl/internal/mock" + "github.com/superfly/flyctl/internal/mpgutil" "github.com/superfly/flyctl/internal/uiex/mpg" mpgv1 "github.com/superfly/flyctl/internal/uiex/mpg/v1" "github.com/superfly/flyctl/wg" @@ -53,7 +61,7 @@ func TestProxyParamsIgnoreCredentials(t *testing.T) { } dialer := &testDialer{} - cluster, params, err := proxyParams(&response, "15432", "test-org", "127.0.0.2", dialer) + cluster, params, err := proxyParams(&response, mpgutil.DefaultPort, "15432", "test-org", "127.0.0.2", dialer) require.NoError(t, err) assert.Same(t, &response.Data, cluster) assert.Equal(t, []string{"15432", "5432"}, params.Ports) @@ -66,42 +74,279 @@ func TestProxyParamsIgnoreCredentials(t *testing.T) { } func TestProxyParamsRequireDirectIP(t *testing.T) { - cluster, params, err := proxyParams(&mpgv1.GetManagedClusterResponse{}, "15432", "test-org", "127.0.0.1", nil) + cluster, params, err := proxyParams(&mpgv1.GetManagedClusterResponse{}, 0, "15432", "test-org", "127.0.0.1", nil) require.EqualError(t, err, "error getting cluster IP") assert.Nil(t, cluster) assert.Nil(t, params) } +func samplePublicCluster() flaps.ManagedPostgresCluster { + return flaps.ManagedPostgresCluster{ + ID: "mpg-123", + Name: "test-cluster", + Status: "ready", + Region: "ord", + Plan: "development", + DiskSizeGB: 10, + Replicas: 1, + Organization: flaps.ManagedPostgresOrganization{Name: "Test Org", Slug: "test-org"}, + Endpoints: flaps.ManagedPostgresEndpoints{ + Primary: struct { + Direct flaps.ManagedPostgresEndpoint `json:"direct"` + Pooler flaps.ManagedPostgresEndpoint `json:"pooler"` + }{ + Direct: flaps.ManagedPostgresEndpoint{Host: "10.0.0.1", Port: 5432}, + Pooler: flaps.ManagedPostgresEndpoint{Host: "10.0.0.2", Port: 6432}, + }, + }, + } +} + +func sampleLegacyCluster() mpgv1.GetManagedClusterResponse { + return mpgv1.GetManagedClusterResponse{ + Data: mpgv1.ManagedCluster{ + Id: "mpg-123", Name: "test-cluster", Region: "ord", Status: "ready", + Plan: "development", Disk: 10, Replicas: 1, + Organization: fly.Organization{Name: "Test Org", Slug: "test-org"}, + IpAssignments: mpg.ManagedClusterIpAssignments{Direct: "10.0.0.1"}, + }, + Credentials: mpgv1.GetManagedClusterCredentialsResponse{ + Status: "ready", User: "app", Password: "secret", + DBName: "app", ConnectionUri: "postgres://app:secret@10.0.0.1:5432/app", + }, + } +} + +func TestGetCluster(t *testing.T) { + tests := []struct { + name string + publicCluster flaps.ManagedPostgresCluster + publicErr error + legacyResponse mpgv1.GetManagedClusterResponse + legacyErr error + wantUseLegacy bool + wantDirectHost string + wantStatus string + wantErr string + wantLegacyCalls int + }{ + { + name: "public success maps to legacy shape", + publicCluster: samplePublicCluster(), + wantUseLegacy: false, + wantDirectHost: "10.0.0.1", + wantStatus: "ready", + wantLegacyCalls: 0, + }, + { + name: "public success with empty host preserves empty", + publicCluster: func() flaps.ManagedPostgresCluster { + c := samplePublicCluster() + c.Endpoints.Primary.Direct.Host = "" + + return c + }(), + wantUseLegacy: false, + wantDirectHost: "", + wantStatus: "ready", + wantLegacyCalls: 0, + }, + { + name: "classified 404 falls back to legacy", + publicCluster: flaps.ManagedPostgresCluster{}, + publicErr: fmt.Errorf("wrapped: %w", &flaps.FlapsError{ + ResponseStatusCode: 404, + OriginalError: errors.New("not found"), + }), + legacyResponse: sampleLegacyCluster(), + wantUseLegacy: true, + wantDirectHost: "10.0.0.1", + wantStatus: "ready", + wantLegacyCalls: 1, + }, + { + name: "classified 404 with legacy failure propagates legacy error", + publicCluster: flaps.ManagedPostgresCluster{}, + publicErr: fmt.Errorf("wrapped: %w", &flaps.FlapsError{ + ResponseStatusCode: 404, + OriginalError: errors.New("not found"), + }), + legacyResponse: mpgv1.GetManagedClusterResponse{}, + legacyErr: errors.New("legacy denied"), + wantErr: "failed retrieving cluster mpg-123: legacy denied", + wantLegacyCalls: 1, + }, + { + name: "403 public error returns without fallback", + publicCluster: flaps.ManagedPostgresCluster{}, + publicErr: &flaps.FlapsError{ResponseStatusCode: 403, OriginalError: errors.New("denied")}, + wantErr: "failed retrieving cluster mpg-123: denied", + wantLegacyCalls: 0, + }, + { + name: "410 public error returns without fallback", + publicCluster: flaps.ManagedPostgresCluster{}, + publicErr: &flaps.FlapsError{ResponseStatusCode: 410, OriginalError: errors.New("gone")}, + wantErr: "failed retrieving cluster mpg-123: gone", + wantLegacyCalls: 0, + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + publicCalls, legacyCalls := 0, 0 + ctx := flapsutil.NewContextWithClient(context.Background(), &mock.FlapsClient{ + GetManagedPostgresClusterFunc: func(_ context.Context, id string) (flaps.ManagedPostgresCluster, error) { + publicCalls++ + require.Equal(t, "mpg-123", id) + + return tt.publicCluster, tt.publicErr + }, + }) + ctx = mpgv1.NewContextWithClient(ctx, &mock.MpgV1Client{ + GetManagedClusterByIdFunc: func(_ context.Context, id string) (mpgv1.GetManagedClusterResponse, error) { + legacyCalls++ + require.Equal(t, "mpg-123", id) + + return tt.legacyResponse, tt.legacyErr + }, + }) + + got, useLegacy, port, err := getCluster(ctx, "mpg-123") + if tt.wantErr != "" { + require.EqualError(t, err, tt.wantErr) + require.Nil(t, got) + require.Equal(t, 1, publicCalls) + require.Equal(t, tt.wantLegacyCalls, legacyCalls) + + return + } + require.NoError(t, err) + require.Equal(t, tt.wantUseLegacy, useLegacy) + require.NotNil(t, got) + require.Equal(t, tt.wantDirectHost, got.Data.IpAssignments.Direct) + require.Equal(t, tt.wantStatus, got.Data.Status) + require.Equal(t, 1, publicCalls) + require.Equal(t, tt.wantLegacyCalls, legacyCalls) + require.Equal(t, 5432, port) + if tt.wantUseLegacy { + _, params, err := proxyParams(got, port, "16380", "test-org", "127.0.0.1", nil) + require.NoError(t, err) + require.Equal(t, []string{"16380", "5432"}, params.Ports) + } + }) + } +} + +func TestGetMpgProxyParamsPublicNeverTouchesCredentials(t *testing.T) { + credCalls, legacyCredCalls := 0, 0 + ctx := flag.NewContext(context.Background(), pflag.NewFlagSet("test", pflag.ContinueOnError)) + ctx = flapsutil.NewContextWithClient(ctx, &mock.FlapsClient{ + GetManagedPostgresClusterFunc: func(context.Context, string) (flaps.ManagedPostgresCluster, error) { + cluster := samplePublicCluster() + cluster.Status = "creating" + cluster.Endpoints.Primary.Direct.Host = "" + + return cluster, nil + }, + GetManagedPostgresUserCredentialsFunc: func(context.Context, string, string) (flaps.ManagedPostgresUserCredentials, error) { + credCalls++ + + return flaps.ManagedPostgresUserCredentials{}, nil + }, + }) + ctx = mpgv1.NewContextWithClient(ctx, &mock.MpgV1Client{ + GetUserCredentialsFunc: func(context.Context, string, string) (mpgv1.GetUserCredentialsResponse, error) { + legacyCredCalls++ + + return mpgv1.GetUserCredentialsResponse{}, nil + }, + }) + + cluster, params, err := GetMpgProxyParams(ctx, "15432", "mpg-123", "test-org") + require.EqualError(t, err, "error getting cluster IP") + require.Nil(t, cluster) + require.Nil(t, params) + require.Zero(t, credCalls) + require.Zero(t, legacyCredCalls) +} + +func TestGetMpgConnectParamsResolvesCredentialsBeforeTunnel(t *testing.T) { + ctx := flapsutil.NewContextWithClient(context.Background(), &mock.FlapsClient{ + GetManagedPostgresClusterFunc: func(context.Context, string) (flaps.ManagedPostgresCluster, error) { + return samplePublicCluster(), nil + }, + GetManagedPostgresUserCredentialsFunc: func(context.Context, string, string) (flaps.ManagedPostgresUserCredentials, error) { + return flaps.ManagedPostgresUserCredentials{Username: mpgutil.DefaultUsername}, nil + }, + }) + + cluster, params, credentials, err := GetMpgConnectParams(ctx, "15432", "", "mpg-123", "test-org") + require.EqualError(t, err, "error getting cluster password") + require.Nil(t, cluster) + require.Nil(t, params) + require.Nil(t, credentials) +} + func TestResolveDefaultConnectCredentials(t *testing.T) { + // Legacy readiness comes from the credential envelope, not cluster status. + const name = "test-cluster" tests := []struct { - name string - credentials mpgv1.GetManagedClusterCredentialsResponse - err string + name string + clusterSt string // response.Data.Status — cluster status. NOT consulted on the legacy path. + credStatus string // credentials.Status — the legacy envelope's own status field (post-fetch). + credPwd string // credentials.Password — empty-password fallback when status checks pass. + err string }{ { - name: "initializing", - credentials: mpgv1.GetManagedClusterCredentialsResponse{Status: "initializing"}, - err: "cluster is still initializing, wait a bit more", + name: "empty password", + clusterSt: "ready", + credStatus: "ready", + credPwd: "", + err: "error getting cluster password", + }, + { + name: "ready cluster with stale initializing credentials refuses", + clusterSt: "ready", + credStatus: "initializing", + credPwd: "secret", + err: "cluster is still initializing, wait a bit more", }, { - name: "error status", - credentials: mpgv1.GetManagedClusterCredentialsResponse{Status: "error", Password: "password"}, - err: "error getting cluster password", + name: "ready cluster with stale error credentials refuses", + clusterSt: "ready", + credStatus: "error", + credPwd: "secret", + err: "error getting cluster password", }, { - name: "empty password", - credentials: mpgv1.GetManagedClusterCredentialsResponse{Status: "ready"}, - err: "error getting cluster password", + name: "creating cluster with ready credentials proceeds", + clusterSt: "creating", + credStatus: "ready", + credPwd: "secret", + err: "", // no error }, } for _, tt := range tests { t.Run(tt.name, func(t *testing.T) { - response := &mpgv1.GetManagedClusterResponse{Credentials: tt.credentials} + response := &mpgv1.GetManagedClusterResponse{ + Data: mpgv1.ManagedCluster{Name: name, Status: tt.clusterSt}, + Credentials: mpgv1.GetManagedClusterCredentialsResponse{ + Status: tt.credStatus, + Password: tt.credPwd, + }, + } + + credentials, err := resolveConnectCredentials(context.Background(), response, true, "") - credentials, err := resolveConnectCredentials(context.Background(), response, "") + if tt.err == "" { + require.NoError(t, err) + require.NotNil(t, credentials) + return + } require.EqualError(t, err, tt.err) assert.Nil(t, credentials) }) @@ -110,7 +355,11 @@ func TestResolveDefaultConnectCredentials(t *testing.T) { func TestResolveExplicitUserConnectCredentials(t *testing.T) { response := &mpgv1.GetManagedClusterResponse{ - Data: mpgv1.ManagedCluster{Id: "cluster-id"}, + Data: mpgv1.ManagedCluster{ + Id: "cluster-id", + Name: "test-cluster", + Status: "ready", + }, Credentials: mpgv1.GetManagedClusterCredentialsResponse{ Status: "initializing", DBName: "default-db", @@ -130,7 +379,7 @@ func TestResolveExplicitUserConnectCredentials(t *testing.T) { } ctx := mpgv1.NewContextWithClient(context.Background(), client) - credentials, err := resolveConnectCredentials(ctx, response, "app-user") + credentials, err := resolveConnectCredentials(ctx, response, true, "app-user") require.NoError(t, err) assert.Equal(t, "app-user", credentials.User) @@ -139,6 +388,7 @@ func TestResolveExplicitUserConnectCredentials(t *testing.T) { } func TestResolveExplicitUserConnectCredentialsErrors(t *testing.T) { + const name = "test-cluster" t.Run("empty password", func(t *testing.T) { client := &mock.MpgV1Client{ GetUserCredentialsFunc: func(context.Context, string, string) (mpgv1.GetUserCredentialsResponse, error) { @@ -147,7 +397,10 @@ func TestResolveExplicitUserConnectCredentialsErrors(t *testing.T) { } ctx := mpgv1.NewContextWithClient(context.Background(), client) - credentials, err := resolveConnectCredentials(ctx, &mpgv1.GetManagedClusterResponse{}, "app-user") + response := &mpgv1.GetManagedClusterResponse{ + Data: mpgv1.ManagedCluster{Name: name, Status: "ready"}, + } + credentials, err := resolveConnectCredentials(ctx, response, true, "app-user") require.EqualError(t, err, "error getting user password") assert.Nil(t, credentials) @@ -161,9 +414,209 @@ func TestResolveExplicitUserConnectCredentialsErrors(t *testing.T) { } ctx := mpgv1.NewContextWithClient(context.Background(), client) - credentials, err := resolveConnectCredentials(ctx, &mpgv1.GetManagedClusterResponse{}, "app-user") + response := &mpgv1.GetManagedClusterResponse{ + Data: mpgv1.ManagedCluster{Name: name, Status: "ready"}, + } + credentials, err := resolveConnectCredentials(ctx, response, true, "app-user") require.EqualError(t, err, "failed retrieving credentials for user app-user: request failed") assert.Nil(t, credentials) }) } + +func TestResolveConnectCredentialsPublic(t *testing.T) { + response := &mpgv1.GetManagedClusterResponse{ + Data: mpgv1.ManagedCluster{Id: "cluster-id", Name: "test-cluster", Status: "ready"}, + } + + t.Run("default user resolves fly-user from public API", func(t *testing.T) { + credCalls := 0 + ctx := flapsutil.NewContextWithClient(context.Background(), &mock.FlapsClient{ + GetManagedPostgresUserCredentialsFunc: func(_ context.Context, clusterID, username string) (flaps.ManagedPostgresUserCredentials, error) { + credCalls++ + require.Equal(t, "cluster-id", clusterID) + require.Equal(t, "fly-user", username) + + return flaps.ManagedPostgresUserCredentials{Username: "fly-user", Password: "p"}, nil + }, + }) + + credentials, err := resolveConnectCredentials(ctx, response, false, "") + require.NoError(t, err) + require.Equal(t, 1, credCalls) + require.Equal(t, "fly-user", credentials.User) + require.Equal(t, "p", credentials.Password) + require.Equal(t, "fly-db", credentials.DBName) + }) + + t.Run("explicit user passes flag value through", func(t *testing.T) { + credCalls := 0 + ctx := flapsutil.NewContextWithClient(context.Background(), &mock.FlapsClient{ + GetManagedPostgresUserCredentialsFunc: func(_ context.Context, clusterID, username string) (flaps.ManagedPostgresUserCredentials, error) { + credCalls++ + require.Equal(t, "cluster-id", clusterID) + require.Equal(t, "alice", username) + + return flaps.ManagedPostgresUserCredentials{Username: "alice", Password: "a"}, nil + }, + }) + + credentials, err := resolveConnectCredentials(ctx, response, false, "alice") + require.NoError(t, err) + require.Equal(t, 1, credCalls) + require.Equal(t, "alice", credentials.User) + require.Equal(t, "a", credentials.Password) + require.Equal(t, "fly-db", credentials.DBName) + require.Equal(t, "postgresql://alice:a@localhost:16380/fly-db", buildConnectURL(credentials, "", "16380")) + require.Equal(t, "postgresql://alice:a@localhost:16380/app-db", buildConnectURL(credentials, "app-db", "16380")) + }) + + t.Run("default user empty password mirrors legacy error", func(t *testing.T) { + ctx := flapsutil.NewContextWithClient(context.Background(), &mock.FlapsClient{ + GetManagedPostgresUserCredentialsFunc: func(context.Context, string, string) (flaps.ManagedPostgresUserCredentials, error) { + return flaps.ManagedPostgresUserCredentials{Username: "fly-user", Password: ""}, nil + }, + }) + + credentials, err := resolveConnectCredentials(ctx, response, false, "") + require.EqualError(t, err, "error getting cluster password") + assert.Nil(t, credentials) + }) + + t.Run("explicit user empty password returns user error", func(t *testing.T) { + ctx := flapsutil.NewContextWithClient(context.Background(), &mock.FlapsClient{ + GetManagedPostgresUserCredentialsFunc: func(context.Context, string, string) (flaps.ManagedPostgresUserCredentials, error) { + return flaps.ManagedPostgresUserCredentials{Username: "alice", Password: ""}, nil + }, + }) + + credentials, err := resolveConnectCredentials(ctx, response, false, "alice") + require.EqualError(t, err, "error getting user password") + assert.Nil(t, credentials) + }) + + t.Run("default user public credentials error propagates", func(t *testing.T) { + ctx := flapsutil.NewContextWithClient(context.Background(), &mock.FlapsClient{ + GetManagedPostgresUserCredentialsFunc: func(context.Context, string, string) (flaps.ManagedPostgresUserCredentials, error) { + return flaps.ManagedPostgresUserCredentials{}, errors.New("denied") + }, + }) + + credentials, err := resolveConnectCredentials(ctx, response, false, "") + require.EqualError(t, err, "failed retrieving credentials for user fly-user: denied") + assert.Nil(t, credentials) + }) + + t.Run("default user 404 preserves initializing error", func(t *testing.T) { + ctx := flapsutil.NewContextWithClient(context.Background(), &mock.FlapsClient{ + GetManagedPostgresUserCredentialsFunc: func(context.Context, string, string) (flaps.ManagedPostgresUserCredentials, error) { + return flaps.ManagedPostgresUserCredentials{}, fmt.Errorf("wrapped: %w", &flaps.FlapsError{ResponseStatusCode: 404, OriginalError: errors.New("not found")}) + }, + }) + + credentials, err := resolveConnectCredentials(ctx, response, false, "") + require.EqualError(t, err, "cluster is still initializing, wait a bit more") + assert.Nil(t, credentials) + }) + + t.Run("explicit user 404 remains a user error", func(t *testing.T) { + ctx := flapsutil.NewContextWithClient(context.Background(), &mock.FlapsClient{ + GetManagedPostgresUserCredentialsFunc: func(context.Context, string, string) (flaps.ManagedPostgresUserCredentials, error) { + return flaps.ManagedPostgresUserCredentials{}, fmt.Errorf("wrapped: %w", &flaps.FlapsError{ResponseStatusCode: 404, OriginalError: errors.New("missing user")}) + }, + }) + + credentials, err := resolveConnectCredentials(ctx, response, false, "alice") + require.EqualError(t, err, "failed retrieving credentials for user alice: wrapped: missing user") + assert.Nil(t, credentials) + }) +} + +func TestResolveConnectCredentialsPublicNonReady(t *testing.T) { + const name = "test-cluster" + tests := []struct { + status string + }{ + {"creating"}, + {"degraded"}, + } + + for _, tt := range tests { + t.Run(tt.status+"/default_user_public", func(t *testing.T) { + c := samplePublicCluster() + c.Status = tt.status + c.Name = name + response, _ := publicToLegacyClusterResponse(c) + + credCalls := 0 + ctx := flapsutil.NewContextWithClient(context.Background(), &mock.FlapsClient{ + GetManagedPostgresUserCredentialsFunc: func(context.Context, string, string) (flaps.ManagedPostgresUserCredentials, error) { + credCalls++ + + return flaps.ManagedPostgresUserCredentials{Username: "fly-user", Password: "p"}, nil + }, + }) + + credentials, err := resolveConnectCredentials(ctx, &response, false, "") + require.NoError(t, err, "non-ready cluster must proceed to credential resolution") + require.Equal(t, 1, credCalls, "must hit the credentials endpoint exactly once") + require.NotNil(t, credentials) + }) + + t.Run(tt.status+"/explicit_user_public", func(t *testing.T) { + c := samplePublicCluster() + c.Status = tt.status + c.Name = name + response, _ := publicToLegacyClusterResponse(c) + + credCalls := 0 + ctx := flapsutil.NewContextWithClient(context.Background(), &mock.FlapsClient{ + GetManagedPostgresUserCredentialsFunc: func(context.Context, string, string) (flaps.ManagedPostgresUserCredentials, error) { + credCalls++ + + return flaps.ManagedPostgresUserCredentials{Username: "alice", Password: "p"}, nil + }, + }) + + credentials, err := resolveConnectCredentials(ctx, &response, false, "alice") + require.NoError(t, err, "non-ready cluster must proceed to credential resolution") + require.Equal(t, 1, credCalls, "must hit the credentials endpoint exactly once") + require.NotNil(t, credentials) + }) + } +} + +func TestProxyParamsPublicPorts(t *testing.T) { + for _, port := range []int{1, 5433, 65535} { + t.Run(strconv.Itoa(port), func(t *testing.T) { + c := samplePublicCluster() + c.Endpoints.Primary.Direct.Port = port + response, advertisedPort := publicToLegacyClusterResponse(c) + + _, params, err := proxyParams(&response, advertisedPort, "0", "test-org", "127.0.0.1", nil) + require.NoError(t, err) + require.Equal(t, "10.0.0.1", params.RemoteHost) + require.Equal(t, []string{"0", strconv.Itoa(port)}, params.Ports) + }) + } +} + +func TestGetMpgProxyParamsRejectsInvalidPublicPort(t *testing.T) { + for _, port := range []int{0, -1, 65536} { + t.Run(strconv.Itoa(port), func(t *testing.T) { + c := samplePublicCluster() + c.Endpoints.Primary.Direct.Port = port + ctx := flag.NewContext(context.Background(), pflag.NewFlagSet("test", pflag.ContinueOnError)) + ctx = flapsutil.NewContextWithClient(ctx, &mock.FlapsClient{ + GetManagedPostgresClusterFunc: func(context.Context, string) (flaps.ManagedPostgresCluster, error) { + return c, nil + }, + }) + // No legacy or tunnel client: invalid public endpoints must stop before either. + cluster, params, err := GetMpgProxyParams(ctx, "0", "mpg-123", "test-org") + require.EqualError(t, err, fmt.Sprintf("invalid cluster port %d: must be between 1 and 65535", c.Endpoints.Primary.Direct.Port)) + require.Nil(t, cluster) + require.Nil(t, params) + }) + } +} diff --git a/internal/command/mpg/v2/run_connect.go b/internal/command/mpg/v2/run_connect.go index 6113a141ea..abaccbf91f 100644 --- a/internal/command/mpg/v2/run_connect.go +++ b/internal/command/mpg/v2/run_connect.go @@ -3,6 +3,7 @@ package cmdv2 import ( "context" "fmt" + "io" "os" "os/exec" "os/signal" @@ -82,9 +83,7 @@ func RunConnect(ctx context.Context, clusterID string, resolvedOrgSlug string, p return err } - if cluster.Status != "ready" { - fmt.Fprintf(io.ErrOut, "%s Cluster is not in ready state, currently: %s\n", aurora.Yellow("WARN"), cluster.Status) - } + maybeWarnNotReady(io.ErrOut, cluster) psqlPath, err := exec.LookPath("psql") if err != nil { @@ -103,15 +102,7 @@ func RunConnect(ctx context.Context, clusterID string, resolvedOrgSlug string, p return err } - user := credentials.User - password := credentials.Password - - // Use selected database or fall back to default from credentials - if db == "" { - db = credentials.DBName - } - - connectUrl := fmt.Sprintf("postgresql://%s:%s@localhost:%s/%s", user, password, localProxyPort, db) + connectUrl := buildConnectURL(credentials, db, localProxyPort) // Allow Ctrl+C signals to hit psql psqlCtx, psqlCancel := context.WithCancel(context.WithoutCancel(ctx)) @@ -174,3 +165,21 @@ func RunConnect(ctx context.Context, clusterID string, resolvedOrgSlug string, p return err } + +// buildConnectURL prefers the selected database over the credential default. +func buildConnectURL(credentials *mpgv2.GetClusterCredentialsResponse, db string, localProxyPort string) string { + if db == "" { + db = credentials.DBName + } + + return fmt.Sprintf("postgresql://%s:%s@localhost:%s/%s", credentials.User, credentials.Password, localProxyPort, db) +} + +// maybeWarnNotReady warns when a cluster is not in ready state. +func maybeWarnNotReady(errOut io.Writer, cluster *mpgv2.ManagedCluster) { + if cluster == nil || cluster.Status == "ready" { + return + } + + fmt.Fprintf(errOut, "%s Cluster is not in ready state, currently: %s\n", aurora.Yellow("WARN"), cluster.Status) +} diff --git a/internal/command/mpg/v2/run_connect_test.go b/internal/command/mpg/v2/run_connect_test.go new file mode 100644 index 0000000000..af7de717bb --- /dev/null +++ b/internal/command/mpg/v2/run_connect_test.go @@ -0,0 +1,40 @@ +package cmdv2 + +import ( + "bytes" + "testing" + + "github.com/stretchr/testify/require" + mpgv2 "github.com/superfly/flyctl/internal/uiex/mpg/v2" +) + +func TestMaybeWarnNotReady(t *testing.T) { + const name = "test-cluster" + tests := []struct { + name string + status string + wantWarn bool + }{ + {name: "ready silent", status: "ready", wantWarn: false}, + {name: "creating warns", status: "creating", wantWarn: true}, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + var buf bytes.Buffer + cluster := &mpgv2.ManagedCluster{Name: name, Status: tt.status} + maybeWarnNotReady(&buf, cluster) + + got := buf.String() + if !tt.wantWarn { + require.Empty(t, got, "no warning expected for status=%q", tt.status) + + return + } + // Check the warning text independently of ANSI color codes. + require.Contains(t, got, "WARN", "warning must contain the literal 'WARN' marker") + require.Contains(t, got, "Cluster is not in ready state, currently: "+tt.status) + require.True(t, bytes.HasSuffix(buf.Bytes(), []byte("\n")), "warning must end with a newline (pre-migration format)") + }) + } +} diff --git a/internal/command/mpg/v2/run_proxy.go b/internal/command/mpg/v2/run_proxy.go index bcff1c67fa..490e85e08c 100644 --- a/internal/command/mpg/v2/run_proxy.go +++ b/internal/command/mpg/v2/run_proxy.go @@ -2,11 +2,18 @@ package cmdv2 import ( "context" + "errors" "fmt" + "strconv" + fly "github.com/superfly/fly-go" + "github.com/superfly/fly-go/flaps" "github.com/superfly/flyctl/agent" "github.com/superfly/flyctl/internal/flag" + "github.com/superfly/flyctl/internal/flapsutil" "github.com/superfly/flyctl/internal/flyutil" + "github.com/superfly/flyctl/internal/mpgutil" + "github.com/superfly/flyctl/internal/uiex/mpg" mpgv2 "github.com/superfly/flyctl/internal/uiex/mpg/v2" "github.com/superfly/flyctl/proxy" ) @@ -29,12 +36,12 @@ func GetMpgProxyParams( clusterID string, resolvedOrgSlug string, ) (*mpgv2.ManagedCluster, *proxy.ConnectParams, error) { - response, err := getCluster(ctx, clusterID) + response, _, port, err := getCluster(ctx, clusterID) if err != nil { return nil, nil, err } - cluster, params, err := buildProxyParams(ctx, response, localProxyPort, resolvedOrgSlug) + cluster, params, err := buildProxyParams(ctx, response, port, localProxyPort, resolvedOrgSlug) if err != nil { return nil, nil, err } @@ -42,8 +49,7 @@ func GetMpgProxyParams( return cluster, params, nil } -// GetMpgConnectParams builds proxy connection parameters and resolves the -// database credentials needed by fly mpg connect. +// GetMpgConnectParams resolves credentials and proxy parameters. func GetMpgConnectParams( ctx context.Context, localProxyPort string, @@ -51,17 +57,17 @@ func GetMpgConnectParams( clusterID string, resolvedOrgSlug string, ) (*mpgv2.ManagedCluster, *proxy.ConnectParams, *mpgv2.GetClusterCredentialsResponse, error) { - response, err := getCluster(ctx, clusterID) + response, useLegacy, port, err := getCluster(ctx, clusterID) if err != nil { return nil, nil, nil, err } - credentials, err := resolveConnectCredentials(ctx, response, username) + credentials, err := resolveConnectCredentials(ctx, response, useLegacy, username) if err != nil { return nil, nil, nil, err } - cluster, params, err := buildProxyParams(ctx, response, localProxyPort, resolvedOrgSlug) + cluster, params, err := buildProxyParams(ctx, response, port, localProxyPort, resolvedOrgSlug) if err != nil { return nil, nil, nil, err } @@ -69,23 +75,62 @@ func GetMpgConnectParams( return cluster, params, credentials, nil } -func getCluster(ctx context.Context, clusterID string) (*mpgv2.GetClusterResponse, error) { - mpgClient := mpgv2.ClientFromContext(ctx) - response, err := mpgClient.GetClusterById(ctx, clusterID) +// getCluster tries the public API, falling back to the legacy client only on 404. +// It returns the credential source and direct endpoint port (5432 for legacy). +func getCluster(ctx context.Context, clusterID string) (*mpgv2.GetClusterResponse, bool, int, error) { + flapsClient := flapsutil.ClientFromContext(ctx) + publicCluster, err := flapsClient.GetManagedPostgresCluster(ctx, clusterID) + if err == nil { + response, port := publicToLegacyClusterResponse(publicCluster) + + return &response, false, port, nil + } + + if !errors.Is(err, flaps.ErrFlapsNotFound) { + return nil, false, 0, fmt.Errorf("failed retrieving cluster %s: %w", clusterID, err) + } + + legacyClient := mpgv2.ClientFromContext(ctx) + response, err := legacyClient.GetClusterById(ctx, clusterID) if err != nil { - return nil, fmt.Errorf("failed retrieving cluster %s: %w", clusterID, err) + return nil, true, 0, fmt.Errorf("failed retrieving cluster %s: %w", clusterID, err) } - return &response, nil + return &response, true, mpgutil.DefaultPort, nil +} + +// publicToLegacyClusterResponse adapts the public cluster to the legacy shape. +// The advertised direct port is returned unchanged for validation. +func publicToLegacyClusterResponse(c flaps.ManagedPostgresCluster) (mpgv2.GetClusterResponse, int) { + port := c.Endpoints.Primary.Direct.Port + + return mpgv2.GetClusterResponse{ + Data: mpgv2.ManagedCluster{ + Id: c.ID, + Name: c.Name, + Status: c.Status, + Region: c.Region, + Plan: c.Plan, + Disk: c.DiskSizeGB, + Replicas: c.Replicas, + Organization: fly.Organization{Name: c.Organization.Name, Slug: c.Organization.Slug}, + IpAssignments: mpg.ManagedClusterIpAssignments{Direct: c.Endpoints.Primary.Direct.Host}, + }, + }, port } +// resolveConnectCredentials uses the same API as the cluster lookup. +// Public credentials default to fly-user and fly-db. func resolveConnectCredentials( ctx context.Context, response *mpgv2.GetClusterResponse, + useLegacy bool, username string, ) (*mpgv2.GetClusterCredentialsResponse, error) { var credentials mpgv2.GetClusterCredentialsResponse - if username != "" { + + switch { + case username != "" && useLegacy: mpgClient := mpgv2.ClientFromContext(ctx) userCreds, err := mpgClient.GetUserCredentials(ctx, response.Data.Id, username) if err != nil { @@ -97,19 +142,56 @@ func resolveConnectCredentials( Password: userCreds.Data.Password, DBName: response.Credentials.DBName, } - } else { + case username != "": + flapsClient := flapsutil.ClientFromContext(ctx) + userCreds, err := flapsClient.GetManagedPostgresUserCredentials(ctx, response.Data.Id, username) + if err != nil { + return nil, fmt.Errorf("failed retrieving credentials for user %s: %w", username, err) + } + + credentials = mpgv2.GetClusterCredentialsResponse{ + User: userCreds.Username, + Password: userCreds.Password, + DBName: mpgutil.DefaultDatabase, + } + case useLegacy: credentials = response.Credentials - } + default: + flapsClient := flapsutil.ClientFromContext(ctx) + userCreds, err := flapsClient.GetManagedPostgresUserCredentials(ctx, response.Data.Id, mpgutil.DefaultUsername) + if err != nil { + if errors.Is(err, flaps.ErrFlapsNotFound) { + return nil, fmt.Errorf("cluster is still initializing, wait a bit more") + } - if username == "" { - if credentials.Status == "initializing" { - return nil, fmt.Errorf("cluster is still initializing, wait a bit more") + return nil, fmt.Errorf("failed retrieving credentials for user %s: %w", mpgutil.DefaultUsername, err) } - if credentials.Status == "error" || credentials.Password == "" { - return nil, fmt.Errorf("error getting cluster password") + credentials = mpgv2.GetClusterCredentialsResponse{ + User: userCreds.Username, + Password: userCreds.Password, + DBName: mpgutil.DefaultDatabase, + } + } + + if useLegacy { + // Only legacy default-user credentials include a status. + if username == "" { + if credentials.Status == "initializing" { + return nil, fmt.Errorf("cluster is still initializing, wait a bit more") + } + + if credentials.Status == "error" || credentials.Password == "" { + return nil, fmt.Errorf("error getting cluster password") + } + } else if credentials.Password == "" { + return nil, fmt.Errorf("error getting user password") } } else if credentials.Password == "" { + if username == "" { + return nil, fmt.Errorf("error getting cluster password") + } + return nil, fmt.Errorf("error getting user password") } @@ -119,10 +201,11 @@ func resolveConnectCredentials( func buildProxyParams( ctx context.Context, response *mpgv2.GetClusterResponse, + port int, localProxyPort string, resolvedOrgSlug string, ) (*mpgv2.ManagedCluster, *proxy.ConnectParams, error) { - cluster, params, err := proxyParams(response, localProxyPort, resolvedOrgSlug, flag.GetBindAddr(ctx), nil) + cluster, params, err := proxyParams(response, port, localProxyPort, resolvedOrgSlug, flag.GetBindAddr(ctx), nil) if err != nil { return nil, nil, err } @@ -147,6 +230,7 @@ func buildProxyParams( func proxyParams( response *mpgv2.GetClusterResponse, + port int, localProxyPort string, resolvedOrgSlug string, bindAddr string, @@ -157,8 +241,12 @@ func proxyParams( return nil, nil, fmt.Errorf("error getting cluster IP") } + if port < 1 || port > 65535 { + return nil, nil, fmt.Errorf("invalid cluster port %d: must be between 1 and 65535", port) + } + return cluster, &proxy.ConnectParams{ - Ports: []string{localProxyPort, "5432"}, + Ports: []string{localProxyPort, strconv.Itoa(port)}, OrganizationSlug: resolvedOrgSlug, Dialer: dialer, BindAddr: bindAddr, diff --git a/internal/command/mpg/v2/run_proxy_test.go b/internal/command/mpg/v2/run_proxy_test.go index 90134aef15..ead8ecdf2d 100644 --- a/internal/command/mpg/v2/run_proxy_test.go +++ b/internal/command/mpg/v2/run_proxy_test.go @@ -3,12 +3,20 @@ package cmdv2 import ( "context" "errors" + "fmt" "net" + "strconv" "testing" + "github.com/spf13/pflag" "github.com/stretchr/testify/assert" "github.com/stretchr/testify/require" + fly "github.com/superfly/fly-go" + "github.com/superfly/fly-go/flaps" + "github.com/superfly/flyctl/internal/flag" + "github.com/superfly/flyctl/internal/flapsutil" "github.com/superfly/flyctl/internal/mock" + "github.com/superfly/flyctl/internal/mpgutil" "github.com/superfly/flyctl/internal/uiex/mpg" mpgv2 "github.com/superfly/flyctl/internal/uiex/mpg/v2" "github.com/superfly/flyctl/wg" @@ -53,7 +61,7 @@ func TestProxyParamsIgnoreCredentials(t *testing.T) { } dialer := &testDialer{} - cluster, params, err := proxyParams(&response, "15432", "test-org", "127.0.0.2", dialer) + cluster, params, err := proxyParams(&response, mpgutil.DefaultPort, "15432", "test-org", "127.0.0.2", dialer) require.NoError(t, err) assert.Same(t, &response.Data, cluster) assert.Equal(t, []string{"15432", "5432"}, params.Ports) @@ -66,42 +74,279 @@ func TestProxyParamsIgnoreCredentials(t *testing.T) { } func TestProxyParamsRequireDirectIP(t *testing.T) { - cluster, params, err := proxyParams(&mpgv2.GetClusterResponse{}, "15432", "test-org", "127.0.0.1", nil) + cluster, params, err := proxyParams(&mpgv2.GetClusterResponse{}, 0, "15432", "test-org", "127.0.0.1", nil) require.EqualError(t, err, "error getting cluster IP") assert.Nil(t, cluster) assert.Nil(t, params) } +func sampleProxyPublicCluster() flaps.ManagedPostgresCluster { + return flaps.ManagedPostgresCluster{ + ID: "mpg-123", + Name: "test-cluster", + Status: "ready", + Region: "ord", + Plan: "development", + DiskSizeGB: 10, + Replicas: 1, + Organization: flaps.ManagedPostgresOrganization{Name: "Test Org", Slug: "test-org"}, + Endpoints: flaps.ManagedPostgresEndpoints{ + Primary: struct { + Direct flaps.ManagedPostgresEndpoint `json:"direct"` + Pooler flaps.ManagedPostgresEndpoint `json:"pooler"` + }{ + Direct: flaps.ManagedPostgresEndpoint{Host: "10.0.0.1", Port: 5432}, + Pooler: flaps.ManagedPostgresEndpoint{Host: "10.0.0.2", Port: 6432}, + }, + }, + } +} + +func sampleProxyLegacyCluster() mpgv2.GetClusterResponse { + return mpgv2.GetClusterResponse{ + Data: mpgv2.ManagedCluster{ + Id: "mpg-123", Name: "test-cluster", Region: "ord", Status: "ready", + Plan: "development", Disk: 10, Replicas: 1, + Organization: fly.Organization{Name: "Test Org", Slug: "test-org"}, + IpAssignments: mpg.ManagedClusterIpAssignments{Direct: "10.0.0.1"}, + }, + Credentials: mpgv2.GetClusterCredentialsResponse{ + Status: "ready", User: "app", Password: "secret", + DBName: "app", ConnectionUri: "postgres://app:secret@10.0.0.1:5432/app", + }, + } +} + +func TestGetCluster(t *testing.T) { + tests := []struct { + name string + publicCluster flaps.ManagedPostgresCluster + publicErr error + legacyResponse mpgv2.GetClusterResponse + legacyErr error + wantUseLegacy bool + wantDirectHost string + wantStatus string + wantErr string + wantLegacyCalls int + }{ + { + name: "public success maps to legacy shape", + publicCluster: sampleProxyPublicCluster(), + wantUseLegacy: false, + wantDirectHost: "10.0.0.1", + wantStatus: "ready", + wantLegacyCalls: 0, + }, + { + name: "public success with empty host preserves empty", + publicCluster: func() flaps.ManagedPostgresCluster { + c := sampleProxyPublicCluster() + c.Endpoints.Primary.Direct.Host = "" + + return c + }(), + wantUseLegacy: false, + wantDirectHost: "", + wantStatus: "ready", + wantLegacyCalls: 0, + }, + { + name: "classified 404 falls back to legacy", + publicCluster: flaps.ManagedPostgresCluster{}, + publicErr: fmt.Errorf("wrapped: %w", &flaps.FlapsError{ + ResponseStatusCode: 404, + OriginalError: errors.New("not found"), + }), + legacyResponse: sampleProxyLegacyCluster(), + wantUseLegacy: true, + wantDirectHost: "10.0.0.1", + wantStatus: "ready", + wantLegacyCalls: 1, + }, + { + name: "classified 404 with legacy failure propagates legacy error", + publicCluster: flaps.ManagedPostgresCluster{}, + publicErr: fmt.Errorf("wrapped: %w", &flaps.FlapsError{ + ResponseStatusCode: 404, + OriginalError: errors.New("not found"), + }), + legacyResponse: mpgv2.GetClusterResponse{}, + legacyErr: errors.New("legacy denied"), + wantErr: "failed retrieving cluster mpg-123: legacy denied", + wantLegacyCalls: 1, + }, + { + name: "403 public error returns without fallback", + publicCluster: flaps.ManagedPostgresCluster{}, + publicErr: &flaps.FlapsError{ResponseStatusCode: 403, OriginalError: errors.New("denied")}, + wantErr: "failed retrieving cluster mpg-123: denied", + wantLegacyCalls: 0, + }, + { + name: "410 public error returns without fallback", + publicCluster: flaps.ManagedPostgresCluster{}, + publicErr: &flaps.FlapsError{ResponseStatusCode: 410, OriginalError: errors.New("gone")}, + wantErr: "failed retrieving cluster mpg-123: gone", + wantLegacyCalls: 0, + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + publicCalls, legacyCalls := 0, 0 + ctx := flapsutil.NewContextWithClient(context.Background(), &mock.FlapsClient{ + GetManagedPostgresClusterFunc: func(_ context.Context, id string) (flaps.ManagedPostgresCluster, error) { + publicCalls++ + require.Equal(t, "mpg-123", id) + + return tt.publicCluster, tt.publicErr + }, + }) + ctx = mpgv2.NewContextWithClient(ctx, &mock.MpgV2Client{ + GetClusterByIdFunc: func(_ context.Context, id string) (mpgv2.GetClusterResponse, error) { + legacyCalls++ + require.Equal(t, "mpg-123", id) + + return tt.legacyResponse, tt.legacyErr + }, + }) + + got, useLegacy, port, err := getCluster(ctx, "mpg-123") + if tt.wantErr != "" { + require.EqualError(t, err, tt.wantErr) + require.Nil(t, got) + require.Equal(t, 1, publicCalls) + require.Equal(t, tt.wantLegacyCalls, legacyCalls) + + return + } + require.NoError(t, err) + require.Equal(t, tt.wantUseLegacy, useLegacy) + require.NotNil(t, got) + require.Equal(t, tt.wantDirectHost, got.Data.IpAssignments.Direct) + require.Equal(t, tt.wantStatus, got.Data.Status) + require.Equal(t, 1, publicCalls) + require.Equal(t, tt.wantLegacyCalls, legacyCalls) + require.Equal(t, 5432, port) + if tt.wantUseLegacy { + _, params, err := proxyParams(got, port, "16380", "test-org", "127.0.0.1", nil) + require.NoError(t, err) + require.Equal(t, []string{"16380", "5432"}, params.Ports) + } + }) + } +} + +func TestGetMpgProxyParamsPublicNeverTouchesCredentials(t *testing.T) { + credCalls, legacyCredCalls := 0, 0 + ctx := flag.NewContext(context.Background(), pflag.NewFlagSet("test", pflag.ContinueOnError)) + ctx = flapsutil.NewContextWithClient(ctx, &mock.FlapsClient{ + GetManagedPostgresClusterFunc: func(context.Context, string) (flaps.ManagedPostgresCluster, error) { + cluster := sampleProxyPublicCluster() + cluster.Status = "creating" + cluster.Endpoints.Primary.Direct.Host = "" + + return cluster, nil + }, + GetManagedPostgresUserCredentialsFunc: func(context.Context, string, string) (flaps.ManagedPostgresUserCredentials, error) { + credCalls++ + + return flaps.ManagedPostgresUserCredentials{}, nil + }, + }) + ctx = mpgv2.NewContextWithClient(ctx, &mock.MpgV2Client{ + GetUserCredentialsFunc: func(context.Context, string, string) (mpgv2.GetUserCredentialsResponse, error) { + legacyCredCalls++ + + return mpgv2.GetUserCredentialsResponse{}, nil + }, + }) + + cluster, params, err := GetMpgProxyParams(ctx, "15432", "mpg-123", "test-org") + require.EqualError(t, err, "error getting cluster IP") + require.Nil(t, cluster) + require.Nil(t, params) + require.Zero(t, credCalls) + require.Zero(t, legacyCredCalls) +} + +func TestGetMpgConnectParamsResolvesCredentialsBeforeTunnel(t *testing.T) { + ctx := flapsutil.NewContextWithClient(context.Background(), &mock.FlapsClient{ + GetManagedPostgresClusterFunc: func(context.Context, string) (flaps.ManagedPostgresCluster, error) { + return sampleProxyPublicCluster(), nil + }, + GetManagedPostgresUserCredentialsFunc: func(context.Context, string, string) (flaps.ManagedPostgresUserCredentials, error) { + return flaps.ManagedPostgresUserCredentials{Username: mpgutil.DefaultUsername}, nil + }, + }) + + cluster, params, credentials, err := GetMpgConnectParams(ctx, "15432", "", "mpg-123", "test-org") + require.EqualError(t, err, "error getting cluster password") + require.Nil(t, cluster) + require.Nil(t, params) + require.Nil(t, credentials) +} + func TestResolveDefaultConnectCredentials(t *testing.T) { + // Legacy readiness comes from the credential envelope, not cluster status. + const name = "test-cluster" tests := []struct { - name string - credentials mpgv2.GetClusterCredentialsResponse - err string + name string + clusterSt string // response.Data.Status — cluster status. NOT consulted on the legacy path. + credStatus string // credentials.Status — the legacy envelope's own status field (post-fetch). + credPwd string // credentials.Password — empty-password fallback when status checks pass. + err string }{ { - name: "initializing", - credentials: mpgv2.GetClusterCredentialsResponse{Status: "initializing"}, - err: "cluster is still initializing, wait a bit more", + name: "empty password", + clusterSt: "ready", + credStatus: "ready", + credPwd: "", + err: "error getting cluster password", + }, + { + name: "ready cluster with stale initializing credentials refuses", + clusterSt: "ready", + credStatus: "initializing", + credPwd: "secret", + err: "cluster is still initializing, wait a bit more", }, { - name: "error status", - credentials: mpgv2.GetClusterCredentialsResponse{Status: "error", Password: "password"}, - err: "error getting cluster password", + name: "ready cluster with stale error credentials refuses", + clusterSt: "ready", + credStatus: "error", + credPwd: "secret", + err: "error getting cluster password", }, { - name: "empty password", - credentials: mpgv2.GetClusterCredentialsResponse{Status: "ready"}, - err: "error getting cluster password", + name: "creating cluster with ready credentials proceeds", + clusterSt: "creating", + credStatus: "ready", + credPwd: "secret", + err: "", // no error }, } for _, tt := range tests { t.Run(tt.name, func(t *testing.T) { - response := &mpgv2.GetClusterResponse{Credentials: tt.credentials} + response := &mpgv2.GetClusterResponse{ + Data: mpgv2.ManagedCluster{Name: name, Status: tt.clusterSt}, + Credentials: mpgv2.GetClusterCredentialsResponse{ + Status: tt.credStatus, + Password: tt.credPwd, + }, + } + + credentials, err := resolveConnectCredentials(context.Background(), response, true, "") - credentials, err := resolveConnectCredentials(context.Background(), response, "") + if tt.err == "" { + require.NoError(t, err) + require.NotNil(t, credentials) + return + } require.EqualError(t, err, tt.err) assert.Nil(t, credentials) }) @@ -110,7 +355,11 @@ func TestResolveDefaultConnectCredentials(t *testing.T) { func TestResolveExplicitUserConnectCredentials(t *testing.T) { response := &mpgv2.GetClusterResponse{ - Data: mpgv2.ManagedCluster{Id: "cluster-id"}, + Data: mpgv2.ManagedCluster{ + Id: "cluster-id", + Name: "test-cluster", + Status: "ready", + }, Credentials: mpgv2.GetClusterCredentialsResponse{ Status: "initializing", DBName: "default-db", @@ -130,7 +379,7 @@ func TestResolveExplicitUserConnectCredentials(t *testing.T) { } ctx := mpgv2.NewContextWithClient(context.Background(), client) - credentials, err := resolveConnectCredentials(ctx, response, "app-user") + credentials, err := resolveConnectCredentials(ctx, response, true, "app-user") require.NoError(t, err) assert.Equal(t, "app-user", credentials.User) @@ -139,6 +388,7 @@ func TestResolveExplicitUserConnectCredentials(t *testing.T) { } func TestResolveExplicitUserConnectCredentialsErrors(t *testing.T) { + const name = "test-cluster" t.Run("empty password", func(t *testing.T) { client := &mock.MpgV2Client{ GetUserCredentialsFunc: func(context.Context, string, string) (mpgv2.GetUserCredentialsResponse, error) { @@ -147,7 +397,10 @@ func TestResolveExplicitUserConnectCredentialsErrors(t *testing.T) { } ctx := mpgv2.NewContextWithClient(context.Background(), client) - credentials, err := resolveConnectCredentials(ctx, &mpgv2.GetClusterResponse{}, "app-user") + response := &mpgv2.GetClusterResponse{ + Data: mpgv2.ManagedCluster{Name: name, Status: "ready"}, + } + credentials, err := resolveConnectCredentials(ctx, response, true, "app-user") require.EqualError(t, err, "error getting user password") assert.Nil(t, credentials) @@ -161,9 +414,209 @@ func TestResolveExplicitUserConnectCredentialsErrors(t *testing.T) { } ctx := mpgv2.NewContextWithClient(context.Background(), client) - credentials, err := resolveConnectCredentials(ctx, &mpgv2.GetClusterResponse{}, "app-user") + response := &mpgv2.GetClusterResponse{ + Data: mpgv2.ManagedCluster{Name: name, Status: "ready"}, + } + credentials, err := resolveConnectCredentials(ctx, response, true, "app-user") require.EqualError(t, err, "failed retrieving credentials for user app-user: request failed") assert.Nil(t, credentials) }) } + +func TestResolveConnectCredentialsPublic(t *testing.T) { + response := &mpgv2.GetClusterResponse{ + Data: mpgv2.ManagedCluster{Id: "cluster-id", Name: "test-cluster", Status: "ready"}, + } + + t.Run("default user resolves fly-user from public API", func(t *testing.T) { + credCalls := 0 + ctx := flapsutil.NewContextWithClient(context.Background(), &mock.FlapsClient{ + GetManagedPostgresUserCredentialsFunc: func(_ context.Context, clusterID, username string) (flaps.ManagedPostgresUserCredentials, error) { + credCalls++ + require.Equal(t, "cluster-id", clusterID) + require.Equal(t, "fly-user", username) + + return flaps.ManagedPostgresUserCredentials{Username: "fly-user", Password: "p"}, nil + }, + }) + + credentials, err := resolveConnectCredentials(ctx, response, false, "") + require.NoError(t, err) + require.Equal(t, 1, credCalls) + require.Equal(t, "fly-user", credentials.User) + require.Equal(t, "p", credentials.Password) + require.Equal(t, "fly-db", credentials.DBName) + }) + + t.Run("explicit user passes flag value through", func(t *testing.T) { + credCalls := 0 + ctx := flapsutil.NewContextWithClient(context.Background(), &mock.FlapsClient{ + GetManagedPostgresUserCredentialsFunc: func(_ context.Context, clusterID, username string) (flaps.ManagedPostgresUserCredentials, error) { + credCalls++ + require.Equal(t, "cluster-id", clusterID) + require.Equal(t, "alice", username) + + return flaps.ManagedPostgresUserCredentials{Username: "alice", Password: "a"}, nil + }, + }) + + credentials, err := resolveConnectCredentials(ctx, response, false, "alice") + require.NoError(t, err) + require.Equal(t, 1, credCalls) + require.Equal(t, "alice", credentials.User) + require.Equal(t, "a", credentials.Password) + require.Equal(t, "fly-db", credentials.DBName) + require.Equal(t, "postgresql://alice:a@localhost:16380/fly-db", buildConnectURL(credentials, "", "16380")) + require.Equal(t, "postgresql://alice:a@localhost:16380/app-db", buildConnectURL(credentials, "app-db", "16380")) + }) + + t.Run("default user empty password mirrors legacy error", func(t *testing.T) { + ctx := flapsutil.NewContextWithClient(context.Background(), &mock.FlapsClient{ + GetManagedPostgresUserCredentialsFunc: func(context.Context, string, string) (flaps.ManagedPostgresUserCredentials, error) { + return flaps.ManagedPostgresUserCredentials{Username: "fly-user", Password: ""}, nil + }, + }) + + credentials, err := resolveConnectCredentials(ctx, response, false, "") + require.EqualError(t, err, "error getting cluster password") + assert.Nil(t, credentials) + }) + + t.Run("explicit user empty password returns user error", func(t *testing.T) { + ctx := flapsutil.NewContextWithClient(context.Background(), &mock.FlapsClient{ + GetManagedPostgresUserCredentialsFunc: func(context.Context, string, string) (flaps.ManagedPostgresUserCredentials, error) { + return flaps.ManagedPostgresUserCredentials{Username: "alice", Password: ""}, nil + }, + }) + + credentials, err := resolveConnectCredentials(ctx, response, false, "alice") + require.EqualError(t, err, "error getting user password") + assert.Nil(t, credentials) + }) + + t.Run("default user public credentials error propagates", func(t *testing.T) { + ctx := flapsutil.NewContextWithClient(context.Background(), &mock.FlapsClient{ + GetManagedPostgresUserCredentialsFunc: func(context.Context, string, string) (flaps.ManagedPostgresUserCredentials, error) { + return flaps.ManagedPostgresUserCredentials{}, errors.New("denied") + }, + }) + + credentials, err := resolveConnectCredentials(ctx, response, false, "") + require.EqualError(t, err, "failed retrieving credentials for user fly-user: denied") + assert.Nil(t, credentials) + }) + + t.Run("default user 404 preserves initializing error", func(t *testing.T) { + ctx := flapsutil.NewContextWithClient(context.Background(), &mock.FlapsClient{ + GetManagedPostgresUserCredentialsFunc: func(context.Context, string, string) (flaps.ManagedPostgresUserCredentials, error) { + return flaps.ManagedPostgresUserCredentials{}, fmt.Errorf("wrapped: %w", &flaps.FlapsError{ResponseStatusCode: 404, OriginalError: errors.New("not found")}) + }, + }) + + credentials, err := resolveConnectCredentials(ctx, response, false, "") + require.EqualError(t, err, "cluster is still initializing, wait a bit more") + assert.Nil(t, credentials) + }) + + t.Run("explicit user 404 remains a user error", func(t *testing.T) { + ctx := flapsutil.NewContextWithClient(context.Background(), &mock.FlapsClient{ + GetManagedPostgresUserCredentialsFunc: func(context.Context, string, string) (flaps.ManagedPostgresUserCredentials, error) { + return flaps.ManagedPostgresUserCredentials{}, fmt.Errorf("wrapped: %w", &flaps.FlapsError{ResponseStatusCode: 404, OriginalError: errors.New("missing user")}) + }, + }) + + credentials, err := resolveConnectCredentials(ctx, response, false, "alice") + require.EqualError(t, err, "failed retrieving credentials for user alice: wrapped: missing user") + assert.Nil(t, credentials) + }) +} + +func TestResolveConnectCredentialsPublicNonReady(t *testing.T) { + const name = "test-cluster" + tests := []struct { + status string + }{ + {"creating"}, + {"degraded"}, + } + + for _, tt := range tests { + t.Run(tt.status+"/default_user_public", func(t *testing.T) { + c := sampleProxyPublicCluster() + c.Status = tt.status + c.Name = name + response, _ := publicToLegacyClusterResponse(c) + + credCalls := 0 + ctx := flapsutil.NewContextWithClient(context.Background(), &mock.FlapsClient{ + GetManagedPostgresUserCredentialsFunc: func(context.Context, string, string) (flaps.ManagedPostgresUserCredentials, error) { + credCalls++ + + return flaps.ManagedPostgresUserCredentials{Username: "fly-user", Password: "p"}, nil + }, + }) + + credentials, err := resolveConnectCredentials(ctx, &response, false, "") + require.NoError(t, err, "non-ready cluster must proceed to credential resolution") + require.Equal(t, 1, credCalls, "must hit the credentials endpoint exactly once") + require.NotNil(t, credentials) + }) + + t.Run(tt.status+"/explicit_user_public", func(t *testing.T) { + c := sampleProxyPublicCluster() + c.Status = tt.status + c.Name = name + response, _ := publicToLegacyClusterResponse(c) + + credCalls := 0 + ctx := flapsutil.NewContextWithClient(context.Background(), &mock.FlapsClient{ + GetManagedPostgresUserCredentialsFunc: func(context.Context, string, string) (flaps.ManagedPostgresUserCredentials, error) { + credCalls++ + + return flaps.ManagedPostgresUserCredentials{Username: "alice", Password: "p"}, nil + }, + }) + + credentials, err := resolveConnectCredentials(ctx, &response, false, "alice") + require.NoError(t, err, "non-ready cluster must proceed to credential resolution") + require.Equal(t, 1, credCalls, "must hit the credentials endpoint exactly once") + require.NotNil(t, credentials) + }) + } +} + +func TestProxyParamsPublicPorts(t *testing.T) { + for _, port := range []int{1, 5433, 65535} { + t.Run(strconv.Itoa(port), func(t *testing.T) { + c := sampleProxyPublicCluster() + c.Endpoints.Primary.Direct.Port = port + response, advertisedPort := publicToLegacyClusterResponse(c) + + _, params, err := proxyParams(&response, advertisedPort, "0", "test-org", "127.0.0.1", nil) + require.NoError(t, err) + require.Equal(t, "10.0.0.1", params.RemoteHost) + require.Equal(t, []string{"0", strconv.Itoa(port)}, params.Ports) + }) + } +} + +func TestGetMpgProxyParamsRejectsInvalidPublicPort(t *testing.T) { + for _, port := range []int{0, -1, 65536} { + t.Run(strconv.Itoa(port), func(t *testing.T) { + c := sampleProxyPublicCluster() + c.Endpoints.Primary.Direct.Port = port + ctx := flag.NewContext(context.Background(), pflag.NewFlagSet("test", pflag.ContinueOnError)) + ctx = flapsutil.NewContextWithClient(ctx, &mock.FlapsClient{ + GetManagedPostgresClusterFunc: func(context.Context, string) (flaps.ManagedPostgresCluster, error) { + return c, nil + }, + }) + // No legacy or tunnel client: invalid public endpoints must stop before either. + cluster, params, err := GetMpgProxyParams(ctx, "0", "mpg-123", "test-org") + require.EqualError(t, err, fmt.Sprintf("invalid cluster port %d: must be between 1 and 65535", c.Endpoints.Primary.Direct.Port)) + require.Nil(t, cluster) + require.Nil(t, params) + }) + } +}