Merge 80aad05 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/postgres/postgres.go

@@ -44,6 +44,7 @@ func New() *cobra.Command {
 		newImport(),
 		newEvents(),
 		newBarman(),
+		newProxy(),
 	)
 
 	return cmd

internal/command/postgres/proxy.go

@@ -0,0 +1,106 @@
+package postgres
+
+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/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) (*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, "", 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, "", selectErr
+	}
+
+	selectedCluster := clustersResponse.Data[index]
+
+	response, err := uiexClient.GetManagedCluster(ctx, selectedCluster.Organization.Slug, selectedCluster.Id)
+	if err != nil {
+		return nil, "", err
+	}
+
+	if response.Password.Status == "initializing" {
+		return nil, "", fmt.Errorf("Cluster is still initializing, wait a bit more")
+	}
+
+	if response.Password.Status == "error" {
+		return nil, "", fmt.Errorf("Error getting cluster password")
+	}
+
+	agentclient, err := agent.Establish(ctx, client)
+	if err != nil {
+		return nil, "", err
+	}
+
+	dialer, err := agentclient.ConnectToTunnel(ctx, orgSlug, "", false)
+	if err != nil {
+		return nil, "", err
+	}
+
+	return &proxy.ConnectParams{
+		Ports:            []string{localProxyPort, "5432"},
+		OrganizationSlug: orgSlug,
+		Dialer:           dialer,
+		RemoteHost:       response.Data.IpAssignments.Direct,
+	}, response.Password.Value, nil
+}

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

internal/uiex/managed_postgres.go

@@ -0,0 +1,104 @@
+package uiex
+
+import (
+	"context"
+	"encoding/json"
+	"fmt"
+	"net/http"
+
+	"github.com/superfly/fly-go"
+	"github.com/superfly/flyctl/internal/config"
+)
+
+type ManagedClusterIpAssignments struct {
+	Direct string `json:"direct"`
+}
+type ManagedCluster struct {
+	Id            string                      `json:"id"`
+	Name          string                      `json:"name"`
+	Region        string                      `json:"region"`
+	Status        string                      `json:"status"`
+	Plan          string                      `json:"plan"`
+	Organization  fly.Organization            `json:"organization"`
+	IpAssignments ManagedClusterIpAssignments `json:"ip_assignments"`
+}
+
+type ListManagedClustersResponse struct {
+	Data []ManagedCluster `json:"data"`
+}
+
+type GetManagedClusterPasswordResponse struct {
+	Status string `json:"status"`
+	Value  string `json:"value"`
+}
+
+type GetManagedClusterResponse struct {
+	Data     ManagedCluster                    `json:"data"`
+	Password GetManagedClusterPasswordResponse `json:"password"`
+}
+
+func (c *Client) ListManagedClusters(ctx context.Context, orgSlug string) (ListManagedClustersResponse, error) {
+	var response ListManagedClustersResponse
+
+	cfg := config.FromContext(ctx)
+	url := fmt.Sprintf("%s/api/v1/organizations/%s/postgres", c.baseUrl, orgSlug)
+
+	req, err := http.NewRequestWithContext(ctx, http.MethodGet, url, nil)
+	if err != nil {
+		return response, fmt.Errorf("failed to create request: %w", err)
+	}
+
+	req.Header.Add("Authorization", "Bearer "+cfg.Tokens.GraphQL())
+	req.Header.Add("Content-Type", "application/json")
+
+	res, err := http.DefaultClient.Do(req)
+	if err != nil {
+		return response, err
+	}
+	defer res.Body.Close()
+
+	switch res.StatusCode {
+	case http.StatusOK:
+		if err = json.NewDecoder(res.Body).Decode(&response); err != nil {
+			return response, fmt.Errorf("failed to decode response, please try again: %w", err)
+		}
+		return response, nil
+	case http.StatusNotFound:
+		return response, err
+	default:
+		return response, err
+	}
+
+}
+
+func (c *Client) GetManagedCluster(ctx context.Context, orgSlug string, id string) (GetManagedClusterResponse, error) {
+	var response GetManagedClusterResponse
+	cfg := config.FromContext(ctx)
+	url := fmt.Sprintf("%s/api/v1/organizations/%s/postgres/%s", c.baseUrl, orgSlug, id)
+
+	req, err := http.NewRequestWithContext(ctx, http.MethodGet, url, nil)
+	if err != nil {
+		return response, fmt.Errorf("failed to create request: %w", err)
+	}
+
+	req.Header.Add("Authorization", "Bearer "+cfg.Tokens.GraphQL())
+	req.Header.Add("Content-Type", "application/json")
+
+	res, err := http.DefaultClient.Do(req)
+	if err != nil {
+		return response, err
+	}
+	defer res.Body.Close()
+
+	switch res.StatusCode {
+	case http.StatusOK:
+		if err = json.NewDecoder(res.Body).Decode(&response); err != nil {
+			return response, fmt.Errorf("failed to decode response, please try again: %w", err)
+		}
+		return response, nil
+	case http.StatusNotFound:
+		return response, err
+	default:
+		return response, err
+	}
+}

internal/uiexutil/client.go

@@ -0,0 +1,27 @@
+package uiexutil
+
+import (
+	"context"
+
+	"github.com/superfly/flyctl/internal/uiex"
+)
+
+type Client interface {
+	ListManagedClusters(ctx context.Context, orgSlug string) (uiex.ListManagedClustersResponse, error)
+	GetManagedCluster(ctx context.Context, orgSlug string, id string) (uiex.GetManagedClusterResponse, error)
+}
+
+type contextKey struct{}
+
+var clientContextKey = &contextKey{}
+
+// NewContextWithClient derives a Context that carries c from ctx.
+func NewContextWithClient(ctx context.Context, c Client) context.Context {
+	return context.WithValue(ctx, clientContextKey, c)
+}
+
+// ClientFromContext returns the Client ctx carries.
+func ClientFromContext(ctx context.Context) Client {
+	c, _ := ctx.Value(clientContextKey).(Client)
+	return c
+}

internal/uiexutil/uiexutil.go

@@ -0,0 +1,20 @@
+package uiexutil
+
+import (
+	"context"
+
+	"github.com/superfly/flyctl/internal/config"
+	"github.com/superfly/flyctl/internal/logger"
+	"github.com/superfly/flyctl/internal/uiex"
+)
+
+func NewClientWithOptions(ctx context.Context, opts uiex.NewClientOpts) (*uiex.Client, error) {
+	if opts.Tokens == nil {
+		opts.Tokens = config.Tokens(ctx)
+	}
+
+	if v := logger.MaybeFromContext(ctx); v != nil && opts.Logger == nil {
+		opts.Logger = v
+	}
+	return uiex.NewWithOptions(ctx, opts)
+}