mpg: route connect and proxy through the public Machines API

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)
+}

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)")
+		})
+	}
+}

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,63 +36,101 @@ 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
 	}
 
 	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,
 	username string,
 	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
 	}
 
 	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,