llmgw/internal/store/tenancy.go

package store

import (
	"context"
	"crypto/rand"
	"database/sql"
	"encoding/base32"
	"errors"
	"fmt"
	"strings"
)

// ErrDuplicateName means a sibling of that name is already there.
var ErrDuplicateName = errors.New("something of that name is already here")

// NewID mints an identifier with a readable prefix.
//
// Random rather than a slug of the name, which is what the tenancy tree used
// and what made two organizations impossible to have: team ids were slugs in
// one flat namespace, PutTeam upserts on id, so creating "platform" under a
// second organization moved the first one's team into it. Nothing failed. The
// team simply changed owner.
func NewID(prefix string) string {
	var buf [10]byte
	_, _ = rand.Read(buf[:])
	return prefix + strings.ToLower(
		base32.StdEncoding.WithPadding(base32.NoPadding).EncodeToString(buf[:]))
}

// TeamPath is a team with the organization it belongs to, for a chooser.
type TeamPath struct {
	ID, Name, OrgID, OrgName string
}

// AllOrgs lists the organizations, by name.
func (s *DB) AllOrgs(ctx context.Context) ([]Org, error) {
	return queryAll(ctx, s.db, `SELECT id, name FROM organizations ORDER BY name`,
		nil, func(rows *sql.Rows) (Org, error) {
			var o Org
			return o, rows.Scan(&o.ID, &o.Name)
		})
}

// AllTeams lists the teams with their organization.
func (s *DB) AllTeams(ctx context.Context) ([]TeamPath, error) {
	return queryAll(ctx, s.db, `
		SELECT t.id, t.name, o.id, o.name
		FROM teams t JOIN organizations o ON o.id = t.org_id
		ORDER BY o.name, t.name`, nil, func(rows *sql.Rows) (TeamPath, error) {
		var t TeamPath
		return t, rows.Scan(&t.ID, &t.Name, &t.OrgID, &t.OrgName)
	})
}

// CreateOrg adds an organization and returns its id.
//
// Create rather than Put. The Put* calls upsert by id so `keys init` can be run
// again without duplicating, which is right for provisioning and wrong for a
// person filling in a form: there, a name that already exists is a mistake
// being made, not a line being re-run.
func (s *DB) CreateOrg(ctx context.Context, name string) (string, error) {
	name = strings.TrimSpace(name)
	if name == "" {
		return "", fmt.Errorf("an organization needs a name")
	}
	if err := s.nameIsFree(ctx,
		`SELECT COUNT(*) FROM organizations WHERE name = ?`, name); err != nil {
		return "", err
	}
	id := NewID("org-")
	return id, s.PutOrg(ctx, Org{ID: id, Name: name})
}

// CreateTeam adds a team to an organization.
//
// Unique within its organization, not globally: two companies both having a
// "platform" team is ordinary, and refusing the second would be the flat
// namespace back again wearing a better error message.
func (s *DB) CreateTeam(ctx context.Context, orgID, name string) (string, error) {
	name = strings.TrimSpace(name)
	if orgID == "" || name == "" {
		return "", fmt.Errorf("a team needs an organization and a name")
	}
	if err := s.nameIsFree(ctx,
		`SELECT COUNT(*) FROM teams WHERE org_id = ? AND name = ?`, orgID, name); err != nil {
		return "", err
	}
	id := NewID("team-")
	return id, s.PutTeam(ctx, Team{ID: id, OrgID: orgID, Name: name})
}

// CreateProject adds a project to a team.
func (s *DB) CreateProject(ctx context.Context, teamID, name string) (string, error) {
	name = strings.TrimSpace(name)
	if teamID == "" || name == "" {
		return "", fmt.Errorf("a project needs a team and a name")
	}
	if err := s.nameIsFree(ctx,
		`SELECT COUNT(*) FROM projects WHERE team_id = ? AND name = ?`, teamID, name); err != nil {
		return "", err
	}
	id := NewID("proj-")
	return id, s.PutProject(ctx, Project{ID: id, TeamID: teamID, Name: name})
}

// nameIsFree refuses a duplicate among siblings.
//
// Checked rather than left to a constraint because there is no unique index to
// lean on: names were never unique, the ids were, and the ids were slugs of the
// names. A query and a race is the honest state of it — two admins creating the
// same team in the same second get two teams, which is visible and fixable,
// where the old behaviour silently reparented somebody else's.
func (s *DB) nameIsFree(ctx context.Context, q string, args ...any) error {
	var n int
	if err := s.db.QueryRowContext(ctx, q, args...).Scan(&n); err != nil {
		return err
	}
	if n > 0 {
		return ErrDuplicateName
	}
	return nil
}

// AllLimits reads every limit that has been set, keyed by scope and id.
//
// All of them in one read rather than one query per group: the quota page shows
// the whole tree, and a row for each of an organization's projects would
// otherwise be a query each.
func (s *DB) AllLimits(ctx context.Context) (map[string]Limits, error) {
	rows, err := queryAll(ctx, s.db, `
		SELECT scope_type, scope_id, rpm, tpm, budget_micros, budget_period, allowed_models, allowed_tools
		FROM limits`, nil, func(rows *sql.Rows) (Limits, error) {
		var l Limits
		var scope string
		var models, tools sql.NullString
		var rpm, tpm sql.NullInt64
		var budget sql.NullInt64
		var period sql.NullString
		if err := rows.Scan(&scope, &l.ScopeID, &rpm, &tpm, &budget, &period, &models, &tools); err != nil {
			return l, err
		}
		l.Scope = Scope(scope)
		if rpm.Valid {
			v := int(rpm.Int64)
			l.RPM = &v
		}
		if tpm.Valid {
			v := int(tpm.Int64)
			l.TPM = &v
		}
		if budget.Valid {
			v := budget.Int64
			l.Budget = &v
		}
		l.BudgetPeriod = Period(period.String)
		am, err := decodeModels(models)
		l.AllowedModels = am
		if err != nil {
			return l, err
		}
		l.AllowedTools, err = decodeModels(tools)
		return l, err
	})
	if err != nil {
		return nil, err
	}
	out := make(map[string]Limits, len(rows))
	for _, l := range rows {
		out[LimitOwner{Scope: l.Scope, ID: l.ScopeID}.CounterKey()] = l
	}
	return out, nil
}

// ClearLimits removes every limit at one scope.
//
// Distinct from setting them all to nothing, which is what an empty form would
// otherwise write: a row of NULLs and no row mean the same to resolution, but
// only one of them says "nobody has set a limit here" when somebody reads the
// table directly.
func (s *DB) ClearLimits(ctx context.Context, scope Scope, id string) error {
	_, err := s.db.ExecContext(ctx,
		`DELETE FROM limits WHERE scope_type = ? AND scope_id = ?`, string(scope), id)
	return err
}