llmgw/internal/store/hierarchy.go

package store

import (
	"context"
	"database/sql"
	"errors"
	"fmt"
	"strings"
	"time"
)

// ErrNotEmpty means something still hangs off what was to be deleted.
var ErrNotEmpty = errors.New("not empty")

// TreeOrg, TreeTeam and TreeProject are the tenancy tree as one value, with
// what hangs off each project counted, for a page or a command that shows it
// whole.
type TreeOrg struct {
	ID, Name string
	// Members counts people who are members of the organization itself.
	Members int
	Teams   []TreeTeam
}

// TreeTeam is a team and its projects.
type TreeTeam struct {
	ID, Name string
	// Members counts people who are members of the team itself.
	Members  int
	Projects []TreeProject
}

// TreeProject is a project and what uses it.
type TreeProject struct {
	ID, Name string
	// Users counts the project's own members, whatever their state; people
	// who reach it through its team or organization are counted there.
	Users int
	// Keys counts every key ever issued in it; ActiveKeys those not revoked.
	Keys, ActiveKeys int
}

// Tree reads the whole tenancy tree, organizations with no teams and teams
// with no projects included.
func (s *DB) Tree(ctx context.Context) ([]TreeOrg, error) {
	orgs, err := s.AllOrgs(ctx)
	if err != nil {
		return nil, err
	}
	teams, err := s.AllTeams(ctx)
	if err != nil {
		return nil, err
	}
	type proj struct {
		TreeProject
		teamID string
	}
	projects, err := queryAll(ctx, s.db, `
		SELECT p.id, p.name, p.team_id,
		  (SELECT COUNT(*) FROM memberships m WHERE m.scope_type = 'project' AND m.scope_id = p.id),
		  (SELECT COUNT(*) FROM api_keys k WHERE k.project_id = p.id),
		  (SELECT COUNT(*) FROM api_keys k WHERE k.project_id = p.id AND k.revoked_at IS NULL)
		FROM projects p ORDER BY p.name`, nil, func(rows *sql.Rows) (proj, error) {
		var p proj
		return p, rows.Scan(&p.ID, &p.Name, &p.teamID, &p.Users, &p.Keys, &p.ActiveKeys)
	})
	if err != nil {
		return nil, err
	}

	byTeam := map[string][]TreeProject{}
	for _, p := range projects {
		byTeam[p.teamID] = append(byTeam[p.teamID], p.TreeProject)
	}
	members := map[string]int{}
	counts, err := queryAll(ctx, s.db,
		`SELECT scope_type || ':' || scope_id, COUNT(*) FROM memberships
		 WHERE scope_type IN ('org', 'team') GROUP BY scope_type, scope_id`, nil,
		func(rows *sql.Rows) (struct {
			k string
			n int
		}, error) {
			var c struct {
				k string
				n int
			}
			return c, rows.Scan(&c.k, &c.n)
		})
	if err != nil {
		return nil, err
	}
	for _, c := range counts {
		members[c.k] = c.n
	}
	byOrg := map[string][]TreeTeam{}
	for _, t := range teams {
		byOrg[t.OrgID] = append(byOrg[t.OrgID], TreeTeam{ID: t.ID, Name: t.Name,
			Members: members["team:"+t.ID], Projects: byTeam[t.ID]})
	}
	out := make([]TreeOrg, 0, len(orgs))
	for _, o := range orgs {
		out = append(out, TreeOrg{ID: o.ID, Name: o.Name, Members: members["org:"+o.ID], Teams: byOrg[o.ID]})
	}
	return out, nil
}

// RenameOrg renames an organization. Names are unique among organizations.
func (s *DB) RenameOrg(ctx context.Context, id, name string) error {
	name = strings.TrimSpace(name)
	if name == "" {
		return errors.New("a name cannot be empty")
	}
	if err := s.nameIsFree(ctx,
		`SELECT COUNT(*) FROM organizations WHERE name = ? AND id <> ?`, name, id); err != nil {
		return err
	}
	return s.updateOne(ctx, `UPDATE organizations SET name = ? WHERE id = ?`, name, id)
}

// RenameTeam renames a team. Names are unique within an organization.
func (s *DB) RenameTeam(ctx context.Context, id, name string) error {
	name = strings.TrimSpace(name)
	if name == "" {
		return errors.New("a name cannot be empty")
	}
	if err := s.nameIsFree(ctx, `SELECT COUNT(*) FROM teams
		WHERE name = ? AND id <> ? AND org_id = (SELECT org_id FROM teams WHERE id = ?)`,
		name, id, id); err != nil {
		return err
	}
	return s.updateOne(ctx, `UPDATE teams SET name = ? WHERE id = ?`, name, id)
}

// RenameProject renames a project. Names are unique within a team.
func (s *DB) RenameProject(ctx context.Context, id, name string) error {
	name = strings.TrimSpace(name)
	if name == "" {
		return errors.New("a name cannot be empty")
	}
	if err := s.nameIsFree(ctx, `SELECT COUNT(*) FROM projects
		WHERE name = ? AND id <> ? AND team_id = (SELECT team_id FROM projects WHERE id = ?)`,
		name, id, id); err != nil {
		return err
	}
	return s.updateOne(ctx, `UPDATE projects SET name = ? WHERE id = ?`, name, id)
}

// MoveTeam puts a team, and everything under it, in another organization.
//
// Its limits go with it, and from now on it inherits the new organization's.
func (s *DB) MoveTeam(ctx context.Context, teamID, orgID string) error {
	if err := s.exists(ctx, "organization", `SELECT COUNT(*) FROM organizations WHERE id = ?`, orgID); err != nil {
		return err
	}
	if err := s.nameIsFree(ctx, `SELECT COUNT(*) FROM teams
		WHERE org_id = ? AND id <> ? AND name = (SELECT name FROM teams WHERE id = ?)`,
		orgID, teamID, teamID); err != nil {
		return err
	}
	return s.updateOne(ctx, `UPDATE teams SET org_id = ? WHERE id = ?`, orgID, teamID)
}

// MoveProject puts a project, with its people and keys, in another team.
func (s *DB) MoveProject(ctx context.Context, projectID, teamID string) error {
	if err := s.exists(ctx, "team", `SELECT COUNT(*) FROM teams WHERE id = ?`, teamID); err != nil {
		return err
	}
	if err := s.nameIsFree(ctx, `SELECT COUNT(*) FROM projects
		WHERE team_id = ? AND id <> ? AND name = (SELECT name FROM projects WHERE id = ?)`,
		teamID, projectID, projectID); err != nil {
		return err
	}
	return s.updateOne(ctx, `UPDATE projects SET team_id = ? WHERE id = ?`, teamID, projectID)
}

// DeleteOrg removes an organization with no teams, and its limits.
func (s *DB) DeleteOrg(ctx context.Context, id string) error {
	return s.deleteEmpty(ctx, ScopeOrg, id, `DELETE FROM organizations WHERE id = ?`,
		emptyCheck{`SELECT COUNT(*) FROM teams WHERE org_id = ?`, "team"},
		emptyCheck{`SELECT COUNT(*) FROM memberships WHERE scope_type = 'org' AND scope_id = ?`, "member"})
}

// DeleteTeam removes a team with no projects, and its limits.
func (s *DB) DeleteTeam(ctx context.Context, id string) error {
	return s.deleteEmpty(ctx, ScopeTeam, id, `DELETE FROM teams WHERE id = ?`,
		emptyCheck{`SELECT COUNT(*) FROM projects WHERE team_id = ?`, "project"},
		emptyCheck{`SELECT COUNT(*) FROM memberships WHERE scope_type = 'team' AND scope_id = ?`, "member"})
}

// DeleteProject removes a project nobody is in and no key was ever issued in.
//
// A revoked key still counts. Usage history names its key and project, and a
// project deleted from under that history leaves rows that belong to nothing;
// such a project can be renamed or left, not removed.
func (s *DB) DeleteProject(ctx context.Context, id string) error {
	return s.deleteEmpty(ctx, ScopeProject, id, `DELETE FROM projects WHERE id = ?`,
		emptyCheck{`SELECT COUNT(*) FROM memberships WHERE scope_type = 'project' AND scope_id = ?`, "member"},
		emptyCheck{`SELECT COUNT(*) FROM api_keys WHERE project_id = ?`, "key"})
}

type emptyCheck struct {
	query, noun string
}

func (s *DB) deleteEmpty(ctx context.Context, scope Scope, id, del string, checks ...emptyCheck) error {
	return s.inTx(ctx, func(c conn) error {
		var busy []string
		for _, ch := range checks {
			var n int
			if err := c.QueryRowContext(ctx, ch.query, id).Scan(&n); err != nil {
				return err
			}
			if n > 0 {
				busy = append(busy, plural(n, ch.noun))
			}
		}
		// Not an emptyCheck: the table is there only on PostgreSQL.
		if n, err := s.kbOwnedBy(ctx, c, scope, id); err != nil {
			return err
		} else if n > 0 {
			busy = append(busy, plural(n, "collection"))
		}
		if len(busy) > 0 {
			return fmt.Errorf("%w: it still has %s", ErrNotEmpty, strings.Join(busy, " and "))
		}
		if _, err := c.ExecContext(ctx,
			`DELETE FROM limits WHERE scope_type IN (?, ?) AND scope_id = ?`,
			string(scope), string(MemberScope(scope)), id); err != nil {
			return err
		}
		if err := s.kbForgetScope(ctx, c, scope, id); err != nil {
			return err
		}
		res, err := c.ExecContext(ctx, del, id)
		if err != nil {
			return err
		}
		if n, _ := res.RowsAffected(); n == 0 {
			return ErrNotFound
		}
		return s.kbDropPartition(ctx, c, scope, id)
	})
}

// DeleteUser removes an account: its sessions end, its keys are revoked, and
// the row goes.
//
// The keys stay, revoked and without an owner, because usage history refers
// to them; the account itself is what the person asked to have removed.
func (s *DB) DeleteUser(ctx context.Context, userID string, at time.Time) error {
	return s.inTx(ctx, func(c conn) error {
		if _, err := c.ExecContext(ctx, `DELETE FROM sessions WHERE user_id = ?`, userID); err != nil {
			return err
		}
		if _, err := c.ExecContext(ctx, `UPDATE api_keys
			SET revoked_at = COALESCE(revoked_at, ?), user_id = NULL WHERE user_id = ?`,
			millis(at), userID); err != nil {
			return err
		}
		// A person's own collections go with the account.
		if err := s.kbForgetScope(ctx, c, ScopeUser, userID); err != nil {
			return err
		}
		res, err := c.ExecContext(ctx, `DELETE FROM users WHERE id = ?`, userID)
		if err != nil {
			return err
		}
		if n, _ := res.RowsAffected(); n == 0 {
			return ErrNotFound
		}
		return nil
	})
}

// KeyRow is a key with where it sits and who made it, for an administrator's
// list of every key.
type KeyRow struct {
	Key
	ProjectPath string
	// OwnerEmail is empty for a key an operator issued in the project.
	OwnerEmail string
}

// AllKeys lists every key, newest first, with its project's path and owner.
func (s *DB) AllKeys(ctx context.Context) ([]KeyRow, error) {
	cols := make([]string, 0, 12)
	for _, c := range strings.Split(keyColumns, ",") {
		cols = append(cols, "k."+strings.TrimSpace(c))
	}
	return queryAll(ctx, s.db, `
		SELECT `+strings.Join(cols, ", ")+`,
		  COALESCE(o.name || ' / ' || t.name || ' / ' || p.name, ''), COALESCE(u.email, '')
		FROM api_keys k
		LEFT JOIN projects p ON p.id = k.project_id
		LEFT JOIN teams t ON t.id = p.team_id
		LEFT JOIN organizations o ON o.id = t.org_id
		LEFT JOIN users u ON u.id = k.user_id
		ORDER BY k.created_at DESC`, nil, func(rows *sql.Rows) (KeyRow, error) {
		var r KeyRow
		k, err := scanKey(keyAndMore{rows, []any{&r.ProjectPath, &r.OwnerEmail}})
		r.Key = k
		return r, err
	})
}

// keyAndMore lets scanKey read a row that carries columns after the key's.
type keyAndMore struct {
	sc   scanner
	more []any
}

func (k keyAndMore) Scan(dest ...any) error { return k.sc.Scan(append(dest, k.more...)...) }

func (s *DB) updateOne(ctx context.Context, q string, args ...any) error {
	res, err := s.db.ExecContext(ctx, q, args...)
	if err != nil {
		return err
	}
	if n, _ := res.RowsAffected(); n == 0 {
		return ErrNotFound
	}
	return nil
}

func (s *DB) exists(ctx context.Context, noun, q string, id string) error {
	var n int
	if err := s.db.QueryRowContext(ctx, q, id).Scan(&n); err != nil {
		return err
	}
	if n == 0 {
		return fmt.Errorf("no %s %q", noun, id)
	}
	return nil
}

func plural(n int, noun string) string {
	if n == 1 {
		return "1 " + noun
	}
	return fmt.Sprintf("%d %ss", n, noun)
}