Add api membership during organization creation

Signed-off-by: Sacha Al Himdani <sacha@getprobo.com>

pkg/auth/service.go

@@ -1574,6 +1574,39 @@ func (s *Service) UpdateUserAPIKeyMemberships(
 	)
 }
 
+func (s *Service) AddAPIKeyMembershipToOrganization(
+	ctx context.Context,
+	tenantID gid.TenantID,
+	userAPIKeyID gid.GID,
+	membershipID gid.GID,
+	organizationID gid.GID,
+	role coredata.APIRole,
+) error {
+	scope := coredata.NewScope(tenantID)
+	now := time.Now()
+
+	return s.pg.WithConn(
+		ctx,
+		func(conn pg.Conn) error {
+			userAPIKeyMembership := &coredata.UserAPIKeyMembership{
+				ID:             gid.New(tenantID, coredata.UserAPIKeyMembershipEntityType),
+				UserAPIKeyID:   userAPIKeyID,
+				MembershipID:   membershipID,
+				Role:           role,
+				OrganizationID: organizationID,
+				CreatedAt:      now,
+				UpdatedAt:      now,
+			}
+
+			if err := userAPIKeyMembership.Insert(ctx, conn, scope); err != nil {
+				return fmt.Errorf("cannot insert user api key membership: %w", err)
+			}
+
+			return nil
+		},
+	)
+}
+
 func (s *Service) UpdateUserAPIKeyName(
 	ctx context.Context,
 	userAPIKeyID gid.GID,

pkg/server/api/console/v1/v1_resolver.go

@@ -1377,6 +1377,7 @@ func (r *membershipConnectionResolver) TotalCount(ctx context.Context, obj *type
 // CreateOrganization is the resolver for the createOrganization field.
 func (r *mutationResolver) CreateOrganization(ctx context.Context, input types.CreateOrganizationInput) (*types.CreateOrganizationPayload, error) {
 	currentUser := UserFromContext(ctx)
+	currentAPIKey := UserAPIKeyFromContext(ctx)
 
 	tenantID := gid.NewTenantID()
 
@@ -1411,6 +1412,25 @@ func (r *mutationResolver) CreateOrganization(ctx context.Context, input types.C
 		panic(fmt.Errorf("cannot add user to organization: %w", err))
 	}
 
+	if currentAPIKey != nil {
+		membership, err := authz.GetMembershipByUserAndOrganizationID(ctx, currentUser.ID, organization.ID)
+		if err != nil {
+			panic(fmt.Errorf("cannot get user membership: %w", err))
+		}
+
+		err = r.authSvc.AddAPIKeyMembershipToOrganization(
+			ctx,
+			tenantID,
+			currentAPIKey.ID,
+			membership.ID,
+			organization.ID,
+			coredata.APIRoleFull,
+		)
+		if err != nil {
+			panic(fmt.Errorf("cannot add API key membership to organization: %w", err))
+		}
+	}
+
 	_, err = prb.Peoples.Create(
 		ctx,
 		probo.CreatePeopleRequest{