Skip to content

Commit 778a18f

Browse files
gargakshitgargakshit
authored
mpg: add user management (#4649)
1 parent 5a396de commit 778a18f

File tree

9 files changed

+987
-33
lines changed

9 files changed

+987
-33
lines changed

‎internal/command/launch/plan/postgres_test.go‎

Lines changed: 20 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -54,6 +54,26 @@ func (m *mockUIEXClient) CreateUser(ctx context.Context, id string, input uiex.C
5454
return uiex.CreateUserResponse{}, nil
5555
}
5656

57+
func (m *mockUIEXClient) CreateUserWithRole(ctx context.Context, id string, input uiex.CreateUserWithRoleInput) (uiex.CreateUserWithRoleResponse, error) {
58+
return uiex.CreateUserWithRoleResponse{}, nil
59+
}
60+
61+
func (m *mockUIEXClient) UpdateUserRole(ctx context.Context, id string, username string, input uiex.UpdateUserRoleInput) (uiex.UpdateUserRoleResponse, error) {
62+
return uiex.UpdateUserRoleResponse{}, nil
63+
}
64+
65+
func (m *mockUIEXClient) DeleteUser(ctx context.Context, id string, username string) error {
66+
return nil
67+
}
68+
69+
func (m *mockUIEXClient) GetUserCredentials(ctx context.Context, id string, username string) (uiex.GetUserCredentialsResponse, error) {
70+
return uiex.GetUserCredentialsResponse{}, nil
71+
}
72+
73+
func (m *mockUIEXClient) ListUsers(ctx context.Context, id string) (uiex.ListUsersResponse, error) {
74+
return uiex.ListUsersResponse{}, nil
75+
}
76+
5777
func (m *mockUIEXClient) ListDatabases(ctx context.Context, id string) (uiex.ListDatabasesResponse, error) {
5878
return uiex.ListDatabasesResponse{}, nil
5979
}

‎internal/command/mpg/attach.go‎

Lines changed: 108 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -3,6 +3,7 @@ package mpg
33
import (
44
"context"
55
"fmt"
6+
"net/url"
67

78
"github.com/spf13/cobra"
89
"github.com/superfly/flyctl/internal/appconfig"
@@ -11,6 +12,8 @@ import (
1112
"github.com/superfly/flyctl/internal/flag"
1213
"github.com/superfly/flyctl/internal/flapsutil"
1314
"github.com/superfly/flyctl/internal/flyutil"
15+
"github.com/superfly/flyctl/internal/prompt"
16+
"github.com/superfly/flyctl/internal/uiex"
1417
"github.com/superfly/flyctl/internal/uiexutil"
1518
"github.com/superfly/flyctl/iostreams"
1619
)
@@ -40,6 +43,16 @@ func newAttach() *cobra.Command {
4043
Default: "DATABASE_URL",
4144
Description: "The name of the environment variable that will be added to the attached app",
4245
},
46+
flag.String{
47+
Name: "database",
48+
Shorthand: "d",
49+
Description: "The database to connect to",
50+
},
51+
flag.String{
52+
Name: "username",
53+
Shorthand: "u",
54+
Description: "The username to connect as",
55+
},
4356
)
4457

4558
return cmd
@@ -65,6 +78,9 @@ func runAttach(ctx context.Context) error {
6578
}
6679

6780
appOrgSlug := app.Organization.RawSlug
81+
if appOrgSlug != "" && clusterId == "" {
82+
fmt.Fprintf(io.Out, "Listing clusters in organization %s\n", appOrgSlug)
83+
}
6884

6985
// Get cluster details to determine which org it belongs to
7086
cluster, _, err := ClusterFromArgOrSelect(ctx, clusterId, appOrgSlug)
@@ -82,11 +98,87 @@ func runAttach(ctx context.Context) error {
8298

8399
uiexClient := uiexutil.ClientFromContext(ctx)
84100

101+
// Username selection: flag > prompt (if interactive) > empty (use default credentials)
102+
username := flag.GetString(ctx, "username")
103+
if username == "" && io.IsInteractive() {
104+
// Prompt for user selection
105+
usersResponse, err := uiexClient.ListUsers(ctx, cluster.Id)
106+
if err != nil {
107+
return fmt.Errorf("failed to list users: %w", err)
108+
}
109+
110+
if len(usersResponse.Data) > 0 {
111+
var userOptions []string
112+
for _, user := range usersResponse.Data {
113+
userOptions = append(userOptions, fmt.Sprintf("%s [%s]", user.Name, user.Role))
114+
}
115+
116+
var userIndex int
117+
err = prompt.Select(ctx, &userIndex, "Select user:", "", userOptions...)
118+
if err != nil {
119+
return err
120+
}
121+
122+
username = usersResponse.Data[userIndex].Name
123+
}
124+
// If no users found, username remains empty and will use default credentials
125+
}
126+
127+
// Database selection priority: flag > prompt result (if interactive) > credentials.DBName
128+
var db string
129+
if database := flag.GetString(ctx, "database"); database != "" {
130+
db = database
131+
} else if io.IsInteractive() {
132+
// Prompt for database selection
133+
databasesResponse, err := uiexClient.ListDatabases(ctx, cluster.Id)
134+
if err != nil {
135+
return fmt.Errorf("failed to list databases: %w", err)
136+
}
137+
138+
if len(databasesResponse.Data) > 0 {
139+
var dbOptions []string
140+
for _, database := range databasesResponse.Data {
141+
dbOptions = append(dbOptions, database.Name)
142+
}
143+
144+
var dbIndex int
145+
err = prompt.Select(ctx, &dbIndex, "Select database:", "", dbOptions...)
146+
if err != nil {
147+
return err
148+
}
149+
150+
db = databasesResponse.Data[dbIndex].Name
151+
}
152+
}
153+
154+
// Get cluster details with credentials
85155
response, err := uiexClient.GetManagedClusterById(ctx, cluster.Id)
86156
if err != nil {
87157
return fmt.Errorf("failed retrieving cluster %s: %w", clusterId, err)
88158
}
89159

160+
// Get credentials - use user-specific endpoint if username provided, otherwise use default
161+
var credentials uiex.GetManagedClusterCredentialsResponse
162+
if username != "" {
163+
userCreds, err := uiexClient.GetUserCredentials(ctx, cluster.Id, username)
164+
if err != nil {
165+
return fmt.Errorf("failed retrieving credentials for user %s: %w", username, err)
166+
}
167+
// Convert user credentials to the standard format
168+
credentials = uiex.GetManagedClusterCredentialsResponse{
169+
User: userCreds.Data.User,
170+
Password: userCreds.Data.Password,
171+
DBName: response.Credentials.DBName, // Use default DB name from cluster credentials
172+
}
173+
} else {
174+
credentials = response.Credentials
175+
}
176+
177+
// Use selected database or fall back to default from credentials
178+
if db == "" {
179+
db = credentials.DBName
180+
}
181+
90182
ctx, flapsClient, _, err := flapsutil.SetClient(ctx, nil, appName)
91183
if err != nil {
92184
return err
@@ -110,15 +202,28 @@ func runAttach(ctx context.Context) error {
110202
}
111203
}
112204

205+
// Build connection URI with selected user and database
206+
// Parse the base connection URI to extract host/port
207+
baseUri := response.Credentials.ConnectionUri
208+
parsedUri, err := url.Parse(baseUri)
209+
if err != nil {
210+
return fmt.Errorf("failed to parse connection URI: %w", err)
211+
}
212+
213+
// Build new connection URI with selected user, password, and database
214+
parsedUri.User = url.UserPassword(credentials.User, credentials.Password)
215+
parsedUri.Path = "/" + db
216+
connectionUri := parsedUri.String()
217+
113218
s := map[string]string{}
114-
s[variableName] = response.Credentials.ConnectionUri
219+
s[variableName] = connectionUri
115220

116221
if err := appsecrets.Update(ctx, flapsClient, app.Name, s, nil); err != nil {
117222
return err
118223
}
119224

120-
fmt.Fprintf(io.Out, "\nPostgres cluster %s is being attached to %s\n", clusterId, appName)
121-
fmt.Fprintf(io.Out, "The following secret was added to %s:\n %s=%s\n", appName, variableName, response.Credentials.ConnectionUri)
225+
fmt.Fprintf(io.Out, "\nPostgres cluster %s is being attached to %s\n", cluster.Id, appName)
226+
fmt.Fprintf(io.Out, "The following secret was added to %s:\n %s=%s\n", appName, variableName, connectionUri)
122227

123228
return nil
124229
}

‎internal/command/mpg/connect.go‎

Lines changed: 91 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -9,6 +9,9 @@ import (
99
"github.com/spf13/cobra"
1010
"github.com/superfly/flyctl/internal/command"
1111
"github.com/superfly/flyctl/internal/flag"
12+
"github.com/superfly/flyctl/internal/prompt"
13+
"github.com/superfly/flyctl/internal/uiex"
14+
"github.com/superfly/flyctl/internal/uiexutil"
1215
"github.com/superfly/flyctl/iostreams"
1316
"github.com/superfly/flyctl/proxy"
1417
)
@@ -29,6 +32,11 @@ func newConnect() (cmd *cobra.Command) {
2932
Shorthand: "d",
3033
Description: "The database to connect to",
3134
},
35+
flag.String{
36+
Name: "username",
37+
Shorthand: "u",
38+
Description: "The username to connect as",
39+
},
3240
)
3341
cmd.Args = cobra.MaximumNArgs(1)
3442

@@ -45,7 +53,86 @@ func runConnect(ctx context.Context) (err error) {
4553

4654
localProxyPort := "16380"
4755

48-
cluster, params, credentials, err := getMpgProxyParams(ctx, localProxyPort)
56+
// Get cluster once (will prompt if needed)
57+
clusterID := flag.FirstArg(ctx)
58+
var cluster *uiex.ManagedCluster
59+
var orgSlug string
60+
61+
if clusterID != "" {
62+
// If cluster ID is provided, fetch directly without prompting for org
63+
uiexClient := uiexutil.ClientFromContext(ctx)
64+
response, err := uiexClient.GetManagedClusterById(ctx, clusterID)
65+
if err != nil {
66+
return fmt.Errorf("failed retrieving cluster %s: %w", clusterID, err)
67+
}
68+
cluster = &response.Data
69+
orgSlug = cluster.Organization.Slug
70+
} else {
71+
// Otherwise, prompt for org/cluster selection
72+
var err error
73+
cluster, orgSlug, err = ClusterFromArgOrSelect(ctx, clusterID, "")
74+
if err != nil {
75+
return err
76+
}
77+
}
78+
79+
// Username selection: flag > prompt (if interactive) > empty (use default credentials)
80+
username := flag.GetString(ctx, "username")
81+
if username == "" && io.IsInteractive() {
82+
// Prompt for user selection
83+
uiexClient := uiexutil.ClientFromContext(ctx)
84+
usersResponse, err := uiexClient.ListUsers(ctx, cluster.Id)
85+
if err != nil {
86+
return fmt.Errorf("failed to list users: %w", err)
87+
}
88+
89+
if len(usersResponse.Data) > 0 {
90+
var userOptions []string
91+
for _, user := range usersResponse.Data {
92+
userOptions = append(userOptions, fmt.Sprintf("%s [%s]", user.Name, user.Role))
93+
}
94+
95+
var userIndex int
96+
err = prompt.Select(ctx, &userIndex, "Select user:", "", userOptions...)
97+
if err != nil {
98+
return err
99+
}
100+
101+
username = usersResponse.Data[userIndex].Name
102+
}
103+
// If no users found, username remains empty and will use default credentials
104+
}
105+
106+
// Database selection priority: flag > prompt result (if interactive) > credentials.DBName
107+
// We'll get credentials from getMpgProxyParams, but need to prompt for database first if needed
108+
var db string
109+
if database := flag.GetString(ctx, "database"); database != "" {
110+
db = database
111+
} else if io.IsInteractive() {
112+
// Prompt for database selection
113+
uiexClient := uiexutil.ClientFromContext(ctx)
114+
databasesResponse, err := uiexClient.ListDatabases(ctx, cluster.Id)
115+
if err != nil {
116+
return fmt.Errorf("failed to list databases: %w", err)
117+
}
118+
119+
if len(databasesResponse.Data) > 0 {
120+
var dbOptions []string
121+
for _, database := range databasesResponse.Data {
122+
dbOptions = append(dbOptions, database.Name)
123+
}
124+
125+
var dbIndex int
126+
err = prompt.Select(ctx, &dbIndex, "Select database:", "", dbOptions...)
127+
if err != nil {
128+
return err
129+
}
130+
131+
db = databasesResponse.Data[dbIndex].Name
132+
}
133+
}
134+
135+
cluster, params, credentials, err := getMpgProxyParamsWithCluster(ctx, localProxyPort, username, cluster.Id, orgSlug)
49136
if err != nil {
50137
return err
51138
}
@@ -67,11 +154,10 @@ func runConnect(ctx context.Context) (err error) {
67154

68155
user := credentials.User
69156
password := credentials.Password
70-
db := credentials.DBName
71157

72-
// Override database name if provided via flag
73-
if database := flag.GetString(ctx, "database"); database != "" {
74-
db = database
158+
// Use selected database or fall back to default from credentials
159+
if db == "" {
160+
db = credentials.DBName
75161
}
76162

77163
connectUrl := fmt.Sprintf("postgresql://%s:%s@localhost:%s/%s", user, password, localProxyPort, db)

‎internal/command/mpg/mpg.go‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -79,6 +79,7 @@ func New() *cobra.Command {
7979
newBackup(),
8080
newRestore(),
8181
newDatabases(),
82+
newUsers(),
8283
)
8384

8485
return cmd

‎internal/command/mpg/mpg_test.go‎

Lines changed: 40 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -27,6 +27,11 @@ type MockUiexClient struct {
2727
GetManagedClusterFunc func(ctx context.Context, orgSlug string, id string) (uiex.GetManagedClusterResponse, error)
2828
GetManagedClusterByIdFunc func(ctx context.Context, id string) (uiex.GetManagedClusterResponse, error)
2929
CreateUserFunc func(ctx context.Context, id string, input uiex.CreateUserInput) (uiex.CreateUserResponse, error)
30+
CreateUserWithRoleFunc func(ctx context.Context, id string, input uiex.CreateUserWithRoleInput) (uiex.CreateUserWithRoleResponse, error)
31+
UpdateUserRoleFunc func(ctx context.Context, id string, username string, input uiex.UpdateUserRoleInput) (uiex.UpdateUserRoleResponse, error)
32+
DeleteUserFunc func(ctx context.Context, id string, username string) error
33+
GetUserCredentialsFunc func(ctx context.Context, id string, username string) (uiex.GetUserCredentialsResponse, error)
34+
ListUsersFunc func(ctx context.Context, id string) (uiex.ListUsersResponse, error)
3035
ListDatabasesFunc func(ctx context.Context, id string) (uiex.ListDatabasesResponse, error)
3136
CreateDatabaseFunc func(ctx context.Context, id string, input uiex.CreateDatabaseInput) (uiex.CreateDatabaseResponse, error)
3237
CreateClusterFunc func(ctx context.Context, input uiex.CreateClusterInput) (uiex.CreateClusterResponse, error)
@@ -72,6 +77,41 @@ func (m *MockUiexClient) CreateUser(ctx context.Context, id string, input uiex.C
7277
return uiex.CreateUserResponse{}, nil
7378
}
7479

80+
func (m *MockUiexClient) CreateUserWithRole(ctx context.Context, id string, input uiex.CreateUserWithRoleInput) (uiex.CreateUserWithRoleResponse, error) {
81+
if m.CreateUserWithRoleFunc != nil {
82+
return m.CreateUserWithRoleFunc(ctx, id, input)
83+
}
84+
return uiex.CreateUserWithRoleResponse{}, nil
85+
}
86+
87+
func (m *MockUiexClient) UpdateUserRole(ctx context.Context, id string, username string, input uiex.UpdateUserRoleInput) (uiex.UpdateUserRoleResponse, error) {
88+
if m.UpdateUserRoleFunc != nil {
89+
return m.UpdateUserRoleFunc(ctx, id, username, input)
90+
}
91+
return uiex.UpdateUserRoleResponse{}, nil
92+
}
93+
94+
func (m *MockUiexClient) DeleteUser(ctx context.Context, id string, username string) error {
95+
if m.DeleteUserFunc != nil {
96+
return m.DeleteUserFunc(ctx, id, username)
97+
}
98+
return nil
99+
}
100+
101+
func (m *MockUiexClient) GetUserCredentials(ctx context.Context, id string, username string) (uiex.GetUserCredentialsResponse, error) {
102+
if m.GetUserCredentialsFunc != nil {
103+
return m.GetUserCredentialsFunc(ctx, id, username)
104+
}
105+
return uiex.GetUserCredentialsResponse{}, nil
106+
}
107+
108+
func (m *MockUiexClient) ListUsers(ctx context.Context, id string) (uiex.ListUsersResponse, error) {
109+
if m.ListUsersFunc != nil {
110+
return m.ListUsersFunc(ctx, id)
111+
}
112+
return uiex.ListUsersResponse{}, nil
113+
}
114+
75115
func (m *MockUiexClient) ListDatabases(ctx context.Context, id string) (uiex.ListDatabasesResponse, error) {
76116
if m.ListDatabasesFunc != nil {
77117
return m.ListDatabasesFunc(ctx, id)

0 commit comments

Comments
 (0)