Support inviting users with custom attributes from UI and CLI

admin/database/database.go

@@ -285,6 +285,7 @@ type DB interface {
 	DeleteOrganizationInvite(ctx context.Context, id string) error
 	CountInvitesForOrganization(ctx context.Context, orgID string) (int, error)
 	UpdateOrganizationInviteRole(ctx context.Context, id, roleID string) error
+	UpdateOrganizationInviteAttributes(ctx context.Context, id string, attributes map[string]any) error
 
 	FindProjectInvites(ctx context.Context, projectID, afterEmail string, limit int) ([]*ProjectInviteWithRole, error)
 	FindProjectInvitesByEmail(ctx context.Context, userEmail string) ([]*ProjectInvite, error)
@@ -1073,19 +1074,21 @@ type MemberUsergroup struct {
 type OrganizationInvite struct {
 	ID              string
 	Email           string
-	OrgID           string    `db:"org_id"`
-	OrgRoleID       string    `db:"org_role_id"`
-	UsergroupIDs    []string  `db:"usergroup_ids"`
-	InvitedByUserID *string   `db:"invited_by_user_id"`
-	CreatedOn       time.Time `db:"created_on"`
+	OrgID           string         `db:"org_id"`
+	OrgRoleID       string         `db:"org_role_id"`
+	UsergroupIDs    []string       `db:"usergroup_ids"`
+	Attributes      map[string]any `db:"attributes"`
+	InvitedByUserID *string        `db:"invited_by_user_id"`
+	CreatedOn       time.Time      `db:"created_on"`
 }
 
 // OrganizationInviteWithRole is a convenience type used for display-friendly representation of an OrganizationInvite.
 type OrganizationInviteWithRole struct {
-	ID        string
-	Email     string
-	RoleName  string  `db:"role_name"`
-	InvitedBy *string `db:"invited_by"`
+	ID         string
+	Email      string
+	RoleName   string         `db:"role_name"`
+	Attributes map[string]any `db:"attributes"`
+	InvitedBy  *string        `db:"invited_by"`
 }
 
 // ProjectInvite represents an outstanding invitation to join a project.
@@ -1161,10 +1164,11 @@ type ProjectWhitelistedDomainWithJoinedRoleNames struct {
 }
 
 type InsertOrganizationInviteOptions struct {
-	Email     string `validate:"email"`
-	InviterID string
-	OrgID     string `validate:"required"`
-	RoleID    string `validate:"required"`
+	Email      string `validate:"email"`
+	InviterID  string
+	OrgID      string `validate:"required"`
+	RoleID     string `validate:"required"`
+	Attributes map[string]any
 }
 
 type InsertProjectInviteOptions struct {

admin/database/postgres/migrations/0096.sql

@@ -0,0 +1 @@
+ALTER TABLE org_invites ADD COLUMN attributes JSONB;

admin/database/postgres/postgres.go

@@ -2602,9 +2602,9 @@ func (c *connection) DeleteProjectMemberService(ctx context.Context, serviceID,
 }
 
 func (c *connection) FindOrganizationInvites(ctx context.Context, orgID, afterEmail string, limit int) ([]*database.OrganizationInviteWithRole, error) {
-	var res []*database.OrganizationInviteWithRole
-	err := c.getDB(ctx).SelectContext(ctx, &res, `
-		SELECT uoi.id, uoi.email, ur.name as role_name, u.email as invited_by
+	var dtos []*organizationInviteWithRoleDTO
+	err := c.getDB(ctx).SelectContext(ctx, &dtos, `
+		SELECT uoi.id, uoi.email, ur.name as role_name, uoi.attributes, u.email as invited_by
 		FROM org_invites uoi
 		JOIN org_roles ur ON uoi.org_role_id = ur.id
 		LEFT JOIN users u ON uoi.invited_by_user_id = u.id
@@ -2614,6 +2614,14 @@ func (c *connection) FindOrganizationInvites(ctx context.Context, orgID, afterEm
 	if err != nil {
 		return nil, parseErr("org invites", err)
 	}
+	res := make([]*database.OrganizationInviteWithRole, len(dtos))
+	for i, dto := range dtos {
+		var err error
+		res[i], err = dto.AsModel()
+		if err != nil {
+			return nil, err
+		}
+	}
 	return res, nil
 }
 
@@ -2657,12 +2665,17 @@ func (c *connection) InsertOrganizationInvite(ctx context.Context, opts *databas
 		return err
 	}
 
+	attrs, err := c.validateAttributes(opts.Attributes)
+	if err != nil {
+		return err
+	}
+
 	var inviterID any
 	if opts.InviterID != "" {
 		inviterID = opts.InviterID
 	}
 
-	_, err := c.getDB(ctx).ExecContext(ctx, "INSERT INTO org_invites (email, invited_by_user_id, org_id, org_role_id) VALUES ($1, $2, $3, $4)", opts.Email, inviterID, opts.OrgID, opts.RoleID)
+	_, err = c.getDB(ctx).ExecContext(ctx, "INSERT INTO org_invites (email, invited_by_user_id, org_id, org_role_id, attributes) VALUES ($1, $2, $3, $4, $5)", opts.Email, inviterID, opts.OrgID, opts.RoleID, attrs)
 	if err != nil {
 		return parseErr("org invite", err)
 	}
@@ -2695,6 +2708,16 @@ func (c *connection) UpdateOrganizationInviteRole(ctx context.Context, id, roleI
 	return checkUpdateRow("org invite", res, err)
 }
 
+func (c *connection) UpdateOrganizationInviteAttributes(ctx context.Context, id string, attributes map[string]any) error {
+	attrs, err := c.validateAttributes(attributes)
+	if err != nil {
+		return err
+	}
+
+	res, err := c.getDB(ctx).ExecContext(ctx, `UPDATE org_invites SET attributes = $1 WHERE id = $2`, attrs, id)
+	return checkUpdateRow("org invite", res, err)
+}
+
 func (c *connection) FindProjectInvites(ctx context.Context, projectID, afterEmail string, limit int) ([]*database.ProjectInviteWithRole, error) {
 	var dtos []*projectInviteWithRoleDTO
 	err := c.getDB(ctx).SelectContext(ctx, &dtos, `
@@ -3713,6 +3736,7 @@ func (c *connection) notificationTokenWithSecretFromDTO(dto *notificationTokenWi
 type organizationInviteDTO struct {
 	*database.OrganizationInvite
 	UsergroupIDs pgtype.TextArray `db:"usergroup_ids"`
+	Attributes   pgtype.JSON      `db:"attributes"`
 }
 
 func (o *organizationInviteDTO) AsModel() (*database.OrganizationInvite, error) {
@@ -3721,9 +3745,38 @@ func (o *organizationInviteDTO) AsModel() (*database.OrganizationInvite, error)
 		return nil, err
 	}
 
+	// Handle Attributes: Normalize NULL JSONB to empty map
+	var attrs map[string]any
+	if err := o.Attributes.AssignTo(&attrs); err != nil {
+		return nil, err
+	}
+	if attrs == nil {
+		attrs = make(map[string]any)
+	}
+	o.OrganizationInvite.Attributes = attrs
+
 	return o.OrganizationInvite, nil
 }
 
+type organizationInviteWithRoleDTO struct {
+	*database.OrganizationInviteWithRole
+	Attributes pgtype.JSON `db:"attributes"`
+}
+
+func (o *organizationInviteWithRoleDTO) AsModel() (*database.OrganizationInviteWithRole, error) {
+	// Handle Attributes: Normalize NULL JSONB to empty map
+	var attrs map[string]any
+	if err := o.Attributes.AssignTo(&attrs); err != nil {
+		return nil, err
+	}
+	if attrs == nil {
+		attrs = make(map[string]any)
+	}
+	o.OrganizationInviteWithRole.Attributes = attrs
+
+	return o.OrganizationInviteWithRole, nil
+}
+
 type authClientDTO struct {
 	*database.AuthClient
 	GrantTypes   pgtype.TextArray `db:"grant_types"`

admin/database/postgres/postgres_test.go

@@ -45,6 +45,7 @@ func TestPostgres(t *testing.T) {
 	t.Run("TestUpsertProjectVariable", func(t *testing.T) { testUpsertProjectVariable(t, db) })
 	t.Run("TestManagedGitRepos", func(t *testing.T) { testManagedGitRepos(t, db) })
 	t.Run("TestOrganizationMemberUserAttributes", func(t *testing.T) { testOrganizationMemberUserAttributes(t, db) })
+	t.Run("TestOrganizationInviteAttributes", func(t *testing.T) { testOrganizationInviteAttributes(t, db) })
 	t.Run("TestAttributeValidation", func(t *testing.T) { testAttributeValidation(t, db) })
 
 	t.Run("TestOrgNameValidation", func(t *testing.T) {
@@ -746,6 +747,84 @@ func testOrganizationMemberUserAttributes(t *testing.T, db database.DB) {
 	require.NoError(t, db.DeleteUser(ctx, user.ID))
 }
 
+func testOrganizationInviteAttributes(t *testing.T, db database.DB) {
+	ctx := context.Background()
+
+	org, err := db.InsertOrganization(ctx, &database.InsertOrganizationOptions{Name: "test-invite-attrs-org"})
+	require.NoError(t, err)
+
+	role, err := db.FindOrganizationRole(ctx, database.OrganizationRoleNameViewer)
+	require.NoError(t, err)
+
+	email := "invitee-attrs@rilldata.com"
+	attributes := map[string]any{"attr1": "value1", "attr2": "value2"}
+
+	t.Run("InsertOrganizationInvite with attributes", func(t *testing.T) {
+		err := db.InsertOrganizationInvite(ctx, &database.InsertOrganizationInviteOptions{
+			Email:      email,
+			OrgID:      org.ID,
+			RoleID:     role.ID,
+			Attributes: attributes,
+		})
+		require.NoError(t, err)
+
+		invite, err := db.FindOrganizationInvite(ctx, org.ID, email)
+		require.NoError(t, err)
+		require.Equal(t, attributes, invite.Attributes)
+
+		invites, err := db.FindOrganizationInvitesByEmail(ctx, email)
+		require.NoError(t, err)
+		require.Len(t, invites, 1)
+		require.Equal(t, attributes, invites[0].Attributes)
+
+		invitesWithRole, err := db.FindOrganizationInvites(ctx, org.ID, "", 10)
+		require.NoError(t, err)
+		require.Len(t, invitesWithRole, 1)
+		require.Equal(t, attributes, invitesWithRole[0].Attributes)
+	})
+
+	t.Run("UpdateOrganizationInviteAttributes", func(t *testing.T) {
+		invite, err := db.FindOrganizationInvite(ctx, org.ID, email)
+		require.NoError(t, err)
+
+		updated := map[string]any{"attr1": "new-value1"}
+		require.NoError(t, db.UpdateOrganizationInviteAttributes(ctx, invite.ID, updated))
+
+		invite, err = db.FindOrganizationInvite(ctx, org.ID, email)
+		require.NoError(t, err)
+		require.Equal(t, updated, invite.Attributes)
+	})
+
+	t.Run("InsertOrganizationInvite without attributes normalizes to empty map", func(t *testing.T) {
+		email2 := "invitee-no-attrs@rilldata.com"
+		err := db.InsertOrganizationInvite(ctx, &database.InsertOrganizationInviteOptions{
+			Email:  email2,
+			OrgID:  org.ID,
+			RoleID: role.ID,
+		})
+		require.NoError(t, err)
+
+		invite, err := db.FindOrganizationInvite(ctx, org.ID, email2)
+		require.NoError(t, err)
+		require.NotNil(t, invite.Attributes)
+		require.Empty(t, invite.Attributes)
+	})
+
+	t.Run("InsertOrganizationInvite rejects invalid attributes", func(t *testing.T) {
+		err := db.InsertOrganizationInvite(ctx, &database.InsertOrganizationInviteOptions{
+			Email:      "invitee-invalid-attrs@rilldata.com",
+			OrgID:      org.ID,
+			RoleID:     role.ID,
+			Attributes: map[string]any{"invalid-key": "value"},
+		})
+		require.Error(t, err)
+		require.Contains(t, err.Error(), "invalid attribute key")
+	})
+
+	// Cleanup
+	require.NoError(t, db.DeleteOrganization(ctx, org.Name))
+}
+
 func testAttributeValidation(t *testing.T, db database.DB) {
 	ctx := context.Background()
 

admin/server/organizations.go

@@ -452,6 +452,11 @@ func (s *Server) AddOrganizationMemberUser(ctx context.Context, req *adminv1.Add
 		}
 	}
 
+	var attrs map[string]any
+	if req.Attributes != nil {
+		attrs = req.Attributes.AsMap()
+	}
+
 	user, err := s.admin.DB.FindUserByEmail(ctx, req.Email)
 	if err != nil {
 		if !errors.Is(err, database.ErrNotFound) {
@@ -460,10 +465,11 @@ func (s *Server) AddOrganizationMemberUser(ctx context.Context, req *adminv1.Add
 
 		// Invite user to join org
 		err := s.admin.DB.InsertOrganizationInvite(ctx, &database.InsertOrganizationInviteOptions{
-			Email:     req.Email,
-			InviterID: invitedByUserID,
-			OrgID:     org.ID,
-			RoleID:    role.ID,
+			Email:      req.Email,
+			InviterID:  invitedByUserID,
+			OrgID:      org.ID,
+			RoleID:     role.ID,
+			Attributes: attrs,
 		})
 		if err != nil {
 			if !errors.Is(err, database.ErrNotUnique) {
@@ -479,6 +485,13 @@ func (s *Server) AddOrganizationMemberUser(ctx context.Context, req *adminv1.Add
 			if err != nil {
 				return nil, err
 			}
+			// Update the invite's attributes only when explicitly provided, to avoid clearing them on a plain re-invite.
+			if req.Attributes != nil {
+				err = s.admin.DB.UpdateOrganizationInviteAttributes(ctx, invite.ID, attrs)
+				if err != nil {
+					return nil, err
+				}
+			}
 			// Fallthrough so we send the email again.
 		}
 
@@ -510,7 +523,7 @@ func (s *Server) AddOrganizationMemberUser(ctx context.Context, req *adminv1.Add
 	}
 
 	// Insert the user in the org and its managed usergroups transactionally.
-	err = s.admin.InsertOrganizationMemberUser(ctx, org.ID, user.ID, role.ID, nil, false)
+	err = s.admin.InsertOrganizationMemberUser(ctx, org.ID, user.ID, role.ID, attrs, false)
 	if err != nil {
 		if !errors.Is(err, database.ErrNotUnique) {
 			return nil, err
@@ -716,17 +729,29 @@ func (s *Server) UpdateOrganizationMemberUserAttributes(ctx context.Context, req
 		return nil, status.Error(codes.PermissionDenied, "not allowed to update org member attributes")
 	}
 
-	user, err := s.admin.DB.FindUserByEmail(ctx, req.Email)
-	if err != nil {
-		return nil, err
-	}
-
 	// Convert protobuf Struct to map[string]any
 	var attributes map[string]any
 	if req.Attributes != nil {
 		attributes = req.Attributes.AsMap()
 	}
 
+	user, err := s.admin.DB.FindUserByEmail(ctx, req.Email)
+	if err != nil {
+		if !errors.Is(err, database.ErrNotFound) {
+			return nil, err
+		}
+		// Check if there is a pending invite for this user
+		invite, err := s.admin.DB.FindOrganizationInvite(ctx, org.ID, req.Email)
+		if err != nil {
+			return nil, err
+		}
+		err = s.admin.DB.UpdateOrganizationInviteAttributes(ctx, invite.ID, attributes)
+		if err != nil {
+			return nil, err
+		}
+		return &adminv1.UpdateOrganizationMemberUserAttributesResponse{}, nil
+	}
+
 	// Update the attributes
 	_, err = s.admin.DB.UpdateOrganizationMemberUserAttributes(ctx, org.ID, user.ID, attributes)
 	if err != nil {