MyAiDesk/GovAI/server/pkg/tenant/tenant.go

package tenant

import (
	"context"
	"database/sql"

	"github.com/google/uuid"
)

// contextKey is the type for context values used by this package.
type contextKey string

const (
	orgIDKey  contextKey = "org_id"
	userIDKey contextKey = "user_id"
)

// WithOrgID stores the organization ID in context.
func WithOrgID(ctx context.Context, orgID string) context.Context {
	return context.WithValue(ctx, orgIDKey, orgID)
}

// GetOrgID retrieves the organization ID from context.
// Returns empty string if not set.
func GetOrgID(ctx context.Context) string {
	if v := ctx.Value(orgIDKey); v != nil {
		return v.(string)
	}
	return ""
}

// WithUserID stores the user ID in context.
func WithUserID(ctx context.Context, userID uuid.UUID) context.Context {
	return context.WithValue(ctx, userIDKey, userID)
}

// GetUserID retrieves the user ID from context.
func GetUserID(ctx context.Context) uuid.UUID {
	if v := ctx.Value(userIDKey); v != nil {
		return v.(uuid.UUID)
	}
	return uuid.Nil
}

// GetUserOrgID queries the database for the current user's org_id.
// This is the standard way to get the caller's organization.
func GetUserOrgID(ctx context.Context, pool interface {
	QueryRowContext(ctx context.Context, query string, args ...any) *sql.Row
}, userID uuid.UUID) (string, error) {
	var orgID string
	err := pool.QueryRowContext(ctx,
		`SELECT COALESCE(org_id::text, '') FROM users WHERE id = $1`,
		userID,
	).Scan(&orgID)
	return orgID, err
}

// IsSuperAdmin checks if the given role is platform super admin.
func IsSuperAdmin(role string) bool {
	return role == "super_admin"
}

// IsAdmin checks if the given role is at least org-level admin.
func IsAdmin(role string) bool {
	return role == "admin" || role == "super_admin"
}