Merge 645cdee into b657dad

internal/command/command.go

@@ -19,6 +19,8 @@ import (
 	"github.com/superfly/flyctl/internal/command/auth/webauth"
 	"github.com/superfly/flyctl/internal/flyutil"
 	"github.com/superfly/flyctl/internal/prompt"
+	"github.com/superfly/flyctl/internal/uiex"
+	"github.com/superfly/flyctl/internal/uiexutil"
 	"github.com/superfly/flyctl/iostreams"
 
 	"github.com/superfly/flyctl/internal/appconfig"
@@ -603,6 +605,24 @@ func RequireSession(ctx context.Context) (context.Context, error) {
 	return ctx, nil
 }
 
+// Apply uiex client to uiex
+func RequireUiex(ctx context.Context) (context.Context, error) {
+	cfg := config.FromContext(ctx)
+
+	if uiexutil.ClientFromContext(ctx) == nil {
+		client, err := uiexutil.NewClientWithOptions(ctx, uiex.NewClientOpts{
+			Logger: logger.FromContext(ctx),
+			Tokens: cfg.Tokens,
+		})
+		if err != nil {
+			return nil, err
+		}
+		ctx = uiexutil.NewContextWithClient(ctx, client)
+	}
+
+	return ctx, nil
+}
+
 func tryOpenUserURL(ctx context.Context, url string) error {
 	io := iostreams.FromContext(ctx)
 

internal/command/mpg/connect.go

@@ -0,0 +1,71 @@
+package mpg
+
+import (
+	"context"
+	"fmt"
+	"os/exec"
+
+	"github.com/spf13/cobra"
+
+	"github.com/superfly/flyctl/iostreams"
+	"github.com/superfly/flyctl/proxy"
+
+	"github.com/superfly/flyctl/internal/command"
+	"github.com/superfly/flyctl/internal/command/orgs"
+	"github.com/superfly/flyctl/internal/flag"
+)
+
+func newConnect() (cmd *cobra.Command) {
+	const (
+		long = `Connect to a Redis database using psql`
+
+		short = long
+		usage = "connect"
+	)
+
+	cmd = command.New(usage, short, long, runConnect, command.RequireSession, command.RequireUiex)
+
+	flag.Add(cmd,
+		flag.Org(),
+	)
+
+	return cmd
+}
+
+func runConnect(ctx context.Context) (err error) {
+	io := iostreams.FromContext(ctx)
+	org, err := orgs.OrgFromFlagOrSelect(ctx)
+	if err != nil {
+		return err
+	}
+
+	localProxyPort := "16380"
+
+	cluster, params, password, err := getMpgProxyParams(ctx, org.Slug, localProxyPort)
+	if err != nil {
+		return err
+	}
+
+	psqlPath, err := exec.LookPath("psql")
+	if err != nil {
+		fmt.Fprintf(io.Out, "Could not find psql in your $PATH. Install it or point your psql at: %s", "someurl")
+		return
+	}
+
+	err = proxy.Start(ctx, params)
+	if err != nil {
+		return err
+	}
+
+	name := fmt.Sprintf("pgdb-%s", cluster.Id)
+	connectUrl := fmt.Sprintf("postgres://%s:%s@localhost:%s/%s?sslmode=require", name, password, localProxyPort, name)
+	cmd := exec.CommandContext(ctx, psqlPath, connectUrl)
+	cmd.Stdout = io.Out
+	cmd.Stderr = io.ErrOut
+	cmd.Stdin = io.In
+
+	cmd.Start()
+	cmd.Wait()
+
+	return
+}

internal/command/mpg/mpg.go

@@ -0,0 +1,25 @@
+package mpg
+
+import (
+	"github.com/spf13/cobra"
+	"github.com/superfly/flyctl/internal/command"
+)
+
+func New() *cobra.Command {
+	const (
+		short = `Manage Managed Postgres clusters.`
+
+		long = short + "\n"
+	)
+
+	cmd := command.New("managed-postgres", short, long, nil)
+
+	cmd.Aliases = []string{"mpg"}
+
+	cmd.AddCommand(
+		newProxy(),
+		newConnect(),
+	)
+
+	return cmd
+}

internal/command/mpg/proxy.go

@@ -0,0 +1,108 @@
+package mpg
+
+import (
+	"context"
+	"fmt"
+
+	"github.com/spf13/cobra"
+	"github.com/superfly/flyctl/agent"
+	"github.com/superfly/flyctl/internal/command"
+	"github.com/superfly/flyctl/internal/command/orgs"
+	"github.com/superfly/flyctl/internal/flag"
+	"github.com/superfly/flyctl/internal/flyutil"
+	"github.com/superfly/flyctl/internal/prompt"
+	"github.com/superfly/flyctl/internal/uiex"
+	"github.com/superfly/flyctl/internal/uiexutil"
+	"github.com/superfly/flyctl/proxy"
+	"github.com/superfly/flyctl/terminal"
+)
+
+func newProxy() (cmd *cobra.Command) {
+	const (
+		long = `Proxy to a MPG database`
+
+		short = long
+		usage = "proxy"
+	)
+
+	cmd = command.New(usage, short, long, runProxy, command.RequireSession, command.RequireUiex)
+
+	flag.Add(cmd,
+		flag.Org(),
+		flag.Region(),
+	)
+
+	return cmd
+}
+
+func runProxy(ctx context.Context) (err error) {
+	org, err := orgs.OrgFromFlagOrSelect(ctx)
+	if err != nil {
+		return err
+	}
+
+	localProxyPort := "16380"
+	_, params, password, err := getMpgProxyParams(ctx, org.Slug, localProxyPort)
+	if err != nil {
+		return err
+	}
+
+	terminal.Infof("Proxying postgres to port \"%s\" with password \"%s\"", localProxyPort, password)
+
+	return proxy.Connect(ctx, params)
+}
+
+func getMpgProxyParams(ctx context.Context, orgSlug string, localProxyPort string) (*uiex.ManagedCluster, *proxy.ConnectParams, string, error) {
+	client := flyutil.ClientFromContext(ctx)
+	uiexClient := uiexutil.ClientFromContext(ctx)
+
+	var index int
+	var options []string
+
+	clustersResponse, err := uiexClient.ListManagedClusters(ctx, orgSlug)
+	if err != nil {
+		return nil, nil, "", err
+	}
+
+	for _, cluster := range clustersResponse.Data {
+		options = append(options, fmt.Sprintf("%s (%s)", cluster.Name, cluster.Region))
+	}
+
+	selectErr := prompt.Select(ctx, &index, "Select a database to connect to", "", options...)
+	if selectErr != nil {
+		return nil, nil, "", selectErr
+	}
+
+	selectedCluster := clustersResponse.Data[index]
+
+	response, err := uiexClient.GetManagedCluster(ctx, selectedCluster.Organization.Slug, selectedCluster.Id)
+	if err != nil {
+		return nil, nil, "", err
+	}
+	cluster := response.Data
+
+	if response.Password.Status == "initializing" {
+		return nil, nil, "", fmt.Errorf("Cluster is still initializing, wait a bit more")
+	}
+
+	if response.Password.Status == "error" {
+		return nil, nil, "", fmt.Errorf("Error getting cluster password")
+	}
+
+	agentclient, err := agent.Establish(ctx, client)
+	if err != nil {
+		return nil, nil, "", err
+	}
+
+	dialer, err := agentclient.ConnectToTunnel(ctx, orgSlug, "", false)
+	if err != nil {
+		return nil, nil, "", err
+	}
+
+	return &cluster, &proxy.ConnectParams{
+		Ports:            []string{localProxyPort, "5432"},
+		OrganizationSlug: orgSlug,
+		Dialer:           dialer,
+		RemoteHost:       cluster.IpAssignments.Direct,
+	}, response.Password.Value, nil
+}

internal/command/root/root.go

@@ -44,6 +44,7 @@ import (
 	"github.com/superfly/flyctl/internal/command/machine"
 	"github.com/superfly/flyctl/internal/command/metrics"
 	"github.com/superfly/flyctl/internal/command/move"
+	"github.com/superfly/flyctl/internal/command/mpg"
 	"github.com/superfly/flyctl/internal/command/mysql"
 	"github.com/superfly/flyctl/internal/command/open"
 	"github.com/superfly/flyctl/internal/command/orgs"
@@ -126,6 +127,7 @@ func New() *cobra.Command {
 		group(ping.New(), "upkeep"),
 		group(proxy.New(), "upkeep"),
 		group(postgres.New(), "dbs_and_extensions"),
+		group(mpg.New(), "dbs_and_extensions"),
 		group(ips.New(), "configuring"),
 		group(secrets.New(), "configuring"),
 		group(ssh.New(), "upkeep"),

internal/uiex/client.go

@@ -0,0 +1,65 @@
+package uiex
+
+import (
+	"context"
+	"fmt"
+	"net/http"
+	"net/url"
+	"os"
+
+	"github.com/superfly/fly-go"
+	"github.com/superfly/fly-go/tokens"
+)
+
+type Client struct {
+	baseUrl    *url.URL
+	tokens     *tokens.Tokens
+	httpClient *http.Client
+	userAgent  string
+}
+
+type NewClientOpts struct {
+	// optional, sent with requests
+	UserAgent string
+
+	// URL used when connecting via usermode wireguard.
+	BaseURL *url.URL
+
+	Tokens *tokens.Tokens
+
+	// optional:
+	Logger fly.Logger
+
+	// optional, used to construct the underlying HTTP client
+	Transport http.RoundTripper
+}
+
+func NewWithOptions(ctx context.Context, opts NewClientOpts) (*Client, error) {
+	var err error
+	uiexBaseURL := os.Getenv("FLY_UIEX_BASE_URL")
+
+	if uiexBaseURL == "" {
+		uiexBaseURL = "https://api.fly.io/graphql"
+	}
+	uiexUrl, err := url.Parse(uiexBaseURL)
+	if err != nil {
+		return nil, fmt.Errorf("invalid FLY_UIEX_BASE_URL '%s' with error: %w", uiexBaseURL, err)
+	}
+
+	httpClient, err := fly.NewHTTPClient(opts.Logger, http.DefaultTransport)
+	if err != nil {
+		return nil, fmt.Errorf("uiex: can't setup HTTP client to %s: %w", uiexUrl.String(), err)
+	}
+
+	userAgent := "flyctl"
+	if opts.UserAgent != "" {
+		userAgent = opts.UserAgent
+	}
+
+	return &Client{
+		baseUrl:    uiexUrl,
+		tokens:     opts.Tokens,
+		httpClient: httpClient,
+		userAgent:  userAgent,
+	}, nil
+}