Files
ai_site/platform/internal/orgunitstore/store.go
2026-07-31 10:19:22 +08:00

353 lines
9.1 KiB
Go
Raw Blame History

This file contains ambiguous Unicode characters
This file contains Unicode characters that might be confused with other characters. If you think that this is intentional, you can safely ignore this warning. Use the Escape button to reveal them.
package orgunitstore
import (
"context"
"database/sql"
"errors"
"fmt"
"strings"
"sync"
"time"
)
// MaxDepth 组织树最大深度(公司下第 1 级为 depth=1
const MaxDepth = 5
type OrgUnit struct {
OrgUnitID int64 `json:"org_unit_id"`
TenantID int64 `json:"tenant_id"`
ParentID int64 `json:"parent_id,omitempty"`
Name string `json:"name"`
Code string `json:"code,omitempty"`
Depth int `json:"depth"`
Path string `json:"path"`
CreatedAt time.Time `json:"created_at"`
}
type CreateInput struct {
ParentID int64
Name string
Code string
}
type UpdateInput struct {
Name *string
Code *string
}
type Store interface {
List(ctx context.Context, tenantID int64) ([]OrgUnit, error)
Get(ctx context.Context, tenantID, orgUnitID int64) (*OrgUnit, error)
Create(ctx context.Context, tenantID int64, in CreateInput) (*OrgUnit, error)
Update(ctx context.Context, tenantID, orgUnitID int64, in UpdateInput) (*OrgUnit, error)
Delete(ctx context.Context, tenantID, orgUnitID int64) error
// DescendantIDs 含自身。
DescendantIDs(ctx context.Context, tenantID, orgUnitID int64) ([]int64, error)
}
type MemoryStore struct {
mu sync.Mutex
byID map[int64]*OrgUnit
seq int64
}
func NewMemoryStore() *MemoryStore {
return &MemoryStore{byID: map[int64]*OrgUnit{}}
}
func (s *MemoryStore) List(_ context.Context, tenantID int64) ([]OrgUnit, error) {
s.mu.Lock()
defer s.mu.Unlock()
out := make([]OrgUnit, 0)
for _, o := range s.byID {
if o.TenantID == tenantID {
out = append(out, *o)
}
}
return out, nil
}
func (s *MemoryStore) Get(_ context.Context, tenantID, orgUnitID int64) (*OrgUnit, error) {
s.mu.Lock()
defer s.mu.Unlock()
o, ok := s.byID[orgUnitID]
if !ok || o.TenantID != tenantID {
return nil, fmt.Errorf("org unit not found")
}
cp := *o
return &cp, nil
}
func (s *MemoryStore) Create(_ context.Context, tenantID int64, in CreateInput) (*OrgUnit, error) {
s.mu.Lock()
defer s.mu.Unlock()
name := strings.TrimSpace(in.Name)
if name == "" {
return nil, fmt.Errorf("name required")
}
code := strings.TrimSpace(in.Code)
depth := 1
pathPrefix := ""
if in.ParentID > 0 {
p, ok := s.byID[in.ParentID]
if !ok || p.TenantID != tenantID {
return nil, fmt.Errorf("parent not found")
}
if p.Depth >= MaxDepth {
return nil, fmt.Errorf("max org depth is %d", MaxDepth)
}
depth = p.Depth + 1
pathPrefix = p.Path
}
if code != "" {
for _, o := range s.byID {
if o.TenantID == tenantID && o.Code == code {
return nil, fmt.Errorf("org code already exists")
}
}
}
s.seq++
id := s.seq
path := fmt.Sprintf("%s/%d", pathPrefix, id)
if pathPrefix == "" {
path = fmt.Sprintf("/%d", id)
}
o := &OrgUnit{
OrgUnitID: id,
TenantID: tenantID,
ParentID: in.ParentID,
Name: name,
Code: code,
Depth: depth,
Path: path,
CreatedAt: time.Now().UTC(),
}
s.byID[id] = o
cp := *o
return &cp, nil
}
func (s *MemoryStore) Update(_ context.Context, tenantID, orgUnitID int64, in UpdateInput) (*OrgUnit, error) {
s.mu.Lock()
defer s.mu.Unlock()
o, ok := s.byID[orgUnitID]
if !ok || o.TenantID != tenantID {
return nil, fmt.Errorf("org unit not found")
}
if in.Name != nil {
n := strings.TrimSpace(*in.Name)
if n == "" {
return nil, fmt.Errorf("name required")
}
o.Name = n
}
if in.Code != nil {
code := strings.TrimSpace(*in.Code)
for _, x := range s.byID {
if x.TenantID == tenantID && x.Code == code && x.OrgUnitID != orgUnitID {
return nil, fmt.Errorf("org code already exists")
}
}
o.Code = code
}
cp := *o
return &cp, nil
}
func (s *MemoryStore) Delete(_ context.Context, tenantID, orgUnitID int64) error {
s.mu.Lock()
defer s.mu.Unlock()
o, ok := s.byID[orgUnitID]
if !ok || o.TenantID != tenantID {
return fmt.Errorf("org unit not found")
}
for _, x := range s.byID {
if x.TenantID == tenantID && x.ParentID == orgUnitID {
return fmt.Errorf("org unit has children")
}
}
delete(s.byID, orgUnitID)
return nil
}
func (s *MemoryStore) DescendantIDs(_ context.Context, tenantID, orgUnitID int64) ([]int64, error) {
s.mu.Lock()
defer s.mu.Unlock()
o, ok := s.byID[orgUnitID]
if !ok || o.TenantID != tenantID {
return nil, fmt.Errorf("org unit not found")
}
out := []int64{orgUnitID}
prefix := o.Path + "/"
for _, x := range s.byID {
if x.TenantID == tenantID && strings.HasPrefix(x.Path, prefix) {
out = append(out, x.OrgUnitID)
}
}
return out, nil
}
type PostgresStore struct {
DB *sql.DB
}
func NewPostgresStore(db *sql.DB) *PostgresStore {
return &PostgresStore{DB: db}
}
func (s *PostgresStore) List(ctx context.Context, tenantID int64) ([]OrgUnit, error) {
rows, err := s.DB.QueryContext(ctx, `
SELECT org_unit_id, tenant_id, COALESCE(parent_id,0), name, COALESCE(code,''), depth, path, created_at
FROM platform_meta.org_units WHERE tenant_id=$1
ORDER BY path`, tenantID)
if err != nil {
return nil, err
}
defer rows.Close()
out := make([]OrgUnit, 0)
for rows.Next() {
var o OrgUnit
if err := rows.Scan(&o.OrgUnitID, &o.TenantID, &o.ParentID, &o.Name, &o.Code, &o.Depth, &o.Path, &o.CreatedAt); err != nil {
return nil, err
}
out = append(out, o)
}
return out, rows.Err()
}
func (s *PostgresStore) Get(ctx context.Context, tenantID, orgUnitID int64) (*OrgUnit, error) {
var o OrgUnit
err := s.DB.QueryRowContext(ctx, `
SELECT org_unit_id, tenant_id, COALESCE(parent_id,0), name, COALESCE(code,''), depth, path, created_at
FROM platform_meta.org_units WHERE org_unit_id=$1 AND tenant_id=$2`, orgUnitID, tenantID,
).Scan(&o.OrgUnitID, &o.TenantID, &o.ParentID, &o.Name, &o.Code, &o.Depth, &o.Path, &o.CreatedAt)
if errors.Is(err, sql.ErrNoRows) {
return nil, fmt.Errorf("org unit not found")
}
return &o, err
}
func (s *PostgresStore) Create(ctx context.Context, tenantID int64, in CreateInput) (*OrgUnit, error) {
name := strings.TrimSpace(in.Name)
if name == "" {
return nil, fmt.Errorf("name required")
}
code := strings.TrimSpace(in.Code)
depth := 1
parentPath := ""
var parentAny any
if in.ParentID > 0 {
p, err := s.Get(ctx, tenantID, in.ParentID)
if err != nil {
return nil, fmt.Errorf("parent not found")
}
if p.Depth >= MaxDepth {
return nil, fmt.Errorf("max org depth is %d", MaxDepth)
}
depth = p.Depth + 1
parentPath = p.Path
parentAny = in.ParentID
}
tx, err := s.DB.BeginTx(ctx, nil)
if err != nil {
return nil, err
}
defer func() { _ = tx.Rollback() }()
var id int64
err = tx.QueryRowContext(ctx, `
INSERT INTO platform_meta.org_units(tenant_id, parent_id, name, code, depth, path)
VALUES($1,$2,$3,$4,$5,'')
RETURNING org_unit_id`, tenantID, parentAny, name, code, depth).Scan(&id)
if err != nil {
if strings.Contains(err.Error(), "unique") || strings.Contains(err.Error(), "duplicate") {
return nil, fmt.Errorf("org code already exists")
}
return nil, err
}
path := fmt.Sprintf("/%d", id)
if parentPath != "" {
path = parentPath + "/" + fmt.Sprintf("%d", id)
}
if _, err := tx.ExecContext(ctx, `UPDATE platform_meta.org_units SET path=$1 WHERE org_unit_id=$2`, path, id); err != nil {
return nil, err
}
if err := tx.Commit(); err != nil {
return nil, err
}
return s.Get(ctx, tenantID, id)
}
func (s *PostgresStore) Update(ctx context.Context, tenantID, orgUnitID int64, in UpdateInput) (*OrgUnit, error) {
cur, err := s.Get(ctx, tenantID, orgUnitID)
if err != nil {
return nil, err
}
name, code := cur.Name, cur.Code
if in.Name != nil {
name = strings.TrimSpace(*in.Name)
if name == "" {
return nil, fmt.Errorf("name required")
}
}
if in.Code != nil {
code = strings.TrimSpace(*in.Code)
}
_, err = s.DB.ExecContext(ctx, `
UPDATE platform_meta.org_units SET name=$1, code=$2 WHERE org_unit_id=$3 AND tenant_id=$4`,
name, code, orgUnitID, tenantID)
if err != nil {
if strings.Contains(err.Error(), "unique") || strings.Contains(err.Error(), "duplicate") {
return nil, fmt.Errorf("org code already exists")
}
return nil, err
}
return s.Get(ctx, tenantID, orgUnitID)
}
func (s *PostgresStore) Delete(ctx context.Context, tenantID, orgUnitID int64) error {
var n int
if err := s.DB.QueryRowContext(ctx, `
SELECT COUNT(1) FROM platform_meta.org_units WHERE tenant_id=$1 AND parent_id=$2`, tenantID, orgUnitID).Scan(&n); err != nil {
return err
}
if n > 0 {
return fmt.Errorf("org unit has children")
}
res, err := s.DB.ExecContext(ctx, `
DELETE FROM platform_meta.org_units WHERE org_unit_id=$1 AND tenant_id=$2`, orgUnitID, tenantID)
if err != nil {
return err
}
aff, _ := res.RowsAffected()
if aff == 0 {
return fmt.Errorf("org unit not found")
}
return nil
}
func (s *PostgresStore) DescendantIDs(ctx context.Context, tenantID, orgUnitID int64) ([]int64, error) {
o, err := s.Get(ctx, tenantID, orgUnitID)
if err != nil {
return nil, err
}
rows, err := s.DB.QueryContext(ctx, `
SELECT org_unit_id FROM platform_meta.org_units
WHERE tenant_id=$1 AND (org_unit_id=$2 OR path LIKE $3)
ORDER BY path`, tenantID, orgUnitID, o.Path+"/%")
if err != nil {
return nil, err
}
defer rows.Close()
out := make([]int64, 0)
for rows.Next() {
var id int64
if err := rows.Scan(&id); err != nil {
return nil, err
}
out = append(out, id)
}
return out, rows.Err()
}