material_manage_system/internal/repository/user_organization.go

package repository

import (
	"database/sql"
	"encoding/json"
	"errors"
	"fmt"
	"sort"
	"strings"

	"github.com/wmms/wmms/internal/model"
)

var ErrInvalidOrganizations = errors.New("公司授权或默认公司无效，请选择启用且已授权的公司")

type userOrganizationState struct {
	UserID                int64   `json:"user_id"`
	OrganizationIDs       []int64 `json:"organization_ids"`
	DefaultOrganizationID *int64  `json:"default_organization_id"`
}

func replaceUserOrganizationsTx(tx *sql.Tx, userID int64, input model.UserOrganizationInput, creating bool) error {
	if input.OrganizationIDs == nil && len(input.DefaultOrganizationID) == 0 && !creating {
		return nil
	}
	ids := []int64{}
	if input.OrganizationIDs != nil {
		for _, id := range *input.OrganizationIDs {
			if id <= 0 {
				return ErrInvalidOrganizations
			}
			ids = append(ids, id)
		}
		ids = uniqueRoleIDs(ids)
		sort.Slice(ids, func(i, j int) bool { return ids[i] < ids[j] })
	}
	old := userOrganizationState{UserID: userID, OrganizationIDs: []int64{}}
	rows, err := tx.Query(`SELECT organization_id, is_default FROM user_organizations WHERE user_id = ? ORDER BY organization_id FOR UPDATE`, userID)
	if err != nil {
		return err
	}
	for rows.Next() {
		var id int64
		var def bool
		if err := rows.Scan(&id, &def); err != nil {
			rows.Close()
			return err
		}
		old.OrganizationIDs = append(old.OrganizationIDs, id)
		if def {
			value := id
			old.DefaultOrganizationID = &value
		}
	}
	err = rows.Err()
	rows.Close()
	if err != nil {
		return err
	}
	if input.OrganizationIDs == nil {
		ids = append(ids, old.OrganizationIDs...)
	}
	// Decode into an independent value so replacing the default cannot mutate
	// the original state retained for the audit record.
	var def *int64
	if old.DefaultOrganizationID != nil {
		value := *old.DefaultOrganizationID
		def = &value
	}
	if len(input.DefaultOrganizationID) > 0 {
		if err := json.Unmarshal(input.DefaultOrganizationID, &def); err != nil {
			return ErrInvalidOrganizations
		}
	}
	if def != nil && *def <= 0 {
		return ErrInvalidOrganizations
	}
	var admin bool
	if err := tx.QueryRow(`SELECT EXISTS(SELECT 1 FROM user_roles ur JOIN roles r ON r.id=ur.role_id WHERE ur.user_id=? AND r.code='super_admin' AND r.enabled=1)`, userID).Scan(&admin); err != nil {
		return err
	}
	contains := func(values []int64, id int64) bool {
		for _, value := range values {
			if value == id {
				return true
			}
		}
		return false
	}
	if !admin {
		if creating && (len(ids) == 0 || def == nil) {
			return ErrInvalidOrganizations
		}
		if len(ids) > 0 && def == nil {
			return ErrInvalidOrganizations
		}
		if def != nil && !contains(ids, *def) {
			return ErrInvalidOrganizations
		}
	} else if def != nil && !contains(ids, *def) {
		ids = append(ids, *def)
		sort.Slice(ids, func(i, j int) bool { return ids[i] < ids[j] })
	}
	if len(ids) > 0 {
		args := make([]interface{}, len(ids))
		for i, id := range ids {
			args[i] = id
		}
		placeholders := strings.TrimSuffix(strings.Repeat("?,", len(ids)), ",")
		rows, err := tx.Query(fmt.Sprintf(`SELECT id, status FROM organizations WHERE id IN (%s) ORDER BY id FOR UPDATE`, placeholders), args...)
		if err != nil {
			return err
		}
		count := 0
		invalid := false
		for rows.Next() {
			var id int64
			var status string
			if err := rows.Scan(&id, &status); err != nil {
				rows.Close()
				return err
			}
			count++
			if status != "active" && (!contains(old.OrganizationIDs, id) || (def != nil && *def == id)) {
				invalid = true
			}
		}
		err = rows.Err()
		rows.Close()
		if err != nil {
			return err
		}
		if invalid || count != len(ids) {
			return ErrInvalidOrganizations
		}
	}
	if _, err := tx.Exec(`DELETE FROM user_organizations WHERE user_id = ?`, userID); err != nil {
		return err
	}
	for _, id := range ids {
		if _, err := tx.Exec(`INSERT INTO user_organizations (user_id, organization_id, is_default) VALUES (?, ?, ?)`, userID, id, def != nil && *def == id); err != nil {
			return err
		}
	}
	next := userOrganizationState{UserID: userID, OrganizationIDs: ids, DefaultOrganizationID: def}
	oldJSON, err := json.Marshal(old)
	if err != nil {
		return err
	}
	newJSON, err := json.Marshal(next)
	if err != nil {
		return err
	}
	return writeOrganizationAudit(tx, &model.OrganizationAuditLog{ActorUserID: input.ActorUserID, Action: "organization_access_changed", OldValue: oldJSON, NewValue: newJSON})
}

func loadOrganizationsForUsers(q interface {
	Query(string, ...interface{}) (*sql.Rows, error)
}, users []*model.User) error {
	if len(users) == 0 {
		return nil
	}
	args := make([]interface{}, len(users))
	byID := map[int64]*model.User{}
	for i, user := range users {
		user.OrganizationIDs = []int64{}
		user.DefaultOrganizationID = nil
		args[i] = user.ID
		byID[user.ID] = user
	}
	placeholders := strings.TrimSuffix(strings.Repeat("?, ", len(users)), ", ")
	rows, err := q.Query(fmt.Sprintf(`SELECT user_id, organization_id, is_default FROM user_organizations WHERE user_id IN (%s) ORDER BY user_id, organization_id`, placeholders), args...)
	if err != nil {
		return err
	}
	defer rows.Close()
	for rows.Next() {
		var userID, id int64
		var def bool
		if err := rows.Scan(&userID, &id, &def); err != nil {
			return err
		}
		if user := byID[userID]; user != nil {
			user.OrganizationIDs = append(user.OrganizationIDs, id)
			if def {
				value := id
				user.DefaultOrganizationID = &value
			}
		}
	}
	return rows.Err()
}