Skip to content

Commit 4216ab2

Browse files
patelspratikpatelspratik
committed
feat: migrate identity stores to dev-plane
1 parent f636937 commit 4216ab2

8 files changed

Lines changed: 405 additions & 138 deletions

File tree

‎pkg/store/devplane.go‎

Lines changed: 198 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,198 @@
1+
package store
2+
3+
import (
4+
"context"
5+
"encoding/json"
6+
"fmt"
7+
"net/http"
8+
9+
nodev1connect "buf.build/gen/go/brevdev/devplane/connectrpc/go/devplaneapi/v1/devplaneapiv1connect"
10+
nodev1 "buf.build/gen/go/brevdev/devplane/protocolbuffers/go/devplaneapi/v1"
11+
"connectrpc.com/connect"
12+
13+
"github.com/brevdev/brev-cli/pkg/config"
14+
"github.com/brevdev/brev-cli/pkg/entity"
15+
breverrors "github.com/brevdev/brev-cli/pkg/errors"
16+
)
17+
18+
type devPlaneUserService interface {
19+
GetCurrentUser(context.Context) (*nodev1.User, string, error)
20+
GetCurrentUserKeys(context.Context) (*entity.UserKeys, error)
21+
GetUser(context.Context, string) (*nodev1.User, error)
22+
SetUserBlocked(context.Context, string, bool) (*nodev1.User, error)
23+
}
24+
25+
type devPlaneOrganizationService interface {
26+
ListOrganizations(context.Context, []string) ([]*nodev1.Organization, error)
27+
GetOrganization(context.Context, string) (*nodev1.Organization, error)
28+
}
29+
30+
type devPlaneServices struct {
31+
user devPlaneUserService
32+
organization devPlaneOrganizationService
33+
}
34+
35+
type authHTTPStoreTransport struct {
36+
store *AuthHTTPStore
37+
base http.RoundTripper
38+
}
39+
40+
func (t *authHTTPStoreTransport) RoundTrip(req *http.Request) (*http.Response, error) {
41+
token, err := t.store.GetAccessToken()
42+
if err != nil {
43+
return nil, breverrors.WrapAndTrace(err)
44+
}
45+
req = req.Clone(req.Context())
46+
req.Header.Set("Authorization", "Bearer "+token)
47+
resp, err := t.base.RoundTrip(req)
48+
if err != nil {
49+
return nil, breverrors.WrapAndTrace(err)
50+
}
51+
return resp, nil
52+
}
53+
54+
func (s *AuthHTTPStore) devPlaneHTTPClient() *http.Client {
55+
return &http.Client{Transport: &authHTTPStoreTransport{store: s, base: http.DefaultTransport}}
56+
}
57+
58+
func (s *AuthHTTPStore) withDevPlaneServices(services *devPlaneServices) *AuthHTTPStore {
59+
copy := *s
60+
copy.devPlane = services
61+
return &copy
62+
}
63+
64+
func (s *AuthHTTPStore) devPlaneServiceClients() *devPlaneServices {
65+
if s.devPlane != nil {
66+
return s.devPlane
67+
}
68+
httpClient := s.devPlaneHTTPClient()
69+
baseURL := config.GlobalConfig.GetBrevPublicAPIURL()
70+
return &devPlaneServices{
71+
user: &generatedDevPlaneUserService{client: nodev1connect.NewUserServiceClient(httpClient, baseURL)},
72+
organization: &generatedDevPlaneOrganizationService{
73+
client: nodev1connect.NewOrganizationServiceClient(httpClient, baseURL),
74+
},
75+
}
76+
}
77+
78+
type generatedDevPlaneUserService struct {
79+
client nodev1connect.UserServiceClient
80+
}
81+
82+
func (s *generatedDevPlaneUserService) GetCurrentUser(ctx context.Context) (*nodev1.User, string, error) {
83+
response, err := s.client.GetCurrentUser(ctx, connect.NewRequest(&nodev1.GetCurrentUserRequest{}))
84+
if err != nil {
85+
return nil, "", breverrors.WrapAndTrace(err)
86+
}
87+
// The public SSH-key response field is added by Set 1. Keep the local
88+
// contract ready and populate it after the landed Buf revision is available.
89+
return response.Msg.GetUser(), "", nil
90+
}
91+
92+
func (s *generatedDevPlaneUserService) GetCurrentUserKeys(context.Context) (*entity.UserKeys, error) {
93+
return nil, fmt.Errorf("GetCurrentUserSSHKeys requires the published Set 1 Buf revision")
94+
}
95+
96+
func (s *generatedDevPlaneUserService) GetUser(ctx context.Context, userID string) (*nodev1.User, error) {
97+
response, err := s.client.GetUser(ctx, connect.NewRequest(&nodev1.GetUserRequest{
98+
UserId: userID,
99+
AttachedDataOptions: &nodev1.UserAttachedDataOptions{
100+
Emails: true,
101+
ExternalIdentities: true,
102+
OrganizationAccesses: true,
103+
},
104+
}))
105+
if err != nil {
106+
return nil, breverrors.WrapAndTrace(err)
107+
}
108+
return response.Msg.GetUser(), nil
109+
}
110+
111+
func (s *generatedDevPlaneUserService) SetUserBlocked(context.Context, string, bool) (*nodev1.User, error) {
112+
return nil, fmt.Errorf("SetUserBlocked requires the published Set 1 Buf revision")
113+
}
114+
115+
type generatedDevPlaneOrganizationService struct {
116+
client nodev1connect.OrganizationServiceClient
117+
}
118+
119+
func (s *generatedDevPlaneOrganizationService) ListOrganizations(ctx context.Context, ids []string) ([]*nodev1.Organization, error) {
120+
organizations := []*nodev1.Organization{}
121+
pageToken := ""
122+
for {
123+
response, err := s.client.ListOrganization(ctx, connect.NewRequest(&nodev1.ListOrganizationRequest{
124+
PageParams: &nodev1.PageParams{PageSize: 1000, PageToken: pageToken},
125+
Options: &nodev1.ListOrganizationOptions{Ids: ids},
126+
}))
127+
if err != nil {
128+
return nil, breverrors.WrapAndTrace(err)
129+
}
130+
organizations = append(organizations, response.Msg.GetItems()...)
131+
pageToken = response.Msg.GetNextPageToken()
132+
if pageToken == "" {
133+
return organizations, nil
134+
}
135+
}
136+
}
137+
138+
func (s *generatedDevPlaneOrganizationService) GetOrganization(ctx context.Context, id string) (*nodev1.Organization, error) {
139+
response, err := s.client.GetOrganization(ctx, connect.NewRequest(&nodev1.GetOrganizationRequest{OrganizationId: id}))
140+
if err != nil {
141+
return nil, breverrors.WrapAndTrace(err)
142+
}
143+
return response.Msg.GetOrganization(), nil
144+
}
145+
146+
func mapDevPlaneUser(user *nodev1.User, publicSSHKey string) (*entity.User, error) {
147+
if user == nil {
148+
return nil, fmt.Errorf("dev-plane returned no user")
149+
}
150+
metadata := struct {
151+
BaseWorkspaceRepo string `json:"baseWorkspaceRepo"`
152+
UserSetupExecPath string `json:"userSetupExecPath"`
153+
WorkspacePassword string `json:"workspacePassword"`
154+
GlobalUserType entity.GlobalUserType `json:"globalUserType"`
155+
IDEConfig entity.IDEConfig `json:"ideConfig"`
156+
OnboardingData map[string]interface{} `json:"onboardingData"`
157+
}{}
158+
if user.GetMetadata() != nil {
159+
data, err := json.Marshal(user.GetMetadata().AsMap())
160+
if err != nil {
161+
return nil, breverrors.WrapAndTrace(err)
162+
}
163+
if err := json.Unmarshal(data, &metadata); err != nil {
164+
return nil, breverrors.WrapAndTrace(err)
165+
}
166+
}
167+
var externalIdentities []*entity.ExternalIdentity
168+
for _, identity := range user.GetExternalIdentities() {
169+
externalIdentities = append(externalIdentities, &entity.ExternalIdentity{
170+
IdentityID: identity.GetIdentityId(),
171+
Provider: identity.GetProvider(),
172+
ExternalID: identity.GetExternalId(),
173+
})
174+
}
175+
return &entity.User{
176+
ID: user.GetUserId(),
177+
PublicKey: publicSSHKey,
178+
Username: user.GetUsername(),
179+
Name: user.GetDisplayName(),
180+
Email: user.GetDefaultEmail(),
181+
ExternalIdentities: externalIdentities,
182+
WorkspacePassword: metadata.WorkspacePassword,
183+
BaseWorkspaceRepo: metadata.BaseWorkspaceRepo,
184+
GlobalUserType: metadata.GlobalUserType,
185+
IdeConfig: metadata.IDEConfig,
186+
OnboardingData: metadata.OnboardingData,
187+
}, nil
188+
}
189+
190+
func mapDevPlaneOrganization(organization *nodev1.Organization) (*entity.Organization, error) {
191+
if organization == nil {
192+
return nil, fmt.Errorf("dev-plane returned no organization")
193+
}
194+
return &entity.Organization{
195+
ID: organization.GetOrganizationId(),
196+
Name: organization.GetDisplayName(),
197+
}, nil
198+
}

‎pkg/store/devplane_migration_test.go‎

Lines changed: 142 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,142 @@
1+
package store
2+
3+
import (
4+
"context"
5+
"testing"
6+
7+
nodev1 "buf.build/gen/go/brevdev/devplane/protocolbuffers/go/devplaneapi/v1"
8+
"github.com/brevdev/brev-cli/pkg/entity"
9+
"github.com/stretchr/testify/require"
10+
"google.golang.org/protobuf/types/known/structpb"
11+
)
12+
13+
type mockDevPlaneUserService struct {
14+
currentUser *nodev1.User
15+
publicSSHKey string
16+
keys *entity.UserKeys
17+
getUser *nodev1.User
18+
gotUserID string
19+
gotBlockedUserID string
20+
gotBlocked bool
21+
}
22+
23+
func (m *mockDevPlaneUserService) GetCurrentUser(context.Context) (*nodev1.User, string, error) {
24+
return m.currentUser, m.publicSSHKey, nil
25+
}
26+
27+
func (m *mockDevPlaneUserService) GetCurrentUserKeys(context.Context) (*entity.UserKeys, error) {
28+
return m.keys, nil
29+
}
30+
31+
func (m *mockDevPlaneUserService) GetUser(_ context.Context, userID string) (*nodev1.User, error) {
32+
m.gotUserID = userID
33+
return m.getUser, nil
34+
}
35+
36+
func (m *mockDevPlaneUserService) SetUserBlocked(_ context.Context, userID string, blocked bool) (*nodev1.User, error) {
37+
m.gotBlockedUserID = userID
38+
m.gotBlocked = blocked
39+
return m.getUser, nil
40+
}
41+
42+
type mockDevPlaneOrganizationService struct {
43+
organizations []*nodev1.Organization
44+
organization *nodev1.Organization
45+
gotIDs []string
46+
gotID string
47+
}
48+
49+
func (m *mockDevPlaneOrganizationService) ListOrganizations(_ context.Context, ids []string) ([]*nodev1.Organization, error) {
50+
m.gotIDs = ids
51+
return m.organizations, nil
52+
}
53+
54+
func (m *mockDevPlaneOrganizationService) GetOrganization(_ context.Context, id string) (*nodev1.Organization, error) {
55+
m.gotID = id
56+
return m.organization, nil
57+
}
58+
59+
func TestIdentityStoresUseDevPlaneServices(t *testing.T) {
60+
metadata, err := structpb.NewStruct(map[string]any{
61+
"baseWorkspaceRepo": "github.com/alice/base",
62+
"globalUserType": "Admin",
63+
"onboardingData": map[string]any{
64+
"usedCLI": true,
65+
},
66+
})
67+
require.NoError(t, err)
68+
apiUser := &nodev1.User{
69+
UserId: "user-1",
70+
Username: "alice",
71+
DisplayName: "Alice",
72+
DefaultEmail: "alice@example.com",
73+
PlatformRole: "user",
74+
Metadata: metadata,
75+
ExternalIdentities: []*nodev1.ExternalIdentity{{
76+
IdentityId: "identity-1",
77+
ExternalId: "external-1",
78+
Provider: "kas",
79+
}},
80+
}
81+
users := &mockDevPlaneUserService{
82+
currentUser: apiUser,
83+
publicSSHKey: "ssh-rsa public",
84+
keys: &entity.UserKeys{PublicKey: "public", PrivateKey: "private"},
85+
getUser: apiUser,
86+
}
87+
store := MakeMockAuthHTTPStore().withDevPlaneServices(&devPlaneServices{user: users})
88+
89+
current, err := store.GetCurrentUser()
90+
require.NoError(t, err)
91+
require.Equal(t, "user-1", current.ID)
92+
require.Equal(t, "Alice", current.Name)
93+
require.Equal(t, "alice@example.com", current.Email)
94+
require.Equal(t, "ssh-rsa public", current.PublicKey)
95+
require.Equal(t, "github.com/alice/base", current.BaseWorkspaceRepo)
96+
require.Equal(t, entity.Admin, current.GlobalUserType)
97+
require.Equal(t, true, current.OnboardingData["usedCLI"])
98+
99+
keys, err := store.GetCurrentUserKeys()
100+
require.NoError(t, err)
101+
require.Equal(t, "private", keys.PrivateKey)
102+
103+
target, err := store.GetUserByID("user-2")
104+
require.NoError(t, err)
105+
require.Equal(t, "user-1", target.ID)
106+
require.Equal(t, "user-2", users.gotUserID)
107+
108+
require.NoError(t, store.BanUser("user-2"))
109+
require.Equal(t, "user-2", users.gotBlockedUserID)
110+
require.True(t, users.gotBlocked)
111+
}
112+
113+
func TestOrganizationStoresUseCurrentUserAccessesAndDevPlane(t *testing.T) {
114+
users := &mockDevPlaneUserService{currentUser: &nodev1.User{
115+
UserId: "user-1",
116+
OrganizationAccesses: []*nodev1.OrganizationAccess{
117+
{OrganizationId: "org-1"},
118+
{OrganizationId: "org-2"},
119+
},
120+
}}
121+
organizations := &mockDevPlaneOrganizationService{
122+
organizations: []*nodev1.Organization{
123+
{OrganizationId: "org-1", DisplayName: "One"},
124+
{OrganizationId: "org-2", DisplayName: "Two"},
125+
},
126+
organization: &nodev1.Organization{OrganizationId: "org-2", DisplayName: "Two"},
127+
}
128+
store := MakeMockAuthHTTPStore().withDevPlaneServices(&devPlaneServices{
129+
user: users,
130+
organization: organizations,
131+
})
132+
133+
listed, err := store.GetOrganizations(nil)
134+
require.NoError(t, err)
135+
require.Equal(t, []entity.Organization{{ID: "org-1", Name: "One"}, {ID: "org-2", Name: "Two"}}, listed)
136+
require.Equal(t, []string{"org-1", "org-2"}, organizations.gotIDs)
137+
138+
organization, err := store.GetOrganization("org-2")
139+
require.NoError(t, err)
140+
require.Equal(t, &entity.Organization{ID: "org-2", Name: "Two"}, organization)
141+
require.Equal(t, "org-2", organizations.gotID)
142+
}

‎pkg/store/http.go‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -50,6 +50,7 @@ type AuthHTTPStore struct {
5050
NoAuthHTTPStore
5151
authHTTPClient *AuthHTTPClient
5252
isRefreshTokenHandlerSet bool
53+
devPlane *devPlaneServices
5354
BasicStore
5455
}
5556

0 commit comments

Comments
 (0)