Database isolation: Balancer (#2407)

* Database isolation: Balancer

Remove naive Balancer

remove debugging lines

Thread dax.Transaction through Controller

Change role to roleType

Swap out Balancer interface with new one

Standardize InvalidTransaction error

Add some interface comments

* Remove type.Worker; replace with type.Address

* Remove database validate from Queryer

This is already being handled in the `CreateTable()` method. Prior
to doing that validation, we were getting a panic, but that's no longer
the case.

* Remove dax.TableQualifier; replace with dax.QualifiedDatabaseID

* Update IDK test to create database

(cherry picked from commit d971cfc269)

api.go

@@ -3354,6 +3354,9 @@ type QueryAPI interface {
 	Query(ctx context.Context, req *QueryRequest) (QueryResponse, error)
 }
 
+// Ensure type implements interface.
+var _ SystemAPI = (*FeatureBaseSystemAPI)(nil)
+
 // FeatureBaseSystemAPI is a wrapper around pilosa.API. It implements the
 // SystemAPI interface
 type FeatureBaseSystemAPI struct {
@@ -3432,3 +3435,54 @@ func (fsapi *FeatureBaseSystemAPI) ClusterNodes() []ClusterNode {
 
 	return result
 }
+
+// Ensure type implements interface.
+var _ SystemAPI = (*NopSystemAPI)(nil)
+
+// NopSystemAPI is a no-op implementation of the SystemAPI.
+type NopSystemAPI struct{}
+
+func (napi *NopSystemAPI) ClusterName() string {
+	return ""
+}
+
+func (napi *NopSystemAPI) Version() string {
+	return ""
+}
+
+func (napi *NopSystemAPI) PlatformDescription() string {
+	return ""
+}
+
+func (napi *NopSystemAPI) PlatformVersion() string {
+	return ""
+}
+
+func (napi *NopSystemAPI) ClusterNodeCount() int {
+	return 0
+}
+
+func (napi *NopSystemAPI) ClusterReplicaCount() int {
+	return 0
+}
+
+func (napi *NopSystemAPI) ShardWidth() int {
+	return 0
+}
+
+func (napi *NopSystemAPI) ClusterState() string {
+	return ""
+}
+
+func (napi *NopSystemAPI) DataDir() string {
+	return ""
+}
+
+func (napi *NopSystemAPI) NodeID() string {
+	return ""
+}
+
+func (napi *NopSystemAPI) ClusterNodes() []ClusterNode {
+	result := make([]ClusterNode, 0)
+	return result
+}

api_directive_test.go

@@ -20,10 +20,10 @@ func TestAPI_Directive(t *testing.T) {
 	api := c.GetPrimary().API
 	ctx := context.Background()
 
-	qual := dax.NewTableQualifier("acme", "db1")
-	tbl1 := daxtest.TestQualifiedTableWithID(t, qual, "1", "tbl1", 12, false)
-	tbl2 := daxtest.TestQualifiedTableWithID(t, qual, "2", "tbl2", 12, false)
-	tbl3 := daxtest.TestQualifiedTableWithID(t, qual, "3", "tbl3", 12, false)
+	qdbid := dax.NewQualifiedDatabaseID("acme", "db1")
+	tbl1 := daxtest.TestQualifiedTableWithID(t, qdbid, "1", "tbl1", 12, false)
+	tbl2 := daxtest.TestQualifiedTableWithID(t, qdbid, "2", "tbl2", 12, false)
+	tbl3 := daxtest.TestQualifiedTableWithID(t, qdbid, "3", "tbl3", 12, false)
 
 	t.Run("Schema", func(t *testing.T) {
 

dax/address.go

@@ -148,6 +148,13 @@ func (a addr) hostPortPath() string {
 	return ret
 }
 
+// Addresses is a sortable slice of Address.
+type Addresses []Address
+
+func (a Addresses) Len() int           { return len(a) }
+func (a Addresses) Less(i, j int) bool { return a[i] < a[j] }
+func (a Addresses) Swap(i, j int)      { a[i], a[j] = a[j], a[i] }
+
 // AddressManager is an interface for any service which needs to maintain a list
 // of addresses, and receive add/remove address requests from other services.
 type AddressManager interface {

dax/boltdb/boltdb.go

@@ -125,6 +125,7 @@ func (db *DB) Close() (err error) {
 // BeginTx starts a transaction and returns a wrapper Tx type. This type
 // provides a reference to the database and a fixed timestamp at the start of
 // the transaction. The timestamp allows us to mock time during tests as well.
+// The wrapper also contains the context.
 func (db *DB) BeginTx(ctx context.Context, writable bool) (*Tx, error) {
 	tx, err := db.db.Begin(writable)
 	if err != nil {
@@ -134,6 +135,7 @@ func (db *DB) BeginTx(ctx context.Context, writable bool) (*Tx, error) {
 	// Return wrapper Tx that includes the transaction start time.
 	return &Tx{
 		Tx:  tx,
+		ctx: ctx,
 		db:  db,
 		now: db.Now().UTC().Truncate(time.Second),
 	}, nil
@@ -142,10 +144,15 @@ func (db *DB) BeginTx(ctx context.Context, writable bool) (*Tx, error) {
 // Tx wraps the SQL Tx object to provide a timestamp at the start of the transaction.
 type Tx struct {
 	*bolt.Tx
+	ctx context.Context
 	db  *DB
 	now time.Time
 }
 
+func (tx *Tx) Context() context.Context {
+	return tx.ctx
+}
+
 func (db *DB) Path() string {
 	return db.filePath
 }

dax/boltdb/directiveversion.go

@@ -1,7 +1,6 @@
 package boltdb
 
 import (
-	"context"
 	"encoding/binary"
 
 	"github.com/featurebasedb/featurebase/v3/dax"
@@ -32,14 +31,13 @@ func NewDirectiveVersion(db *DB) *DirectiveVersion {
 	}
 }
 
-func (d *DirectiveVersion) Increment(ctx context.Context, delta uint64) (uint64, error) {
-	tx, err := d.db.BeginTx(ctx, true)
-	if err != nil {
-		return 0, errors.Wrap(err, "getting transaction")
+func (d *DirectiveVersion) Increment(tx dax.Transaction, delta uint64) (uint64, error) {
+	txx, ok := tx.(*Tx)
+	if !ok {
+		return 0, dax.NewErrInvalidTransaction()
 	}
-	defer tx.Rollback()
 
-	bkt := tx.Bucket(bucketDirective)
+	bkt := txx.Bucket(bucketDirective)
 	if bkt == nil {
 		return 0, errors.Errorf(ErrFmtBucketNotFound, bucketDirective)
 	}
@@ -58,9 +56,5 @@ func (d *DirectiveVersion) Increment(ctx context.Context, delta uint64) (uint64,
 		return 0, errors.Wrap(err, "putting next directive version")
 	}
 
-	if err := tx.Commit(); err != nil {
-		return 0, err
-	}
-
 	return nextVersion, nil
 }

dax/boltdb/node.go

@@ -2,7 +2,6 @@ package boltdb
 
 import (
 	"bytes"
-	"context"
 	"encoding/json"
 	"fmt"
 
@@ -39,14 +38,13 @@ func NewNodeService(db *DB, logger logger.Logger) *NodeService {
 	}
 }
 
-func (s *NodeService) CreateNode(ctx context.Context, addr dax.Address, node *dax.Node) error {
-	tx, err := s.db.BeginTx(ctx, true)
-	if err != nil {
-		return errors.Wrap(err, "getting transaction")
+func (s *NodeService) CreateNode(tx dax.Transaction, addr dax.Address, node *dax.Node) error {
+	txx, ok := tx.(*Tx)
+	if !ok {
+		return dax.NewErrInvalidTransaction()
 	}
-	defer tx.Rollback()
 
-	bkt := tx.Bucket(bucketNodes)
+	bkt := txx.Bucket(bucketNodes)
 	if bkt == nil {
 		return errors.Errorf(ErrFmtBucketNotFound, bucketNodes)
 	}
@@ -60,17 +58,16 @@ func (s *NodeService) CreateNode(ctx context.Context, addr dax.Address, node *da
 		return errors.Wrap(err, "putting node")
 	}
 
-	return tx.Commit()
+	return nil
 }
 
-func (s *NodeService) ReadNode(ctx context.Context, addr dax.Address) (*dax.Node, error) {
-	tx, err := s.db.BeginTx(ctx, false)
-	if err != nil {
-		return nil, errors.Wrap(err, "beginning tx")
+func (s *NodeService) ReadNode(tx dax.Transaction, addr dax.Address) (*dax.Node, error) {
+	txx, ok := tx.(*Tx)
+	if !ok {
+		return nil, dax.NewErrInvalidTransaction()
 	}
-	defer tx.Rollback()
 
-	bkt := tx.Bucket(bucketNodes)
+	bkt := txx.Bucket(bucketNodes)
 	if bkt == nil {
 		return nil, errors.Errorf(ErrFmtBucketNotFound, bucketNodes)
 	}
@@ -88,14 +85,13 @@ func (s *NodeService) ReadNode(ctx context.Context, addr dax.Address) (*dax.Node
 	return node, nil
 }
 
-func (s *NodeService) DeleteNode(ctx context.Context, addr dax.Address) error {
-	tx, err := s.db.BeginTx(ctx, true)
-	if err != nil {
-		return errors.Wrap(err, "beginning tx")
+func (s *NodeService) DeleteNode(tx dax.Transaction, addr dax.Address) error {
+	txx, ok := tx.(*Tx)
+	if !ok {
+		return dax.NewErrInvalidTransaction()
 	}
-	defer tx.Rollback()
 
-	bkt := tx.Bucket(bucketNodes)
+	bkt := txx.Bucket(bucketNodes)
 	if bkt == nil {
 		return errors.Errorf(ErrFmtBucketNotFound, bucketNodes)
 	}
@@ -104,17 +100,16 @@ func (s *NodeService) DeleteNode(ctx context.Context, addr dax.Address) error {
 		return errors.Wrapf(err, "deleting node key: %s", addressKey(addr))
 	}
 
-	return tx.Commit()
+	return nil
 }
 
-func (s *NodeService) Nodes(ctx context.Context) ([]*dax.Node, error) {
-	tx, err := s.db.BeginTx(ctx, false)
-	if err != nil {
-		return nil, errors.Wrap(err, "getting tx")
+func (s *NodeService) Nodes(tx dax.Transaction) ([]*dax.Node, error) {
+	txx, ok := tx.(*Tx)
+	if !ok {
+		return nil, dax.NewErrInvalidTransaction()
 	}
-	defer tx.Rollback()
 
-	nodes, err := s.getNodes(ctx, tx)
+	nodes, err := s.getNodes(txx)
 	if err != nil {
 		return nil, errors.Wrap(err, "getting nodes")
 	}
@@ -122,7 +117,7 @@ func (s *NodeService) Nodes(ctx context.Context) ([]*dax.Node, error) {
 	return nodes, nil
 }
 
-func (s *NodeService) getNodes(ctx context.Context, tx *Tx) ([]*dax.Node, error) {
+func (s *NodeService) getNodes(tx *Tx) ([]*dax.Node, error) {
 	c := tx.Bucket(bucketNodes).Cursor()
 
 	// Deserialize rows into Node objects.

dax/boltdb/node_test.go

@@ -35,21 +35,27 @@ func TestNodeService(t *testing.T) {
 			},
 		}
 
+		tx, err := db.BeginTx(ctx, true)
+		assert.NoError(t, err)
+		defer tx.Rollback()
+
 		// Create node.
-		assert.NoError(t, ns.CreateNode(ctx, node1.Address, node1))
+		assert.NoError(t, ns.CreateNode(tx, node1.Address, node1))
 
 		// Read node.
-		n, err := ns.ReadNode(ctx, node1.Address)
+		n, err := ns.ReadNode(tx, node1.Address)
 		assert.NoError(t, err)
 		assert.Equal(t, node1, n)
 
 		// Delete node.
-		assert.NoError(t, ns.DeleteNode(ctx, node1.Address))
+		assert.NoError(t, ns.DeleteNode(tx, node1.Address))
 
 		// Read node.
-		_, err = ns.ReadNode(ctx, node1.Address)
+		_, err = ns.ReadNode(tx, node1.Address)
 		if assert.Error(t, err) {
 			assert.True(t, errors.Is(err, dax.ErrNodeDoesNotExist))
 		}
+
+		assert.NoError(t, tx.Commit())
 	})
 }

dax/directive.go

@@ -1,7 +1,5 @@
 package dax
 
-import "context"
-
 // Directive contains the instructions, sent from MDS, which a compute node is
 // to follow. A Directive is typically JSON-encoded and POSTed to a compute
 // node's `/directive` endpoint.
@@ -22,7 +20,7 @@ type Directive struct {
 }
 
 type DirectiveVersion interface {
-	Increment(ctx context.Context, delta uint64) (uint64, error)
+	Increment(tx Transaction, delta uint64) (uint64, error)
 }
 
 // DirectiveMethod is used to tell the compute node how it should handle the
@@ -177,5 +175,5 @@ func (d *Directive) IsEmpty() bool {
 type Directives []*Directive
 
 func (d Directives) Len() int           { return len(d) }
-func (d Directives) Less(i, j int) bool { return d[i].Address.String() < d[j].Address.String() }
+func (d Directives) Less(i, j int) bool { return d[i].Version < d[j].Version }
 func (d Directives) Swap(i, j int)      { d[i], d[j] = d[j], d[i] }

dax/errors.go

@@ -7,6 +7,9 @@ import (
 )
 
 const (
+	ErrDatabaseIDExists       errors.Code = "DatabaseIDExists"
+	ErrDatabaseIDDoesNotExist errors.Code = "DatabaseIDDoesNotExist"
+
 	ErrTableIDExists         errors.Code = "TableIDExists"
 	ErrTableKeyExists        errors.Code = "TableKeyExists"
 	ErrTableNameExists       errors.Code = "TableNameExists"
@@ -17,12 +20,28 @@ const (
 	ErrFieldExists       errors.Code = "FieldExists"
 	ErrFieldDoesNotExist errors.Code = "FieldDoesNotExist"
 
+	ErrInvalidTransaction errors.Code = "InvalidTransaction"
+
 	ErrUnimplemented errors.Code = "Unimplemented"
 )
 
 // The following are helper functions for constructing coded errors containing
 // relevant information about the specific error.
 
+func NewErrDatabaseIDExists(qdbid QualifiedDatabaseID) error {
+	return errors.New(
+		ErrDatabaseIDExists,
+		fmt.Sprintf("database ID '%s' already exists", qdbid),
+	)
+}
+
+func NewErrDatabaseIDDoesNotExist(qdbid QualifiedDatabaseID) error {
+	return errors.New(
+		ErrDatabaseIDDoesNotExist,
+		fmt.Sprintf("database ID '%s' does not exist", qdbid),
+	)
+}
+
 func NewErrTableIDDoesNotExist(qtid QualifiedTableID) error {
 	return errors.New(
 		ErrTableIDDoesNotExist,
@@ -78,3 +97,10 @@ func NewErrFieldExists(fieldName FieldName) error {
 		fmt.Sprintf("field '%s' already exists", fieldName),
 	)
 }
+
+func NewErrInvalidTransaction() error {
+	return errors.New(
+		ErrInvalidTransaction,
+		"tx is not a *boltdb.Tx",
+	)
+}

dax/mds/client/client.go

@@ -54,8 +54,8 @@ func (c *Client) TableByID(ctx context.Context, qtid dax.QualifiedTableID) (*dax
 }
 
 // TODO(tlt): collapse TableID into this
-func (c *Client) TableByName(ctx context.Context, qual dax.TableQualifier, tname dax.TableName) (*dax.QualifiedTable, error) {
-	qtid, err := c.TableID(ctx, qual, tname)
+func (c *Client) TableByName(ctx context.Context, qdbid dax.QualifiedDatabaseID, tname dax.TableName) (*dax.QualifiedTable, error) {
+	qtid, err := c.TableID(ctx, qdbid, tname)
 	if err != nil {
 		return nil, errors.Wrap(err, "getting table id")
 	}
@@ -93,14 +93,14 @@ func (c *Client) Table(ctx context.Context, qtid dax.QualifiedTableID) (*dax.Qua
 	return qtable, nil
 }
 
-func (c *Client) TableID(ctx context.Context, qual dax.TableQualifier, name dax.TableName) (dax.QualifiedTableID, error) {
+func (c *Client) TableID(ctx context.Context, qdbid dax.QualifiedDatabaseID, name dax.TableName) (dax.QualifiedTableID, error) {
 	url := fmt.Sprintf("%s/table-id", c.address.WithScheme(defaultScheme))
 
 	dflt := dax.QualifiedTableID{}
 
 	req := dax.QualifiedTableID{
-		TableQualifier: qual,
-		Name:           name,
+		QualifiedDatabaseID: qdbid,
+		Name:                name,
 	}
 
 	// Encode the request.
@@ -130,12 +130,12 @@ func (c *Client) TableID(ctx context.Context, qual dax.TableQualifier, name dax.
 	return qtid, nil
 }
 
-func (c *Client) Tables(ctx context.Context, qual dax.TableQualifier, ids ...dax.TableID) ([]*dax.QualifiedTable, error) {
+func (c *Client) Tables(ctx context.Context, qdbid dax.QualifiedDatabaseID, ids ...dax.TableID) ([]*dax.QualifiedTable, error) {
 	url := fmt.Sprintf("%s/tables", c.address.WithScheme(defaultScheme))
 
 	req := mdshttp.TablesRequest{
-		OrganizationID: qual.OrganizationID,
-		DatabaseID:     qual.DatabaseID,
+		OrganizationID: qdbid.OrganizationID,
+		DatabaseID:     qdbid.DatabaseID,
 		TableIDs:       ids,
 	}
 
@@ -166,6 +166,62 @@ func (c *Client) Tables(ctx context.Context, qual dax.TableQualifier, ids ...dax
 	return qtables, nil
 }
 
+func (c *Client) CreateDatabase(ctx context.Context, qdb *dax.QualifiedDatabase) error {
+	url := fmt.Sprintf("%s/create-database", c.address.WithScheme(defaultScheme))
+
+	// Encode the request.
+	postBody, err := json.Marshal(qdb)
+	if err != nil {
+		return errors.Wrap(err, "marshalling post request")
+	}
+	responseBody := bytes.NewBuffer(postBody)
+
+	// Post the request.
+	resp, err := http.Post(url, "application/json", responseBody)
+	if err != nil {
+		return errors.Wrap(err, "posting create database request")
+	}
+	defer resp.Body.Close()
+
+	if resp.StatusCode != http.StatusOK {
+		b, _ := io.ReadAll(resp.Body)
+		return errors.Errorf("status code: %d: %s", resp.StatusCode, b)
+	}
+
+	return nil
+}
+
+func (c *Client) DatabaseByID(ctx context.Context, qdbid dax.QualifiedDatabaseID) (*dax.QualifiedDatabase, error) {
+	url := fmt.Sprintf("%s/database-by-id", c.address.WithScheme(defaultScheme))
+
+	// Encode the request.
+	postBody, err := json.Marshal(qdbid)
+	if err != nil {
+		return nil, errors.Wrap(err, "marshalling post request")
+	}
+	responseBody := bytes.NewBuffer(postBody)
+
+	// Post the request.
+	c.logger.Debugf("POST database request: url: %s", url)
+	resp, err := http.Post(url, "application/json", responseBody)
+	if err != nil {
+		return nil, errors.Wrap(err, "posting table request")
+	}
+	defer resp.Body.Close()
+
+	if resp.StatusCode != http.StatusOK {
+		b, _ := io.ReadAll(resp.Body)
+		return nil, errors.Errorf("status code: %d: %s", resp.StatusCode, b)
+	}
+
+	var qdb *dax.QualifiedDatabase
+	if err := json.NewDecoder(resp.Body).Decode(&qdb); err != nil {
+		return nil, errors.Wrap(err, "reading response body")
+	}
+
+	return qdb, nil
+}
+
 func (c *Client) CreateTable(ctx context.Context, qtbl *dax.QualifiedTable) error {
 	url := fmt.Sprintf("%s/create-table", c.address.WithScheme(defaultScheme))
 

dax/mds/controller/balancer.go

@@ -8,31 +8,38 @@ import (
 )
 
 type Balancer interface {
-	AddWorker(ctx context.Context, worker fmt.Stringer) ([]dax.WorkerDiff, error)
-	RemoveWorker(ctx context.Context, worker fmt.Stringer) ([]dax.WorkerDiff, error)
-	AddJobs(ctx context.Context, job ...fmt.Stringer) ([]dax.WorkerDiff, error)
-	RemoveJob(ctx context.Context, job fmt.Stringer) ([]dax.WorkerDiff, error)
-	Balance(ctx context.Context) ([]dax.WorkerDiff, error)
-	CurrentState(ctx context.Context) ([]dax.WorkerInfo, error)
-	WorkerState(ctx context.Context, worker dax.Worker) (dax.WorkerInfo, error)
-	WorkersForJobs(ctx context.Context, jobs []dax.Job) ([]dax.WorkerInfo, error)
+	// AddWorker adds a worker to the global pool of available workers.
+	AddWorker(tx dax.Transaction, node *dax.Node) ([]dax.WorkerDiff, error)
 
-	// WorkersForJobPrefix returns all workers and their job
-	// assignments which start with `prefix` for all jobs that start
-	// with `prefix`. If there are free jobs that start with `prefix`
-	// an error is returned.
-	//
-	// The motivating use case is getting all workers for a particular
-	// table so we can execute a query that will hit every shard in a
-	// table. If there are jobs representing shards in that table
-	// which are not assigned to any worker, that means the query
-	// would return incomplete data, so we want to error.
-	WorkersForJobPrefix(ctx context.Context, prefix string) ([]dax.WorkerInfo, error)
+	// RemoveWorker removes a worker from the system. If the worker is currently
+	// assigned to a database and has jobs, it will be removed and its jobs will
+	// be either transferred to other workers or placed on the free job list.
+	RemoveWorker(tx dax.Transaction, addr dax.Address) ([]dax.WorkerDiff, error)
 
-	// RemoveJobs is for e.g. when dropping a table remove all jobs
-	// associated with that table without needing to look up in
-	// advance which shards or partitions are actually present.
-	RemoveJobs(ctx context.Context, prefix string) ([]dax.WorkerDiff, error)
+	// AddJobs adds new jobs for the given database.
+	AddJobs(tx dax.Transaction, roleType dax.RoleType, qtid dax.QualifiedTableID, jobs ...dax.Job) ([]dax.WorkerDiff, error)
+
+	// RemoveJobs removes jobs for the given database.
+	RemoveJobs(tx dax.Transaction, roleType dax.RoleType, qtid dax.QualifiedTableID, jobs ...dax.Job) ([]dax.WorkerDiff, error)
+
+	// BalanceDatabase forces a database balance. TODO(tlt): currently this is
+	// only used in tests, so perhaps we can get rid of it.
+	BalanceDatabase(tx dax.Transaction, qdbid dax.QualifiedDatabaseID) ([]dax.WorkerDiff, error)
+
+	// CurrentState returns the workers and jobs currently active for the given
+	// database.
+	CurrentState(tx dax.Transaction, roleType dax.RoleType, qdbid dax.QualifiedDatabaseID) ([]dax.WorkerInfo, error)
+
+	// WorkerState returns the jobs currently active for the given worker.
+	WorkerState(tx dax.Transaction, roleType dax.RoleType, addr dax.Address) (dax.WorkerInfo, error)
+
+	// WorkersForJobs returns the workers and jobs currently responsible for the
+	// given jobs.
+	WorkersForJobs(tx dax.Transaction, roleType dax.RoleType, qdbid dax.QualifiedDatabaseID, jobs ...dax.Job) ([]dax.WorkerInfo, error)
+
+	// WorkersForTable returns the workers responsible for any job related to
+	// the given table.
+	WorkersForTable(tx dax.Transaction, roleType dax.RoleType, qtid dax.QualifiedTableID) ([]dax.WorkerInfo, error)
 }
 
 // Ensure type implements interface.
@@ -45,34 +52,30 @@ func NewNopBalancer() *NopBalancer {
 	return &NopBalancer{}
 }
 
-func (b *NopBalancer) AddWorker(ctx context.Context, worker fmt.Stringer) ([]dax.WorkerDiff, error) {
+func (b *NopBalancer) AddWorker(tx dax.Transaction, node *dax.Node) ([]dax.WorkerDiff, error) {
 	return []dax.WorkerDiff{}, nil
 }
-func (b *NopBalancer) RemoveWorker(ctx context.Context, worker fmt.Stringer) ([]dax.WorkerDiff, error) {
+func (b *NopBalancer) RemoveWorker(tx dax.Transaction, addr dax.Address) ([]dax.WorkerDiff, error) {
 	return []dax.WorkerDiff{}, nil
 }
-func (b *NopBalancer) AddJobs(ctx context.Context, job ...fmt.Stringer) ([]dax.WorkerDiff, error) {
+func (b *NopBalancer) AddJobs(tx dax.Transaction, roleType dax.RoleType, qtid dax.QualifiedTableID, jobs ...dax.Job) ([]dax.WorkerDiff, error) {
 	return []dax.WorkerDiff{}, nil
 }
-func (b *NopBalancer) RemoveJob(ctx context.Context, job fmt.Stringer) ([]dax.WorkerDiff, error) {
+func (b *NopBalancer) RemoveJobs(tx dax.Transaction, roleType dax.RoleType, qtid dax.QualifiedTableID, jobs ...dax.Job) ([]dax.WorkerDiff, error) {
 	return []dax.WorkerDiff{}, nil
 }
-func (b *NopBalancer) Balance(ctx context.Context) ([]dax.WorkerDiff, error) {
+func (b *NopBalancer) BalanceDatabase(tx dax.Transaction, qdbid dax.QualifiedDatabaseID) ([]dax.WorkerDiff, error) {
 	return []dax.WorkerDiff{}, nil
 }
-func (b *NopBalancer) CurrentState(ctx context.Context) ([]dax.WorkerInfo, error) {
+func (b *NopBalancer) CurrentState(tx dax.Transaction, roleType dax.RoleType, qdbid dax.QualifiedDatabaseID) ([]dax.WorkerInfo, error) {
 	return []dax.WorkerInfo{}, nil
 }
-func (b *NopBalancer) WorkerState(ctx context.Context, worker dax.Worker) (dax.WorkerInfo, error) {
+func (b *NopBalancer) WorkerState(tx dax.Transaction, roleType dax.RoleType, addr dax.Address) (dax.WorkerInfo, error) {
 	return dax.WorkerInfo{}, nil
 }
-func (b *NopBalancer) WorkersForJobs(ctx context.Context, jobs []dax.Job) ([]dax.WorkerInfo, error) {
+func (b *NopBalancer) WorkersForJobs(tx dax.Transaction, roleType dax.RoleType, qdbid dax.QualifiedDatabaseID, jobs ...dax.Job) ([]dax.WorkerInfo, error) {
 	return []dax.WorkerInfo{}, nil
 }
-func (b *NopBalancer) WorkersForJobPrefix(ctx context.Context, prefix string) ([]dax.WorkerInfo, error) {
+func (b *NopBalancer) WorkersForTable(tx dax.Transaction, roleType dax.RoleType, qtid dax.QualifiedTableID) ([]dax.WorkerInfo, error) {
 	return []dax.WorkerInfo{}, nil
 }
-
-func (b *NopBalancer) RemoveJobs(ctx context.Context, prefix string) ([]dax.WorkerDiff, error) {
-	return nil, nil
-}

dax/mds/controller/balancer/balancer.go

@@ -0,0 +1,821 @@
+// Package balancer is an implementation of the controller's Balancer interface.
+package balancer
+
+import (
+	"log"
+	"math"
+	"sort"
+	"strings"
+	"time"
+
+	"github.com/featurebasedb/featurebase/v3/dax"
+	"github.com/featurebasedb/featurebase/v3/dax/mds/controller"
+	"github.com/featurebasedb/featurebase/v3/dax/mds/schemar"
+	"github.com/featurebasedb/featurebase/v3/errors"
+	"github.com/featurebasedb/featurebase/v3/logger"
+)
+
+// Ensure type implements interface.
+var _ controller.Balancer = (*Balancer)(nil)
+
+// Balancer is an implementation of the controller.Balancer interface which
+// isolates workers and jobs by database. It helps manage the relationships
+// between workers and jobs. The logic it uses to balance jobs across workers is
+// very simple; it bases everything off the number of workers and number of
+// jobs. It does not take anything else (such as job size, worker capabilities,
+// etc) into consideration.
+type Balancer struct {
+	// current represents the current state of worker/job assigments.
+	current WorkerJobService
+
+	// freeJobs is the set of jobs which have yet to be assigned to a worker.
+	// This could be because there are no available workers, or because a worker
+	// has been removed and the jobs for which it was responsible have yet to be
+	// reassigned.
+	freeJobs FreeJobService
+
+	freeWorkers FreeWorkerService
+
+	schemar schemar.Schemar
+
+	logger logger.Logger
+}
+
+// New returns a new instance of Balancer.
+func New(fjs FreeJobService, wjs WorkerJobService, fws FreeWorkerService, schemar schemar.Schemar, logger logger.Logger) *Balancer {
+	return &Balancer{
+		current:     wjs,
+		freeJobs:    fjs,
+		freeWorkers: fws,
+		schemar:     schemar,
+		logger:      logger,
+	}
+}
+
+// AddWorker adds the given Node to the Balancer's available worker pool.
+// TODO(tlt): this method takes a Node (as opposed to a Worker) because in the
+// future we may want to maintain separate worker pools based on RoleType
+// (compute, translate, etc.).
+func (b *Balancer) AddWorker(tx dax.Transaction, node *dax.Node) ([]dax.WorkerDiff, error) {
+	addr := node.Address
+	b.logger.Debugf("AddWorker(%s)", addr)
+
+	diffs := NewInternalDiffs()
+
+	// This logic means that a node is used for ALL of the role types specified.
+	// In other words, specifying roleTypes = {compute, translate}, does not
+	// mean that the node can be used as either a compute worker or a translate
+	// worker. It means that it will be used as both.
+	for _, rt := range node.RoleTypes {
+		if err := b.addWorker(tx, rt, addr); err != nil {
+			return nil, errors.Wrapf(err, "adding worker: (%s) %s", rt, addr)
+		}
+	}
+
+	// Process the freeWorkers.
+	// TODO(tlt): this is a little heavy-handed. I'm sure we'll need to be more
+	// intentional about knowing which databases needs workers, as opposed to
+	// this brute force loop over all databases every time.
+	if diff, err := b.balance(tx); err != nil {
+		return nil, errors.Wrapf(err, "balancing new worker: %s", addr)
+	} else {
+		diffs.Merge(diff)
+	}
+
+	return diffs.Output(), nil
+}
+
+// addWorker adds a worker to the free worker list. From there, it can be used
+// by any database which needs a worker.
+func (b *Balancer) addWorker(tx dax.Transaction, roleType dax.RoleType, addr dax.Address) error {
+	// If this worker already exists, don't do anything.
+	if dbkey := b.current.DatabaseForWorker(tx, addr); dbkey != "" {
+		return nil
+	}
+
+	if err := b.freeWorkers.AddWorkers(tx, roleType, addr); err != nil {
+		return errors.Wrap(err, "adding free worker")
+	}
+
+	return nil
+}
+
+func (b *Balancer) assignMinWorkers(tx dax.Transaction, roleType dax.RoleType) (InternalDiffs, error) {
+	// Find out how many free workers we have.
+	freeWorkers, err := b.freeWorkers.ListWorkers(tx, roleType)
+	if err != nil {
+		return nil, errors.Wrap(err, "getting free worker list")
+	}
+	freeWorkerCount := len(freeWorkers)
+
+	// If there are no free workers, return early.
+	if freeWorkerCount == 0 {
+		return InternalDiffs{}, nil
+	}
+
+	// Get all database and their minWorkerCount (Database.Options.WorkersMin).
+	qdbs, err := b.schemar.Databases(tx, "")
+	if err != nil {
+		return nil, errors.Wrap(err, "getting all database")
+	}
+
+	// Create a map[database]int where int is the number of workers required to
+	// reach that database's minWorkerCount. This map will only contain database
+	// which need more workers in order to reach their minimum.
+	m := make(map[dax.QualifiedDatabaseID]int)
+
+	for _, qdb := range qdbs {
+		qdbid := qdb.QualifiedID()
+
+		minWorkers := qdb.Options.WorkersMin
+		if minWorkers == 0 {
+			continue
+		}
+
+		// If the database doesn't have any jobs, there's no need to go any
+		// further. In other words, we don't want to assign a worker to a
+		// database until it has at least one job.
+		if hasJobs, err := b.databaseHasJobs(tx, roleType, qdbid); err != nil {
+			return nil, errors.Wrapf(err, "checking has jobs: (%s) %s", roleType, qdbid)
+		} else if !hasJobs {
+			continue
+		}
+
+		// Get the number of workers assigned to this database.
+		workerCount, err := b.current.WorkerCount(tx, roleType, qdbid)
+		if err != nil {
+			return nil, errors.Wrapf(err, "getting worker count: (%s) %s", roleType, qdbid)
+		}
+
+		diff := minWorkers - workerCount
+
+		// If we have more workers than the min required, or if we have the
+		// exact number of workers , don't do anything for that database.
+		if diff <= 0 {
+			continue
+		}
+
+		m[qdbid] = diff
+	}
+
+	diffs := NewInternalDiffs()
+
+	// Create an ordered slice of map keys so that tests are predicatable.
+	qdbids := make([]dax.QualifiedDatabaseID, 0, len(m))
+	for qdbid := range m {
+		qdbids = append(qdbids, qdbid)
+	}
+	sort.Sort(dax.QualifiedDatabaseIDs(qdbids))
+
+	// For each database, if there are enough free workers to
+	// satisfy its min, then pop that number of workers from the free list. If
+	// not, contine to the next database until either reaching the end of the
+	// database list or until there are no more free workers in the list,
+	// whichever comes first.
+	for _, qdbid := range qdbids {
+		need := m[qdbid]
+
+		if freeWorkerCount == 0 {
+			break
+		}
+
+		if freeWorkerCount >= need {
+			addrs, err := b.freeWorkers.PopWorkers(tx, roleType, need)
+			if err != nil {
+				return nil, errors.Wrapf(err, "popping free worker: (%s)", roleType)
+			}
+
+			if diff, err := b.addDatabaseWorkers(tx, roleType, qdbid, addrs...); err != nil {
+				return nil, errors.Wrapf(err, "adding database workers: (%s) %s, %v", roleType, qdbid, addrs)
+			} else {
+				diffs.Merge(diff)
+			}
+		}
+	}
+
+	return diffs, nil
+}
+
+// addDatabaseWorkers adds workers from the free worker list to the pool of
+// workers for a specific database.
+func (b *Balancer) addDatabaseWorkers(tx dax.Transaction, roleType dax.RoleType, qdbid dax.QualifiedDatabaseID, addrs ...dax.Address) (InternalDiffs, error) {
+	for _, addr := range addrs {
+		if err := b.current.CreateWorker(tx, roleType, qdbid, addr); err != nil {
+			return nil, errors.Wrap(err, "creating worker")
+		}
+	}
+
+	// Process the freeJobs.
+	return b.processFreeJobs(tx, roleType, qdbid)
+}
+
+// databaseHasJobs returns true if the database has at least one job.
+func (b *Balancer) databaseHasJobs(tx dax.Transaction, roleType dax.RoleType, qdbid dax.QualifiedDatabaseID) (bool, error) {
+	// Free jobs.
+	if freeJobs, err := b.freeJobs.ListJobs(tx, roleType, qdbid); err != nil {
+		return false, errors.Wrapf(err, "getting free jobs: (%s) %s", roleType, qdbid)
+	} else if len(freeJobs) > 0 {
+		return true, nil
+	}
+
+	// Assigned jobs.
+	if wis, err := b.current.WorkersJobs(tx, roleType, qdbid); err != nil {
+		return false, errors.Wrapf(err, "getting free jobs: (%s) %s", roleType, qdbid)
+	} else {
+		for _, wi := range wis {
+			if len(wi.Jobs) > 0 {
+				return true, nil
+			}
+		}
+	}
+
+	return false, nil
+}
+
+func (b *Balancer) RemoveWorker(tx dax.Transaction, addr dax.Address) ([]dax.WorkerDiff, error) {
+	diffs := NewInternalDiffs()
+
+	// See if the worker is assigned to a database.
+	dbkey := b.current.DatabaseForWorker(tx, addr)
+	if dbkey == "" {
+		return diffs.Output(), nil
+	}
+	qdbid := dbkey.QualifiedDatabaseID()
+
+	for _, rt := range []dax.RoleType{dax.RoleTypeCompute, dax.RoleTypeTranslate} {
+		// Remove the worker form the free worker list (if it's there).
+		if err := b.freeWorkers.RemoveWorker(tx, rt, addr); err != nil {
+			return nil, errors.Wrapf(err, "removing worker from free list: (%s) %s", rt, addr)
+		}
+
+		if diff, err := b.removeDatabaseWorker(tx, rt, qdbid, addr); err != nil {
+			return nil, errors.Wrapf(err, "removing worker: (%s) %s", rt, addr)
+		} else {
+			diffs.Merge(diff)
+		}
+	}
+
+	// Balance the affected database.
+	if diff, err := b.balanceDatabase(tx, qdbid); err != nil {
+		return nil, errors.Wrapf(err, "balancing database: %s", qdbid)
+	} else {
+		diffs.Merge(diff)
+	}
+
+	return diffs.Output(), nil
+}
+
+func (b *Balancer) removeDatabaseWorker(tx dax.Transaction, roleType dax.RoleType, qdbid dax.QualifiedDatabaseID, addr dax.Address) (InternalDiffs, error) {
+	jobs, err := b.current.ListJobs(tx, roleType, qdbid, addr)
+	if err != nil {
+		return nil, errors.Wrap(err, "listing jobs")
+	}
+
+	// Before removing the worker, mark its jobs as free.
+	if err := b.freeJobs.MergeJobs(tx, roleType, qdbid, jobs); err != nil {
+		return nil, errors.Wrap(err, "merging free jobs")
+	}
+
+	// Remove the worker.
+	if err := b.current.DeleteWorker(tx, roleType, qdbid, addr); err != nil {
+		return nil, errors.Wrap(err, "deleting worker")
+	}
+
+	// Even though this may not be useful to the caller (for example, in the
+	// case where the worker has died and no longer exists), return the diffs
+	// which represent the removal of jobs from the worker.
+	diff := NewInternalDiffs()
+	for _, job := range jobs {
+		diff.Removed(addr, job)
+	}
+
+	return diff, nil
+}
+
+func (b *Balancer) AddJobs(tx dax.Transaction, roleType dax.RoleType, qtid dax.QualifiedTableID, jobs ...dax.Job) ([]dax.WorkerDiff, error) {
+	start := time.Now()
+	defer func() {
+		log.Printf("ELAPSED: Balancer.AddJob: %v", time.Since(start))
+	}()
+
+	switch len(jobs) {
+	case 0:
+		// No jobs so return early.
+		b.logger.Debugf("%s: AddJobs (no jobs provided)", roleType)
+		return []dax.WorkerDiff{}, nil
+	case 1:
+		b.logger.Debugf("%s: AddJobs (%s)", roleType, jobs[0])
+	default:
+		b.logger.Debugf("%s: AddJobs (%d)", roleType, len(jobs))
+	}
+
+	// TODO(tlt): we don't currently use "table" in this method; even though we
+	// pass a table, we're still encoding the tableKey in the job. In theory, we
+	// could exclude tableKey from the job coming into this method, and add it
+	// here.
+	qdbid := qtid.QualifiedDatabaseID
+
+	diff, err := b.addJobs(tx, roleType, qdbid, jobs...)
+	if err != nil {
+		return nil, errors.Wrap(err, "adding job")
+	}
+
+	return diff.Output(), nil
+}
+
+func (b *Balancer) addJobs(tx dax.Transaction, roleType dax.RoleType, qdbid dax.QualifiedDatabaseID, jobs ...dax.Job) (InternalDiffs, error) {
+	diffs := NewInternalDiffs()
+
+	if len(jobs) == 0 {
+		return diffs, nil
+	}
+
+	if cnt, err := b.current.WorkerCount(tx, roleType, qdbid); err != nil {
+		return nil, errors.Wrap(err, "getting worker count")
+	} else if cnt == 0 {
+		if err := b.freeJobs.CreateJobs(tx, roleType, qdbid, jobs...); err != nil {
+			return nil, errors.Wrap(err, "creating free job")
+		}
+
+		// Since we've just added free jobs to the database, try to add a worker
+		// for the database. This case would happen because when a database is
+		// first created, it is not assigned any workers. A database is not
+		// assigned workers until it has at least one job (which this database
+		// now has).
+		if diff, err := b.balanceDatabaseForRole(tx, roleType, qdbid); err != nil {
+			return nil, errors.Wrapf(err, "assigning min workers: (%s)", roleType)
+		} else {
+			diffs.Merge(diff)
+		}
+
+		// Now check, again, to see if the database has a worker.
+		if cnt2, err := b.current.WorkerCount(tx, roleType, qdbid); err != nil {
+			return nil, errors.Wrap(err, "getting worker count, again")
+		} else if cnt2 == 0 {
+			// TODO: we might want to inform the user that a job is in the free list
+			// because there are no workers.
+			return InternalDiffs{}, nil
+		}
+	}
+
+	diff, err := b.addDatabaseJobs(tx, roleType, qdbid, jobs...)
+	if err != nil {
+		return nil, errors.Wrapf(err, "adding database jobs: (%s) %s", roleType, qdbid)
+	}
+	diffs.Merge(diff)
+
+	return diffs, nil
+}
+
+// addDatabaseJobs adds the job for the provided database.
+func (b *Balancer) addDatabaseJobs(tx dax.Transaction, roleType dax.RoleType, qdbid dax.QualifiedDatabaseID, jobs ...dax.Job) (InternalDiffs, error) {
+	workerJobs, err := b.current.WorkersJobs(tx, roleType, qdbid)
+	if err != nil {
+		return nil, errors.Wrapf(err, "getting workers jobs: %s", roleType)
+	}
+	jset := dax.NewSet[dax.Job]()
+	for _, workerInfo := range workerJobs {
+		jset.Merge(dax.NewSet(workerInfo.Jobs...))
+	}
+
+	addrs := make(dax.Addresses, 0, len(workerJobs))
+	jobCounts := make(map[dax.Address]int, 0)
+	for _, v := range workerJobs {
+		addrs = append(addrs, v.Address)
+		jobCounts[v.Address] = len(v.Jobs)
+	}
+
+	diffs := NewInternalDiffs()
+
+	jobsToCreate := make(map[dax.Address][]dax.Job)
+
+	for _, job := range jobs {
+		// Skip any job that already exists.
+		if jset.Contains(job) {
+			continue
+		}
+
+		// Find the worker with the fewest number of jobs and assign it this job.
+		var lowCount int = math.MaxInt
+		var lowWorker dax.Address
+
+		// We loop over addrs here instead of jobCounts because jobCounts is a
+		// map and it can return results in an unexpected order, which is a
+		// problem for testing.
+		for _, addr := range addrs {
+			jobCount := jobCounts[addr]
+			if jobCount < lowCount {
+				lowCount = jobCount
+				lowWorker = addr
+			}
+		}
+
+		jobsToCreate[lowWorker] = append(jobsToCreate[lowWorker], job)
+		jobCounts[lowWorker]++
+	}
+
+	for worker, jobs := range jobsToCreate {
+		if err := b.current.CreateJobs(tx, roleType, qdbid, worker, jobs...); err != nil {
+			return nil, errors.Wrap(err, "creating job")
+		}
+		for _, job := range jobs {
+			diffs.Added(worker, job)
+		}
+	}
+
+	return diffs, nil
+}
+
+func (b *Balancer) RemoveJobs(tx dax.Transaction, roleType dax.RoleType, qtid dax.QualifiedTableID, jobs ...dax.Job) ([]dax.WorkerDiff, error) {
+	qdbid := qtid.QualifiedDatabaseID
+
+	// If no jobs are provided, remove all jobs for table.
+	if len(jobs) == 0 {
+		diffs, err := b.removeJobsForTable(tx, roleType, qtid)
+		if err != nil {
+			return nil, errors.Wrapf(err, "removing jobs for table: (%s) %s", roleType, qtid)
+		}
+		return diffs.Output(), nil
+	}
+
+	diffs := NewInternalDiffs()
+
+	for _, job := range jobs {
+		if diff, err := b.removeJob(tx, roleType, qdbid, job); err != nil {
+			return nil, errors.Wrapf(err, "removing job: (%s) %s, %s", roleType, qdbid, job)
+		} else {
+			diffs.Merge(diff)
+		}
+	}
+
+	return diffs.Output(), nil
+}
+
+func (b *Balancer) removeJobsForTable(tx dax.Transaction, roleType dax.RoleType, qtid dax.QualifiedTableID) (InternalDiffs, error) {
+	idiffs, err := b.current.DeleteJobsForTable(tx, roleType, qtid)
+	if err != nil {
+		return nil, errors.Wrapf(err, "deleting jobs for table: (%s) %s", roleType, qtid)
+	}
+	if err := b.freeJobs.DeleteJobsForTable(tx, roleType, qtid); err != nil {
+		return nil, errors.Wrapf(err, "deleting free jobs for table: (%s) %s", roleType, qtid)
+	}
+	return idiffs, nil
+}
+
+// Balance calls balanceDatabase on every database in the schemar.
+func (b *Balancer) Balance(tx dax.Transaction) ([]dax.WorkerDiff, error) {
+	diffs, err := b.balance(tx)
+	if err != nil {
+		return nil, errors.Wrapf(err, "balancing all")
+	}
+
+	return diffs.Output(), nil
+}
+
+func (b *Balancer) balance(tx dax.Transaction) (InternalDiffs, error) {
+	qdbs, err := b.schemar.Databases(tx, "")
+	if err != nil {
+		return nil, errors.Wrapf(err, "getting all databases")
+	}
+
+	diffs := NewInternalDiffs()
+
+	for _, qdb := range qdbs {
+		qdbid := qdb.QualifiedID()
+		if diff, err := b.balanceDatabase(tx, qdbid); err != nil {
+			return nil, errors.Wrapf(err, "balancing database: %s", qdbid)
+		} else {
+			diffs.Merge(diff)
+		}
+	}
+
+	return diffs, nil
+}
+
+func (b *Balancer) BalanceDatabase(tx dax.Transaction, qdbid dax.QualifiedDatabaseID) ([]dax.WorkerDiff, error) {
+	diffs, err := b.balanceDatabase(tx, qdbid)
+	if err != nil {
+		return nil, errors.Wrapf(err, "balancing database: %s", qdbid)
+	}
+
+	return diffs.Output(), nil
+}
+
+func (b *Balancer) balanceDatabase(tx dax.Transaction, qdbid dax.QualifiedDatabaseID) (InternalDiffs, error) {
+	diffs := NewInternalDiffs()
+
+	for _, role := range []dax.RoleType{dax.RoleTypeCompute, dax.RoleTypeTranslate} {
+		diff, err := b.balanceDatabaseForRole(tx, role, qdbid)
+		if err != nil {
+			return nil, errors.Wrapf(err, "getting worker count: (%s) %s", role, qdbid)
+		}
+		diffs.Merge(diff)
+	}
+
+	return diffs, nil
+}
+
+func (b *Balancer) balanceDatabaseForRole(tx dax.Transaction, roleType dax.RoleType, qdbid dax.QualifiedDatabaseID) (InternalDiffs, error) {
+	diffs := NewInternalDiffs()
+
+	// Before balancing, make sure the database has its minimum number of
+	// workers satisfied.
+	// TODO(tlt): make assignMinWorkers database specific.
+	if diff, err := b.assignMinWorkers(tx, roleType); err != nil {
+		return nil, errors.Wrapf(err, "assigning min workers: (%s) %s", roleType, qdbid)
+	} else {
+		diffs.Merge(diff)
+	}
+
+	// If there are no workers, we can't properly balance.
+	if cnt, err := b.current.WorkerCount(tx, roleType, qdbid); err != nil {
+		return nil, errors.Wrapf(err, "getting worker count: (%s) %s", roleType, qdbid)
+	} else if cnt == 0 {
+		return InternalDiffs{}, nil
+	}
+
+	// Process the freeJobs.
+	if diff, err := b.processFreeJobs(tx, roleType, qdbid); err != nil {
+		return nil, errors.Wrapf(err, "processing free jobs: (%s) %s", roleType, qdbid)
+	} else {
+		diffs.Merge(diff)
+	}
+
+	// Balance the jobs among workers.
+	diff, err := b.balanceDatabaseJobs(tx, roleType, qdbid, diffs)
+	if err != nil {
+		return nil, errors.Wrap(err, "balancing jobs")
+	}
+
+	return diff, nil
+}
+
+func (b *Balancer) CurrentState(tx dax.Transaction, roleType dax.RoleType, qdbid dax.QualifiedDatabaseID) ([]dax.WorkerInfo, error) {
+	return b.current.WorkersJobs(tx, roleType, qdbid)
+}
+
+func (b *Balancer) WorkerState(tx dax.Transaction, roleType dax.RoleType, addr dax.Address) (dax.WorkerInfo, error) {
+	info := dax.WorkerInfo{
+		Address: addr,
+	}
+
+	dbkey := b.current.DatabaseForWorker(tx, addr)
+	if dbkey == "" {
+		return info, nil
+	}
+	qdbid := dbkey.QualifiedDatabaseID()
+
+	jobs, err := b.current.ListJobs(tx, roleType, qdbid, addr)
+	if err != nil {
+		return dax.WorkerInfo{}, errors.Wrapf(err, "listing jobs: (%s) %s, %s", roleType, qdbid, addr)
+	}
+	info.Jobs = jobs
+
+	return info, nil
+}
+
+func (b *Balancer) WorkersForJobs(tx dax.Transaction, roleType dax.RoleType, qdbid dax.QualifiedDatabaseID, jobs ...dax.Job) ([]dax.WorkerInfo, error) {
+	out := make(map[dax.Address]dax.Set[dax.Job])
+
+	workerJobs, err := b.current.WorkersJobs(tx, roleType, qdbid)
+	if err != nil {
+		return nil, errors.Wrapf(err, "getting worker jobs: (%s) %s", roleType, qdbid)
+	}
+	for _, workerInfo := range workerJobs {
+		jset := dax.NewSet(workerInfo.Jobs...)
+
+		matches := dax.NewSet[dax.Job]()
+		for _, job := range jobs {
+			if jset.Contains(job) {
+				matches.Add(job)
+			}
+		}
+
+		if len(matches) > 0 {
+			out[workerInfo.Address] = matches
+		}
+	}
+
+	workers := make([]dax.WorkerInfo, 0, len(out))
+
+	for addr, jset := range out {
+		workers = append(workers, dax.WorkerInfo{
+			Address: addr,
+			Jobs:    jset.Sorted(),
+		})
+	}
+
+	sort.Sort(dax.WorkerInfos(workers))
+
+	return workers, nil
+}
+
+func (b *Balancer) WorkersForTable(tx dax.Transaction, roleType dax.RoleType, qtid dax.QualifiedTableID) ([]dax.WorkerInfo, error) {
+	out := make(map[dax.Address]dax.Set[dax.Job])
+
+	qdbid := qtid.QualifiedDatabaseID
+
+	prefix := string(qtid.Key())
+
+	workerJobs, err := b.current.WorkersJobs(tx, roleType, qdbid)
+	if err != nil {
+		return nil, errors.Wrapf(err, "getting worker jobs: (%s) %s", roleType, qdbid)
+	}
+	for _, workerInfo := range workerJobs {
+
+		matches := dax.NewSet[dax.Job]()
+		for _, job := range workerInfo.Jobs {
+			if strings.HasPrefix(string(job), prefix) {
+				matches.Add(job)
+			}
+		}
+		if len(matches) > 0 {
+			out[workerInfo.Address] = matches
+		}
+	}
+
+	workers := make([]dax.WorkerInfo, 0, len(out))
+
+	for addr, jset := range out {
+		workers = append(workers, dax.WorkerInfo{
+			Address: addr,
+			Jobs:    jset.Sorted(),
+		})
+	}
+
+	sort.Sort(dax.WorkerInfos(workers))
+
+	return workers, nil
+}
+
+func (b *Balancer) removeJob(tx dax.Transaction, roleType dax.RoleType, qdbid dax.QualifiedDatabaseID, job dax.Job) (InternalDiffs, error) {
+	if addr, ok, err := b.workerForJob(tx, roleType, qdbid, job); err != nil {
+		return nil, errors.Wrapf(err, "getting worker for job: %s", job)
+	} else if ok {
+		if err := b.current.DeleteJob(tx, roleType, qdbid, addr, job); err != nil {
+			return nil, errors.Wrapf(err, "deleting job: (%s) %s, %s, %s", roleType, qdbid, addr, job)
+		}
+
+		diffs := NewInternalDiffs()
+		diffs.Removed(addr, job)
+
+		return diffs, nil
+	}
+
+	// Just in case the job is in the free list (and wasn't assigned to a
+	// worker), remove it; there's no need to provide a diff. There should never
+	// be a case where the same job is both in the free list and assigned to a
+	// worker.
+	if err := b.freeJobs.DeleteJob(tx, roleType, qdbid, job); err != nil {
+		return nil, errors.Wrapf(err, "deleting free job: (%s) %s, %s", roleType, qdbid, job)
+	}
+
+	return InternalDiffs{}, nil
+}
+
+// balanceDatabaseJobs moves jobs among workers with the goal of having an equal
+// number of jobs per worker. This method takes an `internalDiffs` as input for
+// cases where some action has preceeded this call which also resulted in
+// `internalDiffs`. Instead of having this method take a value, we could rely on
+// the internalDiffs.merge() method, but we would need to modify that method to
+// be smarter about the order in which it applies the add/remove operations.
+// Until that's in place, we'll pass in a value here.
+func (b *Balancer) balanceDatabaseJobs(tx dax.Transaction, roleType dax.RoleType, qdbid dax.QualifiedDatabaseID, diffs InternalDiffs) (InternalDiffs, error) {
+	numWorkers, err := b.current.WorkerCount(tx, roleType, qdbid)
+	if err != nil {
+		return nil, errors.Wrapf(err, "getting worker count: (%s) %s", roleType, qdbid)
+	}
+	numJobs := 0
+	if addrs, err := b.current.ListWorkers(tx, roleType, qdbid); err != nil {
+		return nil, errors.Wrapf(err, "listing workers: (%s) %s", roleType, qdbid)
+	} else {
+		for _, addr := range addrs {
+			jobCounts, err := b.current.JobCounts(tx, roleType, qdbid, addr)
+			if err != nil {
+				return nil, errors.Wrapf(err, "getting job count: (%s) %s, %s", roleType, qdbid, addr)
+			}
+			numJobs += jobCounts[addr]
+		}
+	}
+
+	minJobsPerWorker := numJobs / numWorkers
+	numWorkersAboveMin := numJobs % numWorkers
+
+	// workerInfos is used now in order to guarantee a sort order.
+	workerInfos, err := b.CurrentState(tx, roleType, qdbid)
+	if err != nil {
+		return nil, errors.Wrapf(err, "getting current state: (%s) %s", roleType, qdbid)
+	}
+
+	// Loop through each worker, and if the number of jobs for the worker
+	// exceeds the target, then remove the job and add it back (which is
+	// effectively how we rebalance a job).
+	for i, workerInfo := range workerInfos {
+		numTargetJobs := minJobsPerWorker
+		if i < numWorkersAboveMin {
+			numTargetJobs += 1
+		}
+
+		jobCounts, err := b.current.JobCounts(tx, roleType, qdbid, workerInfo.Address)
+		if err != nil {
+			return nil, errors.Wrapf(err, "getting job count: (%s) %s, %s", roleType, qdbid, workerInfo.Address)
+		}
+		numCurrentJobs := jobCounts[workerInfo.Address]
+
+		// If we don't need to remove jobs from this worker, then just continue
+		// on to the next worker.
+		if numCurrentJobs <= numTargetJobs {
+			continue
+		}
+
+		sortedJobs, err := b.current.ListJobs(tx, roleType, qdbid, workerInfo.Address)
+		if err != nil {
+			return nil, errors.Wrapf(err, "listing jobs: (%s) %s, %s", roleType, qdbid, workerInfo.Address)
+		}
+
+		// Remove the extra jobs from the end of the list, and add them back
+		// again (which should place them on a worker with fewer jobs).
+		for i := numCurrentJobs - 1; i >= numTargetJobs; i-- {
+			if rj, err := b.removeJob(tx, roleType, qdbid, sortedJobs[i]); err != nil {
+				return nil, errors.Wrapf(err, "removing job: %s", sortedJobs[i])
+			} else {
+				diffs.Merge(rj)
+			}
+			if aj, err := b.addJobs(tx, roleType, qdbid, sortedJobs[i]); err != nil {
+				return nil, errors.Wrapf(err, "adding job: %s", sortedJobs[i])
+			} else {
+				diffs.Merge(aj)
+			}
+		}
+	}
+
+	return diffs, nil
+}
+
+// processFreeJobs assigns all jobs in the free list to a worker.
+func (b *Balancer) processFreeJobs(tx dax.Transaction, roleType dax.RoleType, qdbid dax.QualifiedDatabaseID) (InternalDiffs, error) {
+	diffs := NewInternalDiffs()
+	jobs, err := b.freeJobs.ListJobs(tx, roleType, qdbid)
+	if err != nil {
+		return nil, errors.Wrapf(err, "listing free jobs: %s", roleType)
+	}
+	for _, job := range jobs {
+		if aj, err := b.addDatabaseJobs(tx, roleType, qdbid, job); err != nil {
+			return nil, errors.Wrapf(err, "adding job: %s", job)
+		} else {
+			diffs.Merge(aj)
+		}
+		if err := b.freeJobs.DeleteJob(tx, roleType, qdbid, job); err != nil {
+			return nil, errors.Wrapf(err, "deleting free job: %s", job)
+		}
+	}
+	return diffs, nil
+}
+
+// workerForJob returns the worker currently assigned to the given job.
+func (b *Balancer) workerForJob(tx dax.Transaction, roleType dax.RoleType, qdbid dax.QualifiedDatabaseID, job dax.Job) (dax.Address, bool, error) {
+	workerJobs, err := b.current.WorkersJobs(tx, roleType, qdbid)
+	if err != nil {
+		return "", false, errors.Wrapf(err, "getting workers jobs: (%s) %s", roleType, qdbid)
+	}
+	for _, workerInfo := range workerJobs {
+		jset := dax.NewSet(workerInfo.Jobs...)
+		if jset.Contains(job) {
+			return workerInfo.Address, true, nil
+		}
+	}
+	return "", false, nil
+}
+
+type WorkerJobService interface {
+	WorkersJobs(tx dax.Transaction, roleType dax.RoleType, qdbid dax.QualifiedDatabaseID) ([]dax.WorkerInfo, error)
+
+	WorkerCount(tx dax.Transaction, roleType dax.RoleType, qdbid dax.QualifiedDatabaseID) (int, error)
+	ListWorkers(tx dax.Transaction, roleType dax.RoleType, qdbid dax.QualifiedDatabaseID) (dax.Addresses, error)
+
+	CreateWorker(tx dax.Transaction, roleType dax.RoleType, qdbid dax.QualifiedDatabaseID, addr dax.Address) error
+	DeleteWorker(tx dax.Transaction, roleType dax.RoleType, qdbid dax.QualifiedDatabaseID, addr dax.Address) error
+
+	CreateJobs(tx dax.Transaction, roleType dax.RoleType, qdbid dax.QualifiedDatabaseID, addr dax.Address, job ...dax.Job) error
+	DeleteJob(tx dax.Transaction, roleType dax.RoleType, qdbid dax.QualifiedDatabaseID, addr dax.Address, job dax.Job) error
+	DeleteJobsForTable(tx dax.Transaction, roleType dax.RoleType, qtid dax.QualifiedTableID) (InternalDiffs, error)
+	JobCounts(tx dax.Transaction, roleType dax.RoleType, qdbid dax.QualifiedDatabaseID, addr ...dax.Address) (map[dax.Address]int, error)
+	ListJobs(tx dax.Transaction, roleType dax.RoleType, qdbid dax.QualifiedDatabaseID, addr dax.Address) (dax.Jobs, error)
+
+	DatabaseForWorker(tx dax.Transaction, addr dax.Address) dax.DatabaseKey
+}
+
+type FreeJobService interface {
+	CreateJobs(tx dax.Transaction, roleType dax.RoleType, qdbid dax.QualifiedDatabaseID, job ...dax.Job) error
+	DeleteJob(tx dax.Transaction, roleType dax.RoleType, qdbid dax.QualifiedDatabaseID, job dax.Job) error
+	DeleteJobsForTable(tx dax.Transaction, roleType dax.RoleType, qtid dax.QualifiedTableID) error
+	ListJobs(tx dax.Transaction, roleType dax.RoleType, qdbid dax.QualifiedDatabaseID) (dax.Jobs, error)
+	MergeJobs(tx dax.Transaction, roleType dax.RoleType, qdbid dax.QualifiedDatabaseID, jobs dax.Jobs) error
+}
+
+type FreeWorkerService interface {
+	AddWorkers(tx dax.Transaction, roleType dax.RoleType, addrs ...dax.Address) error
+	RemoveWorker(tx dax.Transaction, roleType dax.RoleType, addr dax.Address) error
+	PopWorkers(tx dax.Transaction, roleType dax.RoleType, num int) ([]dax.Address, error)
+	ListWorkers(tx dax.Transaction, roleType dax.RoleType) (dax.Addresses, error)
+}

dax/mds/controller/balancer/boltdb/balancer.go

@@ -0,0 +1,813 @@
+// Package boltdb contains the boltdb implementation of the Balancer interface.
+package boltdb
+
+import (
+	"bytes"
+	"encoding/json"
+	"fmt"
+	"strings"
+
+	"github.com/featurebasedb/featurebase/v3/dax"
+	"github.com/featurebasedb/featurebase/v3/dax/boltdb"
+	"github.com/featurebasedb/featurebase/v3/dax/mds/controller"
+	"github.com/featurebasedb/featurebase/v3/dax/mds/controller/balancer"
+	"github.com/featurebasedb/featurebase/v3/dax/mds/schemar"
+	"github.com/featurebasedb/featurebase/v3/errors"
+	"github.com/featurebasedb/featurebase/v3/logger"
+)
+
+var (
+	bucketBalancer = boltdb.Bucket("balancer")
+)
+
+// BalancerBuckets defines the buckets used by this package. It can be
+// called during setup to create the buckets ahead of time.
+var BalancerBuckets []boltdb.Bucket = []boltdb.Bucket{
+	bucketBalancer,
+}
+
+// NewBalancer returns a new instance of controller.Balancer.
+func NewBalancer(db *boltdb.DB, schemar schemar.Schemar, logger logger.Logger) controller.Balancer {
+	fjs := newFreeJobService(db)
+	wjs := newWorkerJobService(db, logger)
+	fws := newFreeWorkerService(db)
+
+	return balancer.New(fjs, wjs, fws, schemar, logger)
+}
+
+// Ensure type implements interface.
+var _ balancer.WorkerJobService = (*workerJobService)(nil)
+
+type workerJobService struct {
+	db     *boltdb.DB
+	logger logger.Logger
+}
+
+func newWorkerJobService(db *boltdb.DB, logger logger.Logger) *workerJobService {
+	return &workerJobService{
+		db:     db,
+		logger: logger,
+	}
+}
+
+func (w *workerJobService) WorkersJobs(tx dax.Transaction, roleType dax.RoleType, qdbid dax.QualifiedDatabaseID) ([]dax.WorkerInfo, error) {
+	workerInfos, err := w.getWorkerInfos(tx, roleType, qdbid)
+	if err != nil {
+		return nil, errors.Wrapf(err, "getting worker infos: %s", roleType)
+	}
+
+	return workerInfos, nil
+}
+
+func (w *workerJobService) WorkerCount(tx dax.Transaction, roleType dax.RoleType, qdbid dax.QualifiedDatabaseID) (int, error) {
+	workers, err := w.getWorkers(tx, roleType, qdbid)
+	if err != nil {
+		return 0, errors.Wrapf(err, "getting workers: %s", roleType)
+	}
+
+	return len(workers), nil
+}
+
+func (w *workerJobService) ListWorkers(tx dax.Transaction, roleType dax.RoleType, qdbid dax.QualifiedDatabaseID) (dax.Addresses, error) {
+	return w.getWorkers(tx, roleType, qdbid)
+}
+
+func (w *workerJobService) getWorkers(tx dax.Transaction, roleType dax.RoleType, qdbid dax.QualifiedDatabaseID) (dax.Addresses, error) {
+	txx, ok := tx.(*boltdb.Tx)
+	if !ok {
+		return nil, dax.NewErrInvalidTransaction()
+	}
+
+	c := txx.Bucket(bucketBalancer).Cursor()
+
+	// Deserialize rows into Worker objects.
+	addrs := make(dax.Addresses, 0)
+
+	prefix := []byte(fmt.Sprintf(prefixFmtWorkersDB, roleType, qdbid.Key()))
+	for k, v := c.Seek(prefix); k != nil && bytes.HasPrefix(k, prefix); k, v = c.Next() {
+		if v == nil {
+			w.logger.Printf("nil value for key: %s", k)
+			continue
+		}
+
+		addr, err := keyWorker(k)
+		if err != nil {
+			return nil, errors.Wrapf(err, "getting worker from key: %s", k)
+		}
+
+		addrs = append(addrs, addr)
+	}
+
+	return addrs, nil
+}
+
+func (w *workerJobService) getWorkerInfos(tx dax.Transaction, roleType dax.RoleType, qdbid dax.QualifiedDatabaseID) (dax.WorkerInfos, error) {
+	txx, ok := tx.(*boltdb.Tx)
+	if !ok {
+		return nil, dax.NewErrInvalidTransaction()
+	}
+
+	c := txx.Bucket(bucketBalancer).Cursor()
+
+	// Deserialize rows into WorkerInfo objects.
+	workerInfos := make(dax.WorkerInfos, 0)
+
+	prefix := []byte(fmt.Sprintf(prefixFmtWorkersDB, roleType, qdbid.Key()))
+	for k, v := c.Seek(prefix); k != nil && bytes.HasPrefix(k, prefix); k, v = c.Next() {
+		addr, err := keyWorker(k)
+		if err != nil {
+			return nil, errors.Wrapf(err, "getting worker from key: %s", k)
+		}
+
+		jobs := dax.NewSet[dax.Job]()
+		if v != nil {
+			jobs, err = decodeJobSet(v)
+			if err != nil {
+				return nil, errors.Wrap(err, "decoding job set")
+			}
+		}
+
+		workerInfo := dax.WorkerInfo{
+			Address: addr,
+			Jobs:    jobs.Sorted(),
+		}
+
+		workerInfos = append(workerInfos, workerInfo)
+	}
+
+	return workerInfos, nil
+}
+
+func (w *workerJobService) CreateWorker(tx dax.Transaction, roleType dax.RoleType, qdbid dax.QualifiedDatabaseID, addr dax.Address) error {
+	txx, ok := tx.(*boltdb.Tx)
+	if !ok {
+		return dax.NewErrInvalidTransaction()
+	}
+
+	bkt := txx.Bucket(bucketBalancer)
+	if bkt == nil {
+		return errors.Errorf(boltdb.ErrFmtBucketNotFound, bucketBalancer)
+	}
+
+	// If this worker already exists, don't do anything.
+	wrkr := bkt.Get(workerDBKey(roleType, qdbid, addr))
+	if wrkr != nil {
+		return nil
+	}
+
+	val := []byte("[]")
+	if err := bkt.Put(workerDBKey(roleType, qdbid, addr), val); err != nil {
+		return errors.Wrapf(err, "putting db worker: %s, %s", qdbid, addr)
+	}
+
+	if err := bkt.Put(workerAssignedKey(addr), []byte(qdbid.Key())); err != nil {
+		return errors.Wrapf(err, "putting assigned worker: %s, %s", qdbid, addr)
+	}
+
+	return nil
+}
+
+func (w *workerJobService) DeleteWorker(tx dax.Transaction, roleType dax.RoleType, qdbid dax.QualifiedDatabaseID, addr dax.Address) error {
+	txx, ok := tx.(*boltdb.Tx)
+	if !ok {
+		return dax.NewErrInvalidTransaction()
+	}
+
+	bkt := txx.Bucket(bucketBalancer)
+	if bkt == nil {
+		return errors.Errorf(boltdb.ErrFmtBucketNotFound, bucketBalancer)
+	}
+
+	if err := bkt.Delete(workerDBKey(roleType, qdbid, addr)); err != nil {
+		return errors.Wrapf(err, "deleting node key: %s", workerDBKey(roleType, qdbid, addr))
+	}
+
+	if err := bkt.Delete(workerAssignedKey(addr)); err != nil {
+		return errors.Wrapf(err, "deleting assigned worker: %s", workerAssignedKey(addr))
+	}
+
+	return nil
+}
+
+func (w *workerJobService) CreateJobs(tx dax.Transaction, roleType dax.RoleType, qdbid dax.QualifiedDatabaseID, addr dax.Address, jobs ...dax.Job) error {
+	txx, ok := tx.(*boltdb.Tx)
+	if !ok {
+		return dax.NewErrInvalidTransaction()
+	}
+
+	bkt := txx.Bucket(bucketBalancer)
+	if bkt == nil {
+		return errors.Errorf(boltdb.ErrFmtBucketNotFound, bucketBalancer)
+	}
+
+	jobset := dax.NewSet[dax.Job]()
+	var err error
+
+	// get worker
+	wrkr := bkt.Get(workerDBKey(roleType, qdbid, addr))
+	if wrkr != nil {
+		jobset, err = decodeJobSet(wrkr)
+		if err != nil {
+			return errors.Wrap(err, "decoding job set")
+		}
+	}
+
+	for _, job := range jobs {
+		jobset.Add(job)
+	}
+	val, err := encodeJobSet(jobset)
+	if err != nil {
+		return errors.Wrap(err, "encoding job set")
+	}
+
+	if err := bkt.Put(workerDBKey(roleType, qdbid, addr), val); err != nil {
+		return errors.Wrap(err, "putting worker")
+	}
+
+	return nil
+}
+
+func (w *workerJobService) DeleteJob(tx dax.Transaction, roleType dax.RoleType, qdbid dax.QualifiedDatabaseID, addr dax.Address, job dax.Job) error {
+	txx, ok := tx.(*boltdb.Tx)
+	if !ok {
+		return dax.NewErrInvalidTransaction()
+	}
+
+	bkt := txx.Bucket(bucketBalancer)
+	if bkt == nil {
+		return errors.Errorf(boltdb.ErrFmtBucketNotFound, bucketBalancer)
+	}
+
+	// get worker
+	wrkr := bkt.Get(workerDBKey(roleType, qdbid, addr))
+	if wrkr == nil {
+		return nil
+	}
+
+	jobset, err := decodeJobSet(wrkr)
+	if err != nil {
+		return errors.Wrap(err, "decoding job set")
+	}
+	if !jobset.Contains(job) {
+		return nil
+	}
+
+	jobset.Remove(job)
+	val, err := encodeJobSet(jobset)
+	if err != nil {
+		return errors.Wrap(err, "encoding job set")
+	}
+
+	if err := bkt.Put(workerDBKey(roleType, qdbid, addr), val); err != nil {
+		return errors.Wrap(err, "putting worker")
+	}
+
+	return nil
+}
+
+func (w *workerJobService) DeleteJobsForTable(tx dax.Transaction, roleType dax.RoleType, qtid dax.QualifiedTableID) (balancer.InternalDiffs, error) {
+	txx, ok := tx.(*boltdb.Tx)
+	if !ok {
+		return nil, dax.NewErrInvalidTransaction()
+	}
+
+	bkt := txx.Bucket(bucketBalancer)
+	if bkt == nil {
+		return nil, errors.Errorf(boltdb.ErrFmtBucketNotFound, bucketBalancer)
+	}
+
+	qdbid := qtid.QualifiedDatabaseID
+	prefix := string(qtid.Key())
+
+	workers, err := w.getWorkers(tx, roleType, qdbid)
+	if err != nil {
+		return nil, errors.Wrap(err, "getting workers")
+	}
+
+	idiffs := balancer.NewInternalDiffs()
+	for _, worker := range workers {
+		// get worker
+		wrkr := bkt.Get(workerDBKey(roleType, qdbid, worker))
+		if wrkr == nil {
+			panic("didn't find worker that should... definitely exist")
+		}
+		jobset, err := decodeJobSet(wrkr)
+		if err != nil {
+			return nil, errors.Wrap(err, "decoding job set")
+		}
+
+		jobs := jobset.RemoveByPrefix(prefix)
+		for _, job := range jobs {
+			idiffs.Removed(worker, job)
+		}
+		val, err := encodeJobSet(jobset)
+		if err != nil {
+			return nil, errors.Wrap(err, "encoding job set")
+		}
+
+		if err := bkt.Put(workerDBKey(roleType, qdbid, worker), val); err != nil {
+			return nil, errors.Wrap(err, "putting worker")
+		}
+
+	}
+
+	return idiffs, nil
+}
+
+func (w *workerJobService) ListJobs(tx dax.Transaction, roleType dax.RoleType, qdbid dax.QualifiedDatabaseID, addr dax.Address) (dax.Jobs, error) {
+	txx, ok := tx.(*boltdb.Tx)
+	if !ok {
+		return nil, dax.NewErrInvalidTransaction()
+	}
+
+	bkt := txx.Bucket(bucketBalancer)
+	if bkt == nil {
+		return nil, errors.Errorf(boltdb.ErrFmtBucketNotFound, bucketBalancer)
+	}
+
+	jobset := dax.NewSet[dax.Job]()
+	var err error
+
+	// get worker
+	wrkr := bkt.Get(workerDBKey(roleType, qdbid, addr))
+	if wrkr != nil {
+		jobset, err = decodeJobSet(wrkr)
+		if err != nil {
+			return nil, errors.Wrap(err, "decoding job set")
+		}
+	}
+
+	return jobset.Sorted(), nil
+}
+
+func (w *workerJobService) JobCounts(tx dax.Transaction, roleType dax.RoleType, qdbid dax.QualifiedDatabaseID, addrs ...dax.Address) (map[dax.Address]int, error) {
+	txx, ok := tx.(*boltdb.Tx)
+	if !ok {
+		return nil, dax.NewErrInvalidTransaction()
+	}
+
+	bkt := txx.Bucket(bucketBalancer)
+	if bkt == nil {
+		return nil, errors.Errorf(boltdb.ErrFmtBucketNotFound, bucketBalancer)
+	}
+
+	m := make(map[dax.Address]int)
+
+	for _, addr := range addrs {
+		jobset := dax.NewSet[dax.Job]()
+		var err error
+
+		// get worker
+		wrkr := bkt.Get(workerDBKey(roleType, qdbid, addr))
+		if wrkr != nil {
+			jobset, err = decodeJobSet(wrkr)
+			if err != nil {
+				return nil, errors.Wrap(err, "decoding job set")
+			}
+		}
+
+		m[addr] = len(jobset)
+	}
+
+	return m, nil
+}
+
+func (w *workerJobService) DatabaseForWorker(tx dax.Transaction, addr dax.Address) dax.DatabaseKey {
+	txx, ok := tx.(*boltdb.Tx)
+	if !ok {
+		return "" // TODO(tlt): return error here?
+	}
+
+	bkt := txx.Bucket(bucketBalancer)
+	if bkt == nil {
+		return ""
+	}
+
+	wrkr := bkt.Get(workerAssignedKey(addr))
+
+	return dax.DatabaseKey(wrkr)
+}
+
+// encodeJobSet encode the jobSet into a JSON array of strings.
+func encodeJobSet(jobSet dax.Set[dax.Job]) ([]byte, error) {
+	arr := jobSet.Sorted()
+	b, err := json.Marshal(arr)
+	if err != nil {
+		return nil, errors.Wrap(err, "marshalling json")
+	}
+	return b, nil
+}
+
+// decodeJobSet decode the string (a JSON array of strings) into jobSet.
+func decodeJobSet(v []byte) (dax.Set[dax.Job], error) {
+	var arr []string
+	err := json.Unmarshal(v, &arr)
+	if err != nil {
+		return nil, errors.Wrap(err, "unmarshalling json")
+	}
+
+	js := dax.NewSet[dax.Job]()
+	for _, s := range arr {
+		js.Add(dax.Job(s))
+	}
+
+	return js, nil
+}
+
+// encodeWorkerSet encode the workerSet into a JSON array of strings.
+func encodeWorkerSet(workerSet dax.Set[dax.Address]) ([]byte, error) {
+	arr := workerSet.Sorted()
+	b, err := json.Marshal(arr)
+	if err != nil {
+		return nil, errors.Wrap(err, "marshalling json")
+	}
+	return b, nil
+}
+
+// decodeWorkerSet decode the string (a JSON array of strings) into workerSet.
+func decodeWorkerSet(v []byte) (dax.Set[dax.Address], error) {
+	var arr []string
+	err := json.Unmarshal(v, &arr)
+	if err != nil {
+		return nil, errors.Wrap(err, "unmarshalling json")
+	}
+
+	ws := dax.NewSet[dax.Address]()
+	for _, s := range arr {
+		ws.Add(dax.Address(s))
+	}
+
+	return ws, nil
+}
+
+// Ensure type implements interface.
+var _ balancer.FreeJobService = (*freeJobService)(nil)
+
+type freeJobService struct {
+	db *boltdb.DB
+}
+
+func newFreeJobService(db *boltdb.DB) *freeJobService {
+	return &freeJobService{
+		db: db,
+	}
+}
+
+func (f *freeJobService) CreateJobs(tx dax.Transaction, roleType dax.RoleType, qdbid dax.QualifiedDatabaseID, jobs ...dax.Job) error {
+	return f.MergeJobs(tx, roleType, qdbid, jobs)
+}
+
+func (f *freeJobService) DeleteJob(tx dax.Transaction, roleType dax.RoleType, qdbid dax.QualifiedDatabaseID, job dax.Job) error {
+	txx, ok := tx.(*boltdb.Tx)
+	if !ok {
+		return dax.NewErrInvalidTransaction()
+	}
+
+	bkt := txx.Bucket(bucketBalancer)
+	if bkt == nil {
+		return errors.Errorf(boltdb.ErrFmtBucketNotFound, bucketBalancer)
+	}
+
+	// get free jobs
+	fjs := bkt.Get(freeJobKey(roleType, qdbid))
+	if fjs == nil {
+		return nil
+	}
+
+	jobset, err := decodeJobSet(fjs)
+	if err != nil {
+		return errors.Wrap(err, "decoding job set")
+	}
+	if !jobset.Contains(job) {
+		return nil
+	}
+
+	jobset.Remove(job)
+	val, err := encodeJobSet(jobset)
+	if err != nil {
+		return errors.Wrap(err, "encoding job set")
+	}
+
+	if err := bkt.Put(freeJobKey(roleType, qdbid), val); err != nil {
+		return errors.Wrap(err, "putting free job")
+	}
+
+	return nil
+}
+
+func (f *freeJobService) DeleteJobsForTable(tx dax.Transaction, roleType dax.RoleType, qtid dax.QualifiedTableID) error {
+	txx, ok := tx.(*boltdb.Tx)
+	if !ok {
+		return dax.NewErrInvalidTransaction()
+	}
+
+	bkt := txx.Bucket(bucketBalancer)
+	if bkt == nil {
+		return errors.Errorf(boltdb.ErrFmtBucketNotFound, bucketBalancer)
+	}
+
+	qdbid := qtid.QualifiedDatabaseID
+	prefix := string(qtid.Key())
+
+	// get free jobs
+	fjs := bkt.Get(freeJobKey(roleType, qdbid))
+	if fjs == nil {
+		return nil
+	}
+
+	jobset, err := decodeJobSet(fjs)
+	if err != nil {
+		return errors.Wrap(err, "decoding job set")
+	}
+
+	jobset.RemoveByPrefix(prefix)
+	val, err := encodeJobSet(jobset)
+	if err != nil {
+		return errors.Wrap(err, "encoding job set")
+	}
+
+	if err := bkt.Put(freeJobKey(roleType, qdbid), val); err != nil {
+		return errors.Wrap(err, "putting free job")
+	}
+
+	return nil
+}
+
+func (f *freeJobService) ListJobs(tx dax.Transaction, roleType dax.RoleType, qdbid dax.QualifiedDatabaseID) (dax.Jobs, error) {
+	txx, ok := tx.(*boltdb.Tx)
+	if !ok {
+		return nil, dax.NewErrInvalidTransaction()
+	}
+
+	bkt := txx.Bucket(bucketBalancer)
+	if bkt == nil {
+		return nil, errors.Errorf(boltdb.ErrFmtBucketNotFound, bucketBalancer)
+	}
+
+	jobset := dax.NewSet[dax.Job]()
+	var err error
+
+	// get free jobs
+	fjs := bkt.Get(freeJobKey(roleType, qdbid))
+	if fjs != nil {
+		jobset, err = decodeJobSet(fjs)
+		if err != nil {
+			return nil, errors.Wrap(err, "decoding job set")
+		}
+	}
+
+	return jobset.Sorted(), nil
+}
+
+func (f *freeJobService) MergeJobs(tx dax.Transaction, roleType dax.RoleType, qdbid dax.QualifiedDatabaseID, jobs dax.Jobs) error {
+	txx, ok := tx.(*boltdb.Tx)
+	if !ok {
+		return dax.NewErrInvalidTransaction()
+	}
+
+	bkt := txx.Bucket(bucketBalancer)
+	if bkt == nil {
+		return errors.Errorf(boltdb.ErrFmtBucketNotFound, bucketBalancer)
+	}
+
+	jobset := dax.NewSet[dax.Job]()
+	var err error
+
+	// get free jobs
+	fjs := bkt.Get(freeJobKey(roleType, qdbid))
+	if fjs != nil {
+		jobset, err = decodeJobSet(fjs)
+		if err != nil {
+			return errors.Wrap(err, "decoding job set")
+		}
+	}
+
+	for _, j := range jobs {
+		jobset.Add(j)
+	}
+	val, err := encodeJobSet(jobset)
+	if err != nil {
+		return errors.Wrap(err, "encoding job set")
+	}
+
+	if err := bkt.Put(freeJobKey(roleType, qdbid), val); err != nil {
+		return errors.Wrap(err, "putting free job")
+	}
+
+	return nil
+}
+
+//////////////////////////////////////////////////////
+
+// Ensure type implements interface.
+var _ balancer.FreeWorkerService = (*freeWorkerService)(nil)
+
+type freeWorkerService struct {
+	db *boltdb.DB
+}
+
+func newFreeWorkerService(db *boltdb.DB) *freeWorkerService {
+	return &freeWorkerService{
+		db: db,
+	}
+}
+
+func (f *freeWorkerService) AddWorkers(tx dax.Transaction, roleType dax.RoleType, addres ...dax.Address) error {
+	txx, ok := tx.(*boltdb.Tx)
+	if !ok {
+		return dax.NewErrInvalidTransaction()
+	}
+
+	bkt := txx.Bucket(bucketBalancer)
+	if bkt == nil {
+		return errors.Errorf(boltdb.ErrFmtBucketNotFound, bucketBalancer)
+	}
+
+	workerset := dax.NewSet[dax.Address]()
+	var err error
+
+	// get free workers
+	fws := bkt.Get(freeWorkerKey(roleType))
+	if fws != nil {
+		workerset, err = decodeWorkerSet(fws)
+		if err != nil {
+			return errors.Wrap(err, "decoding worker set")
+		}
+	}
+
+	for _, w := range addres {
+		workerset.Add(w)
+	}
+	val, err := encodeWorkerSet(workerset)
+	if err != nil {
+		return errors.Wrap(err, "encoding worker set")
+	}
+
+	if err := bkt.Put(freeWorkerKey(roleType), val); err != nil {
+		return errors.Wrap(err, "putting free worker")
+	}
+
+	return nil
+}
+
+func (f *freeWorkerService) RemoveWorker(tx dax.Transaction, roleType dax.RoleType, addr dax.Address) error {
+	txx, ok := tx.(*boltdb.Tx)
+	if !ok {
+		return dax.NewErrInvalidTransaction()
+	}
+
+	workers, err := f.ListWorkers(tx, roleType)
+	if err != nil {
+		return errors.Wrap(err, "listing free workers")
+	}
+
+	// Create a workerset containing the free workers which remain after
+	// removing num workers.
+	workerset := dax.NewSet[dax.Address]()
+	for _, w := range workers {
+		workerset.Add(w)
+	}
+
+	if !workerset.Contains(addr) {
+		return nil
+	}
+
+	// Remove the worker.
+	workerset.Remove(addr)
+
+	bkt := txx.Bucket(bucketBalancer)
+	if bkt == nil {
+		return errors.Errorf(boltdb.ErrFmtBucketNotFound, bucketBalancer)
+	}
+
+	val, err := encodeWorkerSet(workerset)
+	if err != nil {
+		return errors.Wrap(err, "encoding worker set")
+	}
+
+	if err := bkt.Put(freeWorkerKey(roleType), val); err != nil {
+		return errors.Wrap(err, "putting free worker")
+	}
+
+	return nil
+}
+
+func (f *freeWorkerService) PopWorkers(tx dax.Transaction, roleType dax.RoleType, num int) ([]dax.Address, error) {
+	txx, ok := tx.(*boltdb.Tx)
+	if !ok {
+		return nil, dax.NewErrInvalidTransaction()
+	}
+
+	workers, err := f.ListWorkers(tx, roleType)
+	if err != nil {
+		return nil, errors.Wrap(err, "listing free workers")
+	}
+
+	if len(workers) < num {
+		return nil, errors.Errorf("not enough free workers to pop: wanted %d, have: %d", num, len(workers))
+	}
+
+	// Get num workers from the list.
+	workersToAssign := workers[0:num]
+
+	// Create a workerset containing the free workers which remain after
+	// removing num workers.
+	workerset := dax.NewSet[dax.Address]()
+	for _, worker := range workers[num:] {
+		workerset.Add(worker)
+	}
+
+	bkt := txx.Bucket(bucketBalancer)
+	if bkt == nil {
+		return nil, errors.Errorf(boltdb.ErrFmtBucketNotFound, bucketBalancer)
+	}
+
+	val, err := encodeWorkerSet(workerset)
+	if err != nil {
+		return nil, errors.Wrap(err, "encoding worker set")
+	}
+
+	if err := bkt.Put(freeWorkerKey(roleType), val); err != nil {
+		return nil, errors.Wrap(err, "putting free worker")
+	}
+
+	return workersToAssign, nil
+}
+
+func (f *freeWorkerService) ListWorkers(tx dax.Transaction, roleType dax.RoleType) (dax.Addresses, error) {
+	txx, ok := tx.(*boltdb.Tx)
+	if !ok {
+		return nil, dax.NewErrInvalidTransaction()
+	}
+
+	bkt := txx.Bucket(bucketBalancer)
+	if bkt == nil {
+		return nil, errors.Errorf(boltdb.ErrFmtBucketNotFound, bucketBalancer)
+	}
+
+	workerset := dax.NewSet[dax.Address]()
+	var err error
+
+	// get free workers
+	fws := bkt.Get(freeWorkerKey(roleType))
+	if fws != nil {
+		workerset, err = decodeWorkerSet(fws)
+		if err != nil {
+			return nil, errors.Wrap(err, "decoding worker set")
+		}
+	}
+
+	return workerset.Sorted(), nil
+}
+
+//////////////////////////////////////////////////////
+
+const (
+	prefixFmtWorkersDB       = "workers/role/%s/db/%s/" // %s - role, dbKey
+	prefixFmtWorkersAssigned = "workers/assigned/"
+
+	prefixFmtFreeJobs    = "freejobs/role/%s/db/%s" // %s - role, dbKey
+	prefixFmtFreeWorkers = "freeworkers/role/%s"    // %s - role
+)
+
+// workerDBKey returns a key based on worker.
+//
+// Format: workers/role/[role]/db/[dbKey]/[worker] = [job1, job2, ...]
+func workerDBKey(roleType dax.RoleType, qdbid dax.QualifiedDatabaseID, addr dax.Address) []byte {
+	key := fmt.Sprintf(prefixFmtWorkersDB+"%s", roleType, qdbid.Key(), addr)
+	return []byte(key)
+}
+
+// workerAssignedKey returns a key based on worker.
+//
+// Format: workers/assigned/[worker] = dbKey
+func workerAssignedKey(addr dax.Address) []byte {
+	key := fmt.Sprintf(prefixFmtWorkersAssigned+"%s", addr)
+	return []byte(key)
+}
+
+// keyWorker gets the worker out of the key.
+func keyWorker(key []byte) (dax.Address, error) {
+	parts := strings.SplitN(string(key), "/", 6)
+	if len(parts) != 6 {
+		return "", errors.New(errors.ErrUncoded, "worker key format expected: `workers/role/[role]/db/[db]/worker`")
+	}
+
+	return dax.Address(parts[5]), nil
+}
+
+// freeJobKey returns a key for all freeJobs.
+//
+// Format: freejobs/role/[role]/db/[dbKey] = [job1, job2, ...]
+func freeJobKey(roleType dax.RoleType, qdbid dax.QualifiedDatabaseID) []byte {
+	key := fmt.Sprintf(prefixFmtFreeJobs, roleType, qdbid.Key())
+	return []byte(key)
+}
+
+// freeWorkerKey returns a key for all freeWorkers.
+//
+// Format: freeworkers/role/[role] = [worker1, worker2, ...]
+func freeWorkerKey(roleType dax.RoleType) []byte {
+	key := fmt.Sprintf(prefixFmtFreeWorkers, roleType)
+	return []byte(key)
+}

dax/mds/controller/balancer/types.go

@@ -1,4 +1,4 @@
-package naive
+package balancer
 
 import (
 	"sort"
@@ -20,36 +20,36 @@ func newJobSetDiffs() jobSetDiffs {
 	}
 }
 
-type InternalDiffs map[dax.Worker]jobSetDiffs
+type InternalDiffs map[dax.Address]jobSetDiffs
 
 func NewInternalDiffs() InternalDiffs {
 	return make(InternalDiffs)
 }
 
-func (d InternalDiffs) Added(worker dax.Worker, job dax.Job) {
-	if _, ok := d[worker]; !ok {
-		d[worker] = newJobSetDiffs()
+func (d InternalDiffs) Added(address dax.Address, job dax.Job) {
+	if _, ok := d[address]; !ok {
+		d[address] = newJobSetDiffs()
 	}
 
 	// Before adding the job, make sure we haven't indicated that it has been
 	// removed prior to this. If it has, we need to invalidate that "remove"
 	// instruction.
-	d[worker].removed.Remove(job)
+	d[address].removed.Remove(job)
 
-	d[worker].added.Add(job)
+	d[address].added.Add(job)
 }
 
-func (d InternalDiffs) Removed(worker dax.Worker, job dax.Job) {
-	if _, ok := d[worker]; !ok {
-		d[worker] = newJobSetDiffs()
+func (d InternalDiffs) Removed(address dax.Address, job dax.Job) {
+	if _, ok := d[address]; !ok {
+		d[address] = newJobSetDiffs()
 	}
 
 	// Before removing the job, make sure we haven't indicated that it has been
 	// added prior to this. If it has, we need to invalidate that "add"
 	// instruction.
-	d[worker].added.Remove(job)
+	d[address].added.Remove(job)
 
-	d[worker].removed.Add(job)
+	d[address].removed.Add(job)
 }
 
 func (d InternalDiffs) Merge(d2 InternalDiffs) {
@@ -69,7 +69,7 @@ func (d InternalDiffs) Output() []dax.WorkerDiff {
 
 	i := 0
 	for k, v := range d {
-		out[i].WorkerID = k
+		out[i].Address = k
 		out[i].AddedJobs = v.added.Sorted()
 		out[i].RemovedJobs = v.removed.Sorted()
 		i++

dax/mds/controller/config.go

@@ -11,10 +11,10 @@ import (
 type NewBalancerFn func(string, logger.Logger) Balancer
 
 type Config struct {
-	Director          Director
-	Schemar           schemar.Schemar
-	ComputeBalancer   Balancer
-	TranslateBalancer Balancer
+	Director Director
+	Schemar  schemar.Schemar
+
+	Balancer Balancer
 
 	StorageMethod string
 	BoltDB        *boltdb.DB

dax/mds/controller/errors.go

@@ -96,6 +96,6 @@ func NewErrInvalidRequest(msg string) error {
 func NewErrUnassignedJobs(jobs []dax.Job) error {
 	return errors.New(
 		ErrCodeUnassignedJobs,
-		fmt.Sprintf("found %d unassigned jobs", len(jobs)),
+		fmt.Sprintf("found %d unassigned jobs: %+v", len(jobs), jobs),
 	)
 }

dax/mds/controller/naive/balancer.go

@@ -1,599 +0,0 @@
-// Package naive contains a naive implementation of the Balancer interface.
-package naive
-
-import (
-	"context"
-	"fmt"
-	"log"
-	"math"
-	"sort"
-	"strings"
-	"sync"
-	"time"
-
-	"github.com/featurebasedb/featurebase/v3/dax"
-	"github.com/featurebasedb/featurebase/v3/dax/mds/controller"
-	"github.com/featurebasedb/featurebase/v3/errors"
-	"github.com/featurebasedb/featurebase/v3/logger"
-)
-
-// Ensure type implements interface.
-var _ controller.Balancer = (*Balancer)(nil)
-
-// Balancer is a naive implementation of the controller.Balancer interface. It
-// helps manage the relationships between workers and jobs. The logic it uses to
-// balance jobs across workers is very simple; it bases everything off the
-// number of workers and number of jobs. It does not take anything else (such as
-// job size, worker capabilities, etc) into consideration.
-type Balancer struct {
-	mu sync.RWMutex
-
-	// name is used in logging to help identify the balancer responsible for the
-	// log.
-	name string
-
-	// current represents the current state of worker/job assigments.
-	current WorkerJobService
-
-	// freeJobs is the set of jobs which have yet to be assigned to a worker.
-	// This could be because there are no available workers, or because a worker
-	// has been removed and the jobs for which it was responsible have yet to be
-	// reassigned.
-	freeJobs FreeJobService
-
-	logger logger.Logger
-}
-
-type WorkerJobService interface {
-	WorkersJobs(ctx context.Context, balancerName string) ([]dax.WorkerInfo, error)
-
-	WorkerCount(ctx context.Context, balancerName string) (int, error)
-	ListWorkers(ctx context.Context, balancerName string) (dax.Workers, error)
-	WorkerExists(ctx context.Context, balancerName string, worker dax.Worker) (bool, error)
-	CreateWorker(ctx context.Context, balancerName string, worker dax.Worker) error
-	DeleteWorker(ctx context.Context, balancerName string, worker dax.Worker) error
-
-	CreateJobs(ctx context.Context, balancerName string, worker dax.Worker, job ...dax.Job) error
-	DeleteJob(ctx context.Context, balancerName string, worker dax.Worker, job dax.Job) error
-	DeleteJobs(ctx context.Context, balancerName, prefix string) (InternalDiffs, error)
-	JobCounts(ctx context.Context, balancerName string, worker ...dax.Worker) (map[dax.Worker]int, error)
-	ListJobs(ctx context.Context, balancerName string, worker dax.Worker) (dax.Jobs, error)
-}
-
-// TODO: I don't think all these method names need "Free" in them.
-type FreeJobService interface {
-	CreateFreeJobs(ctx context.Context, balancerName string, job ...dax.Job) error
-	DeleteFreeJob(ctx context.Context, balancerName string, job dax.Job) error
-	ListFreeJobs(ctx context.Context, balancerName string) (dax.Jobs, error)
-	MergeFreeJobs(ctx context.Context, balancerName string, jobs dax.Jobs) error
-	DeleteFreeJobs(ctx context.Context, balancerName, prefix string) error
-}
-
-// New returns a new instance of Balancer.
-func New(name string, fjs FreeJobService, wjs WorkerJobService, logger logger.Logger) *Balancer {
-	return &Balancer{
-		name:     name,
-		current:  wjs,
-		freeJobs: fjs,
-		logger:   logger,
-	}
-}
-
-// AddWorker adds a worker to the Balancer's worker pool. This may cause the
-// Balancer to assign existing jobs that are currently in the free list to the
-// worker. Also, the worker will immediately be available for assignments of new
-// jobs.
-func (b *Balancer) AddWorker(ctx context.Context, worker fmt.Stringer) ([]dax.WorkerDiff, error) {
-	b.logger.Debugf("%s: AddWorker(%s)", b.name, worker.String())
-	b.mu.Lock()
-	defer b.mu.Unlock()
-
-	diff, err := b.addWorker(ctx, dax.Worker(worker.String()))
-	if err != nil {
-		return nil, errors.Wrap(err, "adding worker")
-	}
-
-	return diff.Output(), nil
-}
-
-func (b *Balancer) addWorker(ctx context.Context, worker dax.Worker) (InternalDiffs, error) {
-	// If this worker already exists, don't do anything.
-	if exists, err := b.current.WorkerExists(ctx, b.name, worker); err != nil {
-		return nil, errors.Wrap(err, "checking if worker exists")
-	} else if exists {
-		return InternalDiffs{}, nil
-	}
-
-	if err := b.current.CreateWorker(ctx, b.name, worker); err != nil {
-		return nil, errors.Wrap(err, "creating worker")
-	}
-
-	// Process the freeJobs.
-	return b.processFreeJobs(ctx)
-}
-
-// ReplaceWorker is meant to avoid the job re-assignment caused by performing a
-// RemoveWorker followed by an AddWorker. In this case, it does both in one step
-// so that it's more likely that the jobs will just get transferred directly
-// over. NOT IMPLEMENTED YET.
-// func (b *Balancer) ReplaceWorker(fromWorker string, toWorker string) []WorkerDiff {
-// 	b.mu.Lock()
-// 	defer b.mu.Unlock()
-
-// 	return []WorkerDiff{}
-// }
-
-// RemoveWorker removes a worker from the worker pool and moves any of its
-// currently assigned jobs to the free list. If the intention is to remove a
-// worker and reassign its jobs to other workers, then RemoveWorker() should be
-// followed by Balance().
-func (b *Balancer) RemoveWorker(ctx context.Context, worker fmt.Stringer) ([]dax.WorkerDiff, error) {
-	b.mu.Lock()
-	defer b.mu.Unlock()
-
-	diff, err := b.removeWorker(ctx, dax.Worker(worker.String()))
-	if err != nil {
-		return nil, errors.Wrap(err, "removing worker")
-	}
-
-	return diff.Output(), nil
-}
-
-func (b *Balancer) removeWorker(ctx context.Context, worker dax.Worker) (InternalDiffs, error) {
-	// If this worker doesn't exist, don't do anything else.
-	if exists, err := b.current.WorkerExists(ctx, b.name, worker); err != nil {
-		return nil, errors.Wrap(err, "checking if worker exists")
-	} else if !exists {
-		return InternalDiffs{}, nil
-	}
-
-	jobs, err := b.current.ListJobs(ctx, b.name, worker)
-	if err != nil {
-		return nil, errors.Wrap(err, "listing jobs")
-	}
-
-	// Before removing the worker, mark its jobs as free.
-	if err := b.freeJobs.MergeFreeJobs(ctx, b.name, jobs); err != nil {
-		return nil, errors.Wrap(err, "merging free jobs")
-	}
-
-	// Remove the worker.
-	if err := b.current.DeleteWorker(ctx, b.name, worker); err != nil {
-		return nil, errors.Wrap(err, "deleting worker")
-	}
-
-	// Even though this may not be useful to the caller (for example, in the
-	// case where the worker has died and no longer exists), return the diffs
-	// which represent the removal of jobs from the worker.
-	diff := NewInternalDiffs()
-	for _, job := range jobs {
-		diff.Removed(worker, job)
-	}
-
-	return diff, nil
-}
-
-// AddJobs adds one or more jobs to an existing worker. If there are no existing
-// workers, the jobs are placed into the free list and will be assigned to a
-// worker once one becomes available.
-func (b *Balancer) AddJobs(ctx context.Context, jobs ...fmt.Stringer) ([]dax.WorkerDiff, error) {
-	start := time.Now()
-	defer func() {
-		log.Printf("ELAPSED: Balancer.AddJob: %v", time.Since(start))
-	}()
-
-	jobsToAdd := make([]dax.Job, 0, len(jobs))
-	for _, job := range jobs {
-		jobsToAdd = append(jobsToAdd, dax.Job(job.String()))
-	}
-
-	if len(jobsToAdd) == 1 {
-		b.logger.Debugf("%s: AddJobs (%s)", b.name, jobsToAdd[0])
-	} else {
-		b.logger.Debugf("%s: AddJobs (%d)", b.name, len(jobsToAdd))
-	}
-
-	b.mu.Lock()
-	defer b.mu.Unlock()
-
-	diff, err := b.addJobs(ctx, jobsToAdd...)
-	if err != nil {
-		return nil, errors.Wrap(err, "adding job")
-	}
-
-	return diff.Output(), nil
-}
-
-func (b *Balancer) addJobs(ctx context.Context, jobs ...dax.Job) (InternalDiffs, error) {
-	if cnt, err := b.current.WorkerCount(ctx, b.name); err != nil {
-		return nil, errors.Wrap(err, "getting worker count")
-	} else if cnt == 0 {
-		if err := b.freeJobs.CreateFreeJobs(ctx, b.name, jobs...); err != nil {
-			return nil, errors.Wrap(err, "creating free job")
-		}
-		// TODO: we might want to inform the user that a job is in the free list
-		// because there are no workers.
-		return InternalDiffs{}, nil
-	}
-
-	workerJobs, err := b.current.WorkersJobs(ctx, b.name)
-	if err != nil {
-		return nil, errors.Wrapf(err, "getting workers jobs: %s", b.name)
-	}
-	jset := dax.NewSet[dax.Job]()
-	for _, workerInfo := range workerJobs {
-		jset.Merge(dax.NewSet(workerInfo.Jobs...))
-	}
-
-	workerIDs := make(dax.Workers, 0, len(workerJobs))
-	jobCounts := make(map[dax.Worker]int, 0)
-	for _, v := range workerJobs {
-		workerIDs = append(workerIDs, v.ID)
-		jobCounts[v.ID] = len(v.Jobs)
-	}
-
-	diffs := NewInternalDiffs()
-
-	jobsToCreate := make(map[dax.Worker][]dax.Job)
-
-	for _, job := range jobs {
-		// Skip any job that already exists.
-		if jset.Contains(job) {
-			continue
-		}
-
-		// Find the worker with the fewest number of jobs and assign it this job.
-		var lowCount int = math.MaxInt
-		var lowWorker dax.Worker
-
-		// We loop over workerIDs here instead of jobCounts because jobCounts is
-		// a map and it can return results in an unexpected order, which is a
-		// problem for testing.
-		for _, worker := range workerIDs {
-			jobCount := jobCounts[worker]
-			if jobCount < lowCount {
-				lowCount = jobCount
-				lowWorker = worker
-			}
-		}
-
-		jobsToCreate[lowWorker] = append(jobsToCreate[lowWorker], job)
-		jobCounts[lowWorker]++
-	}
-
-	for worker, jobs := range jobsToCreate {
-		if err := b.current.CreateJobs(ctx, b.name, worker, jobs...); err != nil {
-			return nil, errors.Wrap(err, "creating job")
-		}
-		for _, job := range jobs {
-			diffs.Added(worker, job)
-		}
-	}
-
-	return diffs, nil
-}
-
-// RemoveJob removes a job from the worker to which is was assigned. If the job
-// is not currently assigned to a worker, but it is in the free list, then it
-// will be removed from the free list.
-func (b *Balancer) RemoveJob(ctx context.Context, job fmt.Stringer) ([]dax.WorkerDiff, error) {
-	b.mu.Lock()
-	defer b.mu.Unlock()
-
-	diff, err := b.removeJob(ctx, dax.Job(job.String()))
-	if err != nil {
-		return nil, errors.Wrapf(err, "removing job: %s", job)
-	}
-
-	return diff.Output(), nil
-}
-
-func (b *Balancer) RemoveJobs(ctx context.Context, prefix string) ([]dax.WorkerDiff, error) {
-	b.mu.Lock()
-	defer b.mu.Unlock()
-
-	idiffs, err := b.current.DeleteJobs(ctx, b.name, prefix)
-	if err != nil {
-		return nil, errors.Wrap(err, "deleting worker jobs")
-	}
-	if err := b.freeJobs.DeleteFreeJobs(ctx, b.name, prefix); err != nil {
-		return nil, errors.Wrap(err, "deleting free jobs")
-	}
-	return idiffs.Output(), nil
-}
-
-func (b *Balancer) removeJob(ctx context.Context, job dax.Job) (InternalDiffs, error) {
-	if worker, ok, err := b.workerForJob(ctx, job); err != nil {
-		return nil, errors.Wrapf(err, "getting worker for job: %s", job)
-	} else if ok {
-		if err := b.current.DeleteJob(ctx, b.name, worker, job); err != nil {
-			return nil, errors.Wrapf(err, "deleting job: %s", job)
-		}
-
-		diffs := NewInternalDiffs()
-		diffs.Removed(worker, job)
-
-		return diffs, nil
-	}
-
-	// Just in case the job is in the free list (and wasn't assigned to a
-	// worker), remove it; there's no need to provide a diff. There should never
-	// be a case where the same job is both in the free list and assigned to a
-	// worker.
-	if err := b.freeJobs.DeleteFreeJob(ctx, b.name, job); err != nil {
-		return nil, errors.Wrapf(err, "deleting free job: %s", job)
-	}
-
-	return InternalDiffs{}, nil
-}
-
-// Balance ensures that all jobs are being handled by a worker by assigning jobs
-// in the free list to workers, and by moving job assignments around in order to
-// balance the load on workers.
-func (b *Balancer) Balance(ctx context.Context) ([]dax.WorkerDiff, error) {
-	b.mu.Lock()
-	defer b.mu.Unlock()
-
-	// If there are no workers, we can't properly balance.
-	if cnt, err := b.current.WorkerCount(ctx, b.name); err != nil {
-		return nil, errors.Wrapf(err, "getting worker count: %s", b.name)
-	} else if cnt == 0 {
-		return []dax.WorkerDiff{}, nil
-	}
-
-	// Process the freeJobs.
-	diffs, err := b.processFreeJobs(ctx)
-	if err != nil {
-		return nil, errors.Wrapf(err, "processing free jobs: %s", b.name)
-	}
-
-	// Balance the jobs among workers.
-	diff, err := b.balance(ctx, diffs)
-	if err != nil {
-		return nil, errors.Wrap(err, "balancing jobs")
-	}
-
-	return diff.Output(), nil
-}
-
-// balance moves jobs among workers with the goal of having an equal number of
-// jobs per worker. This method takes an `internalDiffs` as input for cases
-// where some action has preceeded this call which also resulted in
-// `internalDiffs`. Instead of having this method take a value, we could rely on
-// the internalDiffs.merge() method, but we would need to modify that method to
-// be smarter about the order in which it applies the add/remove operations.
-// Until that's in place, we'll pass in a value here.
-func (b *Balancer) balance(ctx context.Context, diffs InternalDiffs) (InternalDiffs, error) {
-	numWorkers, err := b.current.WorkerCount(ctx, b.name)
-	if err != nil {
-		return nil, errors.Wrapf(err, "getting worker count: %s", b.name)
-	}
-	numJobs := 0
-	if workers, err := b.current.ListWorkers(ctx, b.name); err != nil {
-		return nil, errors.Wrapf(err, "listing workers: %s", b.name)
-	} else {
-		for _, worker := range workers {
-			jobCounts, err := b.current.JobCounts(ctx, b.name, worker)
-			if err != nil {
-				return nil, errors.Wrapf(err, "getting job count: %s", worker)
-			}
-			numJobs += jobCounts[worker]
-		}
-	}
-
-	minJobsPerWorker := numJobs / numWorkers
-	numWorkersAboveMin := numJobs % numWorkers
-
-	// sortedWorkerInfos is used now in order to guarantee a sort order.
-	sortedWorkerInfos, err := b.currentState(ctx, true)
-	if err != nil {
-		return nil, errors.Wrapf(err, "getting current state: %s", b.name)
-	}
-
-	// Loop through each worker, and if the number of jobs for the worker
-	// exceeds the target, then remove the job and add it back (which is
-	// effectively how we rebalance a job).
-	for i, workerInfo := range sortedWorkerInfos {
-		numTargetJobs := minJobsPerWorker
-		if i < numWorkersAboveMin {
-			numTargetJobs += 1
-		}
-
-		jobCounts, err := b.current.JobCounts(ctx, b.name, workerInfo.ID)
-		if err != nil {
-			return nil, errors.Wrapf(err, "getting job count: %s", workerInfo.ID)
-		}
-		numCurrentJobs := jobCounts[workerInfo.ID]
-
-		// If we don't need to remove jobs from this worker, then just continue
-		// on to the next worker.
-		if numCurrentJobs <= numTargetJobs {
-			continue
-		}
-
-		sortedJobs, err := b.current.ListJobs(ctx, b.name, workerInfo.ID)
-		if err != nil {
-			return nil, errors.Wrapf(err, "listing jobs: %s", workerInfo.ID)
-		}
-
-		// Remove the extra jobs from the end of the list, and add them back
-		// again (which should place them on a worker with fewer jobs).
-		for i := numCurrentJobs - 1; i >= numTargetJobs; i-- {
-			if rj, err := b.removeJob(ctx, sortedJobs[i]); err != nil {
-				return nil, errors.Wrapf(err, "removing job: %s", sortedJobs[i])
-			} else {
-				diffs.Merge(rj)
-			}
-			if aj, err := b.addJobs(ctx, sortedJobs[i]); err != nil {
-				return nil, errors.Wrapf(err, "adding job: %s", sortedJobs[i])
-			} else {
-				diffs.Merge(aj)
-			}
-		}
-	}
-
-	return diffs, nil
-}
-
-// CurrentState returns the current state of worker and job assignments. Note
-// that there could be unassigned jobs which are not captured in this output.
-// Calling Balance() would force any unassigned jobs to be assigned (assuming
-// there is at least one worker), and the output would then reflect that.
-func (b *Balancer) CurrentState(ctx context.Context) ([]dax.WorkerInfo, error) {
-	b.mu.RLock()
-	defer b.mu.RUnlock()
-
-	return b.currentState(ctx, true)
-}
-
-func (b *Balancer) currentState(ctx context.Context, sorted bool) ([]dax.WorkerInfo, error) {
-	return b.current.WorkersJobs(ctx, b.name)
-}
-
-// WorkerState returns the current state of job assignments for a given worker.
-func (b *Balancer) WorkerState(ctx context.Context, worker dax.Worker) (dax.WorkerInfo, error) {
-	b.mu.RLock()
-	defer b.mu.RUnlock()
-
-	return b.workerState(ctx, worker)
-}
-
-func (b *Balancer) workerState(ctx context.Context, worker dax.Worker) (dax.WorkerInfo, error) {
-	if exists, err := b.current.WorkerExists(ctx, b.name, worker); err != nil {
-		return dax.WorkerInfo{}, errors.Wrapf(err, "checking worker exists: %s", worker)
-	} else if !exists {
-		return dax.WorkerInfo{
-			ID: dax.Worker(worker),
-		}, nil
-	}
-
-	jobs, err := b.current.ListJobs(ctx, b.name, worker)
-	if err != nil {
-		return dax.WorkerInfo{}, errors.Wrapf(err, "listing jobs: %s", worker)
-	}
-
-	return dax.WorkerInfo{
-		ID:   dax.Worker(worker),
-		Jobs: jobs,
-	}, nil
-}
-
-// WorkersForJobs returns the list of workers for the given jobs. If a given job
-// is not currently assigned to a worker, it will be ignored.
-func (b *Balancer) WorkersForJobs(ctx context.Context, jobs []dax.Job) ([]dax.WorkerInfo, error) {
-	b.mu.RLock()
-	defer b.mu.RUnlock()
-
-	return b.workersForJobs(ctx, jobs)
-}
-
-func (b *Balancer) workersForJobs(ctx context.Context, jobs []dax.Job) ([]dax.WorkerInfo, error) {
-	out := make(map[dax.Worker]dax.Set[dax.Job])
-
-	workerJobs, err := b.current.WorkersJobs(ctx, b.name)
-	if err != nil {
-		return nil, errors.Wrapf(err, "getting worker jobs: %s", b.name)
-	}
-	for _, workerInfo := range workerJobs {
-		jset := dax.NewSet(workerInfo.Jobs...)
-
-		matches := dax.NewSet[dax.Job]()
-		for _, job := range jobs {
-			if jset.Contains(job) {
-				matches.Add(job)
-			}
-		}
-
-		if len(matches) > 0 {
-			out[workerInfo.ID] = matches
-		}
-	}
-
-	workers := make([]dax.WorkerInfo, len(out))
-
-	i := 0
-	for w, jset := range out {
-		workers[i] = dax.WorkerInfo{
-			ID:   dax.Worker(w),
-			Jobs: jset.Sorted(),
-		}
-		i++
-	}
-
-	sort.Sort(dax.WorkerInfos(workers))
-
-	return workers, nil
-}
-
-func (b *Balancer) WorkersForJobPrefix(ctx context.Context, prefix string) ([]dax.WorkerInfo, error) {
-	b.mu.RLock()
-	defer b.mu.RUnlock()
-
-	jobs, err := b.freeJobs.ListFreeJobs(ctx, b.name)
-	if err != nil {
-		return nil, errors.Wrap(err, "listing free jobs")
-	}
-	for _, job := range jobs {
-		if strings.HasPrefix(string(job), prefix) {
-			return nil, errors.Errorf("found free job '%s' matching prefix '%s'", job, prefix)
-		}
-	}
-
-	workerJobs, err := b.current.WorkersJobs(ctx, b.name)
-	if err != nil {
-		return nil, errors.Wrapf(err, "getting worker jobs: %s", b.name)
-	}
-
-	result := make([]dax.WorkerInfo, 0)
-	for _, workerInfo := range workerJobs {
-		matchedJobs := make([]dax.Job, 0)
-		for _, job := range workerInfo.Jobs {
-			if strings.HasPrefix(string(job), prefix) {
-				matchedJobs = append(matchedJobs, job)
-			}
-		}
-		if len(matchedJobs) > 0 {
-			result = append(result, dax.WorkerInfo{
-				ID:   workerInfo.ID,
-				Jobs: matchedJobs,
-			})
-		}
-	}
-	return result, nil
-
-}
-
-// processFreeJobs assigns all jobs in the free list to a worker.
-func (b *Balancer) processFreeJobs(ctx context.Context) (InternalDiffs, error) {
-	diffs := NewInternalDiffs()
-	jobs, err := b.freeJobs.ListFreeJobs(ctx, b.name)
-	if err != nil {
-		return nil, errors.Wrapf(err, "listing free jobs: %s", b.name)
-	}
-	for _, job := range jobs {
-		if aj, err := b.addJobs(ctx, job); err != nil {
-			return nil, errors.Wrapf(err, "adding job: %s", job)
-		} else {
-			diffs.Merge(aj)
-		}
-		if err := b.freeJobs.DeleteFreeJob(ctx, b.name, job); err != nil {
-			return nil, errors.Wrapf(err, "deleting free job: %s", job)
-		}
-	}
-	return diffs, nil
-}
-
-// workerForJob returns the worker currently assigned to the given job.
-func (b *Balancer) workerForJob(ctx context.Context, job dax.Job) (dax.Worker, bool, error) {
-	workerJobs, err := b.current.WorkersJobs(ctx, b.name)
-	if err != nil {
-		return "", false, errors.Wrapf(err, "getting workers jobs: %s", b.name)
-	}
-	for _, workerInfo := range workerJobs {
-		jset := dax.NewSet(workerInfo.Jobs...)
-		if jset.Contains(job) {
-			return workerInfo.ID, true, nil
-		}
-	}
-	return "", false, nil
-}

dax/mds/controller/naive/balancer_test.go

@@ -1,751 +0,0 @@
-package naive_test
-
-import (
-	"context"
-	"fmt"
-	"os"
-	"testing"
-
-	"github.com/featurebasedb/featurebase/v3/dax"
-	daxbolt "github.com/featurebasedb/featurebase/v3/dax/boltdb"
-	"github.com/featurebasedb/featurebase/v3/dax/mds/controller/naive/boltdb"
-	testbolt "github.com/featurebasedb/featurebase/v3/dax/test/boltdb"
-	"github.com/featurebasedb/featurebase/v3/logger"
-	"github.com/stretchr/testify/assert"
-)
-
-func newBoltBalancer(t *testing.T) (*daxbolt.DB, func()) {
-	db := testbolt.MustOpenDB(t)
-	assert.NoError(t, db.InitializeBuckets(boltdb.NaiveBalancerBuckets...))
-
-	return db, func() {
-		testbolt.MustCloseDB(t, db)
-		testbolt.CleanupDB(t, db.Path())
-	}
-}
-
-func TestBalancer(t *testing.T) {
-	ctx := context.Background()
-	t.Run("SingleWorker", func(t *testing.T) {
-		db, cleanup := newBoltBalancer(t)
-		defer cleanup()
-		bal := boltdb.NewBalancer("test", db, logger.NewStandardLogger(os.Stderr))
-		// addJob is a wrapper around bal.AddJobs() which we added when the
-		// function signature of bal.AddJobs changed to take multiple jobs (and
-		// it therefore no longer satisfied the fn type in this test).
-		addJob := func(ctx context.Context, job fmt.Stringer) ([]dax.WorkerDiff, error) {
-			return bal.AddJobs(ctx, job)
-		}
-		tests := []struct {
-			fn       func(context.Context, fmt.Stringer) ([]dax.WorkerDiff, error)
-			input    string
-			expDiff  []dax.WorkerDiff
-			expState []dax.WorkerInfo
-		}{
-			{
-				// Add job.
-				fn:       addJob,
-				input:    "p2",
-				expDiff:  []dax.WorkerDiff{},
-				expState: []dax.WorkerInfo{},
-			},
-			{
-				// Add worker.
-				fn:    bal.AddWorker,
-				input: "n1",
-				expDiff: []dax.WorkerDiff{
-					{
-						WorkerID:    "n1",
-						AddedJobs:   []dax.Job{"p2"},
-						RemovedJobs: []dax.Job{},
-					},
-				},
-				expState: []dax.WorkerInfo{
-					{
-						ID:   "n1",
-						Jobs: []dax.Job{"p2"},
-					},
-				},
-			},
-			{
-				// Add another job out of order.
-				fn:    addJob,
-				input: "p1",
-				expDiff: []dax.WorkerDiff{
-					{
-						WorkerID:    "n1",
-						AddedJobs:   []dax.Job{"p1"},
-						RemovedJobs: []dax.Job{},
-					},
-				},
-				expState: []dax.WorkerInfo{
-					{
-						ID:   "n1",
-						Jobs: []dax.Job{"p1", "p2"},
-					},
-				},
-			},
-			{
-				// Add another job.
-				fn:    addJob,
-				input: "p3",
-				expDiff: []dax.WorkerDiff{
-					{
-						WorkerID:    "n1",
-						AddedJobs:   []dax.Job{"p3"},
-						RemovedJobs: []dax.Job{},
-					},
-				},
-				expState: []dax.WorkerInfo{
-					{
-						ID:   "n1",
-						Jobs: []dax.Job{"p1", "p2", "p3"},
-					},
-				},
-			},
-			{
-				// Add a duplicate job.
-				fn:      addJob,
-				input:   "p2",
-				expDiff: []dax.WorkerDiff{},
-				expState: []dax.WorkerInfo{
-					{
-						ID:   "n1",
-						Jobs: []dax.Job{"p1", "p2", "p3"},
-					},
-				},
-			},
-		}
-		for i, test := range tests {
-			t.Run(fmt.Sprintf("test-%d", i), func(t *testing.T) {
-				diff, err := test.fn(ctx, newStringWrapper(test.input))
-				assert.NoError(t, err)
-				assert.Equal(t, test.expDiff, diff)
-
-				cs, err := bal.CurrentState(ctx)
-				assert.NoError(t, err)
-				assert.Equal(t, test.expState, cs)
-			})
-		}
-	})
-
-	t.Run("MultipleWorkers", func(t *testing.T) {
-		db, cleanup := newBoltBalancer(t)
-		defer cleanup()
-		bal := boltdb.NewBalancer("test", db, logger.NewStandardLogger(os.Stderr))
-		// addJob is a wrapper around bal.AddJobs() which we added when the
-		// function signature of bal.AddJobs changed to take multiple jobs (and
-		// it therefore no longer satisfied the fn type in this test).
-		addJob := func(ctx context.Context, job fmt.Stringer) ([]dax.WorkerDiff, error) {
-			return bal.AddJobs(ctx, job)
-		}
-		tests := []struct {
-			fn       func(context.Context, fmt.Stringer) ([]dax.WorkerDiff, error)
-			input    string
-			balance  bool
-			expDiff  []dax.WorkerDiff
-			expState []dax.WorkerInfo
-		}{
-			{
-				// Balance when empty.
-				balance:  true,
-				expDiff:  []dax.WorkerDiff{},
-				expState: []dax.WorkerInfo{},
-			},
-			{
-				// Add worker.
-				fn:      bal.AddWorker,
-				input:   "n2",
-				expDiff: []dax.WorkerDiff{},
-				expState: []dax.WorkerInfo{
-					{
-						ID:   "n2",
-						Jobs: []dax.Job{},
-					},
-				},
-			},
-			{
-				// Add worker again.
-				fn:      bal.AddWorker,
-				input:   "n2",
-				expDiff: []dax.WorkerDiff{},
-				expState: []dax.WorkerInfo{
-					{
-						ID:   "n2",
-						Jobs: []dax.Job{},
-					},
-				},
-			},
-			{
-				// Add a second worker.
-				fn:      bal.AddWorker,
-				input:   "n1",
-				expDiff: []dax.WorkerDiff{},
-				expState: []dax.WorkerInfo{
-					{
-						ID:   "n1",
-						Jobs: []dax.Job{},
-					},
-					{
-						ID:   "n2",
-						Jobs: []dax.Job{},
-					},
-				},
-			},
-			{
-				// Add job.
-				fn:    addJob,
-				input: "p2",
-				expDiff: []dax.WorkerDiff{
-					{
-						WorkerID:    "n1",
-						AddedJobs:   []dax.Job{"p2"},
-						RemovedJobs: []dax.Job{},
-					},
-				},
-				expState: []dax.WorkerInfo{
-					{
-						ID:   "n1",
-						Jobs: []dax.Job{"p2"},
-					},
-					{
-						ID:   "n2",
-						Jobs: []dax.Job{},
-					},
-				},
-			},
-			{
-				// Add job.
-				fn:    addJob,
-				input: "p3",
-				expDiff: []dax.WorkerDiff{
-					{
-						WorkerID:    "n2",
-						AddedJobs:   []dax.Job{"p3"},
-						RemovedJobs: []dax.Job{},
-					},
-				},
-				expState: []dax.WorkerInfo{
-					{
-						ID:   "n1",
-						Jobs: []dax.Job{"p2"},
-					},
-					{
-						ID:   "n2",
-						Jobs: []dax.Job{"p3"},
-					},
-				},
-			},
-			{
-				// Add job.
-				fn:    addJob,
-				input: "p1",
-				expDiff: []dax.WorkerDiff{
-					{
-						WorkerID:    "n1",
-						AddedJobs:   []dax.Job{"p1"},
-						RemovedJobs: []dax.Job{},
-					},
-				},
-				expState: []dax.WorkerInfo{
-					{
-						ID:   "n1",
-						Jobs: []dax.Job{"p1", "p2"},
-					},
-					{
-						ID:   "n2",
-						Jobs: []dax.Job{"p3"},
-					},
-				},
-			},
-			{
-				// Add a third worker.
-				fn:      bal.AddWorker,
-				input:   "n0",
-				expDiff: []dax.WorkerDiff{},
-				expState: []dax.WorkerInfo{
-					{
-						ID:   "n0",
-						Jobs: []dax.Job{},
-					},
-					{
-						ID:   "n1",
-						Jobs: []dax.Job{"p1", "p2"},
-					},
-					{
-						ID:   "n2",
-						Jobs: []dax.Job{"p3"},
-					},
-				},
-			},
-			{
-				// Add job.
-				fn:    addJob,
-				input: "p4",
-				expDiff: []dax.WorkerDiff{
-					{
-						WorkerID:    "n0",
-						AddedJobs:   []dax.Job{"p4"},
-						RemovedJobs: []dax.Job{},
-					},
-				},
-				expState: []dax.WorkerInfo{
-					{
-						ID:   "n0",
-						Jobs: []dax.Job{"p4"},
-					},
-					{
-						ID:   "n1",
-						Jobs: []dax.Job{"p1", "p2"},
-					},
-					{
-						ID:   "n2",
-						Jobs: []dax.Job{"p3"},
-					},
-				},
-			},
-			{
-				// Add job.
-				fn:    addJob,
-				input: "p5",
-				expDiff: []dax.WorkerDiff{
-					{
-						WorkerID:    "n0",
-						AddedJobs:   []dax.Job{"p5"},
-						RemovedJobs: []dax.Job{},
-					},
-				},
-				expState: []dax.WorkerInfo{
-					{
-						ID:   "n0",
-						Jobs: []dax.Job{"p4", "p5"},
-					},
-					{
-						ID:   "n1",
-						Jobs: []dax.Job{"p1", "p2"},
-					},
-					{
-						ID:   "n2",
-						Jobs: []dax.Job{"p3"},
-					},
-				},
-			},
-			{
-				// Add job.
-				fn:    addJob,
-				input: "p0",
-				expDiff: []dax.WorkerDiff{
-					{
-						WorkerID:    "n2",
-						AddedJobs:   []dax.Job{"p0"},
-						RemovedJobs: []dax.Job{},
-					},
-				},
-				expState: []dax.WorkerInfo{
-					{
-						ID:   "n0",
-						Jobs: []dax.Job{"p4", "p5"},
-					},
-					{
-						ID:   "n1",
-						Jobs: []dax.Job{"p1", "p2"},
-					},
-					{
-						ID:   "n2",
-						Jobs: []dax.Job{"p0", "p3"},
-					},
-				},
-			},
-			{
-				// Add job.
-				fn:    addJob,
-				input: "p6",
-				expDiff: []dax.WorkerDiff{
-					{
-						WorkerID:    "n0",
-						AddedJobs:   []dax.Job{"p6"},
-						RemovedJobs: []dax.Job{},
-					},
-				},
-				expState: []dax.WorkerInfo{
-					{
-						ID:   "n0",
-						Jobs: []dax.Job{"p4", "p5", "p6"},
-					},
-					{
-						ID:   "n1",
-						Jobs: []dax.Job{"p1", "p2"},
-					},
-					{
-						ID:   "n2",
-						Jobs: []dax.Job{"p0", "p3"},
-					},
-				},
-			},
-			{
-				// Add job.
-				fn:    addJob,
-				input: "p7",
-				expDiff: []dax.WorkerDiff{
-					{
-						WorkerID:    "n1",
-						AddedJobs:   []dax.Job{"p7"},
-						RemovedJobs: []dax.Job{},
-					},
-				},
-				expState: []dax.WorkerInfo{
-					{
-						ID:   "n0",
-						Jobs: []dax.Job{"p4", "p5", "p6"},
-					},
-					{
-						ID:   "n1",
-						Jobs: []dax.Job{"p1", "p2", "p7"},
-					},
-					{
-						ID:   "n2",
-						Jobs: []dax.Job{"p0", "p3"},
-					},
-				},
-			},
-
-			//////////////////// Remove /////////////////////////
-
-			{
-				// Remove nonexistent worker.
-				fn:      bal.RemoveWorker,
-				input:   "nonexistent",
-				expDiff: []dax.WorkerDiff{},
-				expState: []dax.WorkerInfo{
-					{
-						ID:   "n0",
-						Jobs: []dax.Job{"p4", "p5", "p6"},
-					},
-					{
-						ID:   "n1",
-						Jobs: []dax.Job{"p1", "p2", "p7"},
-					},
-					{
-						ID:   "n2",
-						Jobs: []dax.Job{"p0", "p3"},
-					},
-				},
-			},
-			{
-				// Remove worker.
-				fn:    bal.RemoveWorker,
-				input: "n1",
-				expDiff: []dax.WorkerDiff{
-					{
-						WorkerID:    "n1",
-						AddedJobs:   []dax.Job{},
-						RemovedJobs: []dax.Job{"p1", "p2", "p7"},
-					},
-				},
-				expState: []dax.WorkerInfo{
-					{
-						ID:   "n0",
-						Jobs: []dax.Job{"p4", "p5", "p6"},
-					},
-					{
-						ID:   "n2",
-						Jobs: []dax.Job{"p0", "p3"},
-					},
-				},
-			},
-
-			{
-				// Remove job (from free list).
-				fn:      bal.RemoveJob,
-				input:   "p2",
-				expDiff: []dax.WorkerDiff{},
-				expState: []dax.WorkerInfo{
-					{
-						ID:   "n0",
-						Jobs: []dax.Job{"p4", "p5", "p6"},
-					},
-					{
-						ID:   "n2",
-						Jobs: []dax.Job{"p0", "p3"},
-					},
-				},
-			},
-
-			{
-				// Balance after remove.
-				balance: true,
-				expDiff: []dax.WorkerDiff{
-					{
-						WorkerID:    "n0",
-						AddedJobs:   []dax.Job{"p7"},
-						RemovedJobs: []dax.Job{},
-					},
-					{
-						WorkerID:    "n2",
-						AddedJobs:   []dax.Job{"p1"},
-						RemovedJobs: []dax.Job{},
-					},
-				},
-				expState: []dax.WorkerInfo{
-					{
-						ID:   "n0",
-						Jobs: []dax.Job{"p4", "p5", "p6", "p7"},
-					},
-					{
-						ID:   "n2",
-						Jobs: []dax.Job{"p0", "p1", "p3"},
-					},
-				},
-			},
-
-			{
-				// Remove job.
-				fn:    bal.RemoveJob,
-				input: "p1",
-				expDiff: []dax.WorkerDiff{
-					{
-						WorkerID:    "n2",
-						AddedJobs:   []dax.Job{},
-						RemovedJobs: []dax.Job{"p1"},
-					},
-				},
-				expState: []dax.WorkerInfo{
-					{
-						ID:   "n0",
-						Jobs: []dax.Job{"p4", "p5", "p6", "p7"},
-					},
-					{
-						ID:   "n2",
-						Jobs: []dax.Job{"p0", "p3"},
-					},
-				},
-			},
-		}
-		for i, test := range tests {
-			t.Run(fmt.Sprintf("test-%d", i), func(t *testing.T) {
-				var diff []dax.WorkerDiff
-				var err error
-				if test.balance {
-					diff, err = bal.Balance(ctx)
-				} else {
-					diff, err = test.fn(ctx, newStringWrapper(test.input))
-				}
-				assert.NoError(t, err)
-				assert.Equal(t, test.expDiff, diff)
-
-				cs, err := bal.CurrentState(ctx)
-				assert.NoError(t, err)
-				assert.Equal(t, test.expState, cs)
-			})
-		}
-	})
-
-	t.Run("WorkerState", func(t *testing.T) {
-		db, cleanup := newBoltBalancer(t)
-		defer cleanup()
-		bal := boltdb.NewBalancer("test", db, logger.NewStandardLogger(os.Stderr))
-
-		_, err := bal.AddWorker(ctx, newStringWrapper("n1"))
-		assert.NoError(t, err)
-		_, err = bal.AddJobs(ctx, newStringWrapper("p1"))
-		assert.NoError(t, err)
-
-		exp := dax.WorkerInfo{
-			ID:   "n1",
-			Jobs: []dax.Job{"p1"},
-		}
-		ws, err := bal.WorkerState(ctx, "n1")
-		assert.NoError(t, err)
-		assert.Equal(t, exp, ws)
-
-		// Worker doesn't exist.
-		exp = dax.WorkerInfo{
-			ID: "x1",
-		}
-		ws, err = bal.WorkerState(ctx, "x1")
-		assert.NoError(t, err)
-		assert.Equal(t, exp, ws)
-	})
-
-	t.Run("WorkersForJobs", func(t *testing.T) {
-		db, cleanup := newBoltBalancer(t)
-		defer cleanup()
-		bal := boltdb.NewBalancer("test", db, logger.NewStandardLogger(os.Stderr))
-
-		_, err := bal.AddWorker(ctx, newStringWrapper("n1"))
-		assert.NoError(t, err)
-		_, err = bal.AddWorker(ctx, newStringWrapper("n2"))
-		assert.NoError(t, err)
-		for i := 0; i < 12; i++ {
-			_, err = bal.AddJobs(ctx, newStringWrapper(fmt.Sprintf("p%d", i)))
-			assert.NoError(t, err)
-		}
-
-		exp := dax.WorkerInfo{
-			ID:   "n1",
-			Jobs: []dax.Job{"p0", "p10", "p2", "p4", "p6", "p8"},
-		}
-		ws, err := bal.WorkerState(ctx, "n1")
-		assert.NoError(t, err)
-		assert.Equal(t, exp, ws)
-
-		exp = dax.WorkerInfo{
-			ID:   "n2",
-			Jobs: []dax.Job{"p1", "p11", "p3", "p5", "p7", "p9"},
-		}
-		ws, err = bal.WorkerState(ctx, "n2")
-		assert.NoError(t, err)
-		assert.Equal(t, exp, ws)
-
-		tests := []struct {
-			jobs []dax.Job
-			exp  []dax.WorkerInfo
-		}{
-			{
-				jobs: []dax.Job{"p0"},
-				exp: []dax.WorkerInfo{
-					{ID: "n1", Jobs: []dax.Job{"p0"}},
-				},
-			},
-			{
-				jobs: []dax.Job{"p0", "p4"},
-				exp: []dax.WorkerInfo{
-					{ID: "n1", Jobs: []dax.Job{"p0", "p4"}},
-				},
-			},
-			{
-				jobs: []dax.Job{"p0", "p4", "p999"},
-				exp: []dax.WorkerInfo{
-					{ID: "n1", Jobs: []dax.Job{"p0", "p4"}},
-				},
-			},
-			{
-				jobs: []dax.Job{"p0", "p1"},
-				exp: []dax.WorkerInfo{
-					{ID: "n1", Jobs: []dax.Job{"p0"}},
-					{ID: "n2", Jobs: []dax.Job{"p1"}},
-				},
-			},
-			{
-				jobs: []dax.Job{"p5", "p0", "p1", "p8"},
-				exp: []dax.WorkerInfo{
-					{ID: "n1", Jobs: []dax.Job{"p0", "p8"}},
-					{ID: "n2", Jobs: []dax.Job{"p1", "p5"}},
-				},
-			},
-		}
-		for i, test := range tests {
-			t.Run(fmt.Sprintf("test-%d", i), func(t *testing.T) {
-				workers, err := bal.WorkersForJobs(ctx, test.jobs)
-				assert.NoError(t, err)
-				assert.Equal(t, test.exp, workers)
-			})
-		}
-
-		// Some tests for WorkersForJobPrefix
-		workers, err := bal.WorkersForJobPrefix(ctx, "p1")
-		assert.NoError(t, err)
-		assert.ElementsMatch(t, []dax.WorkerInfo{
-			{ID: "n1", Jobs: []dax.Job{"p10"}},
-			{ID: "n2", Jobs: []dax.Job{"p1", "p11"}},
-		}, workers)
-
-		workers, err = bal.WorkersForJobPrefix(ctx, "p2")
-		assert.NoError(t, err)
-		assert.ElementsMatch(t, []dax.WorkerInfo{
-			{ID: "n1", Jobs: []dax.Job{"p2"}},
-		}, workers)
-
-		workers, err = bal.WorkersForJobPrefix(ctx, "pp")
-		assert.NoError(t, err)
-		assert.ElementsMatch(t, []dax.WorkerInfo{}, workers)
-
-	})
-
-	t.Run("Balance", func(t *testing.T) {
-		db, cleanup := newBoltBalancer(t)
-		defer cleanup()
-		bal := boltdb.NewBalancer("test", db, logger.NewStandardLogger(os.Stderr))
-
-		// Add two workers with some jobs evenly spread across them.
-		_, err := bal.AddWorker(ctx, newStringWrapper("n1"))
-		assert.NoError(t, err)
-		_, err = bal.AddWorker(ctx, newStringWrapper("n2"))
-		assert.NoError(t, err)
-		for i := 0; i < 13; i++ {
-			_, err = bal.AddJobs(ctx, newStringWrapper(fmt.Sprintf("p%d", i)))
-			assert.NoError(t, err)
-		}
-
-		exp := dax.WorkerInfo{
-			ID:   "n1",
-			Jobs: []dax.Job{"p0", "p10", "p12", "p2", "p4", "p6", "p8"},
-		}
-		ws, err := bal.WorkerState(ctx, "n1")
-		assert.NoError(t, err)
-		assert.Equal(t, exp, ws)
-
-		exp = dax.WorkerInfo{
-			ID:   "n2",
-			Jobs: []dax.Job{"p1", "p11", "p3", "p5", "p7", "p9"},
-		}
-		ws, err = bal.WorkerState(ctx, "n2")
-		assert.NoError(t, err)
-		assert.Equal(t, exp, ws)
-
-		// Now, add a worker and confirm that it currently has no jobs assigned
-		// to it.
-		_, err = bal.AddWorker(ctx, newStringWrapper("n3"))
-		assert.NoError(t, err)
-		exp = dax.WorkerInfo{
-			ID:   "n3",
-			Jobs: []dax.Job{},
-		}
-		ws, err = bal.WorkerState(ctx, "n3")
-		assert.NoError(t, err)
-		assert.Equal(t, exp, ws)
-
-		// Finally, call Balance() and confirm that the appropriate jobs got
-		// reassigned.
-		_, err = bal.Balance(ctx)
-		assert.NoError(t, err)
-
-		exp = dax.WorkerInfo{
-			ID:   "n1",
-			Jobs: []dax.Job{"p0", "p10", "p12", "p2", "p4"},
-		}
-		ws, err = bal.WorkerState(ctx, "n1")
-		assert.NoError(t, err)
-		assert.Equal(t, exp, ws)
-
-		exp = dax.WorkerInfo{
-			ID:   "n2",
-			Jobs: []dax.Job{"p1", "p11", "p3", "p5"},
-		}
-		ws, err = bal.WorkerState(ctx, "n2")
-		assert.NoError(t, err)
-		assert.Equal(t, exp, ws)
-
-		exp = dax.WorkerInfo{
-			ID:   "n3",
-			Jobs: []dax.Job{"p6", "p7", "p8", "p9"},
-		}
-		ws, err = bal.WorkerState(ctx, "n3")
-		assert.NoError(t, err)
-		assert.Equal(t, exp, ws)
-	})
-}
-
-type stringWrapper struct {
-	s string
-}
-
-func newStringWrapper(s string) *stringWrapper {
-	return &stringWrapper{
-		s: s,
-	}
-}
-
-func (s *stringWrapper) String() string {
-	return s.s
-}

dax/mds/controller/naive/boltdb/balancer.go

@@ -1,606 +0,0 @@
-// Package boltdb contains the boltdb implementation of the Balancer interface.
-package boltdb
-
-import (
-	"bytes"
-	"context"
-	"encoding/json"
-	"fmt"
-	"strings"
-
-	"github.com/featurebasedb/featurebase/v3/dax"
-	"github.com/featurebasedb/featurebase/v3/dax/boltdb"
-	"github.com/featurebasedb/featurebase/v3/dax/mds/controller"
-	"github.com/featurebasedb/featurebase/v3/dax/mds/controller/naive"
-	"github.com/featurebasedb/featurebase/v3/errors"
-	"github.com/featurebasedb/featurebase/v3/logger"
-)
-
-var (
-	bucketNaiveBalancer = boltdb.Bucket("naiveBalancer")
-)
-
-// NaiveBalancerBuckets defines the buckets used by this package. It can be
-// called during setup to create the buckets ahead of time.
-var NaiveBalancerBuckets []boltdb.Bucket = []boltdb.Bucket{
-	bucketNaiveBalancer,
-}
-
-// NewBalancer returns a new instance of controller.Balancer.
-func NewBalancer(name string, db *boltdb.DB, logger logger.Logger) controller.Balancer {
-	fjs := newFreeJobService(db)
-	wjs := newWorkerJobService(db, logger)
-
-	return naive.New(name, fjs, wjs, logger)
-}
-
-// Ensure type implements interface.
-var _ naive.WorkerJobService = (*workerJobService)(nil)
-
-type workerJobService struct {
-	db     *boltdb.DB
-	logger logger.Logger
-}
-
-func newWorkerJobService(db *boltdb.DB, logger logger.Logger) *workerJobService {
-	return &workerJobService{
-		db:     db,
-		logger: logger,
-	}
-}
-
-func (w *workerJobService) WorkersJobs(ctx context.Context, balancerName string) ([]dax.WorkerInfo, error) {
-	tx, err := w.db.BeginTx(ctx, false)
-	if err != nil {
-		return nil, errors.Wrap(err, "getting tx")
-	}
-	defer tx.Rollback()
-
-	workerInfos, err := getWorkerInfos(ctx, tx, balancerName)
-	if err != nil {
-		return nil, errors.Wrapf(err, "getting worker infos: %s", balancerName)
-	}
-
-	return workerInfos, nil
-}
-
-func (w *workerJobService) WorkerCount(ctx context.Context, balancerName string) (int, error) {
-	tx, err := w.db.BeginTx(ctx, false)
-	if err != nil {
-		return 0, errors.Wrap(err, "getting tx")
-	}
-	defer tx.Rollback()
-
-	workers, err := w.getWorkers(ctx, tx, balancerName)
-	if err != nil {
-		return 0, errors.Wrapf(err, "getting workers: %s", balancerName)
-	}
-
-	return len(workers), nil
-}
-
-func (w *workerJobService) ListWorkers(ctx context.Context, balancerName string) (dax.Workers, error) {
-	tx, err := w.db.BeginTx(ctx, false)
-	if err != nil {
-		return nil, errors.Wrap(err, "beginning tx")
-	}
-	defer tx.Rollback()
-
-	workers, err := w.getWorkers(ctx, tx, balancerName)
-	if err != nil {
-		return nil, errors.Wrapf(err, "getting workers: %s", balancerName)
-	}
-
-	return workers, nil
-}
-
-func (w *workerJobService) getWorkers(ctx context.Context, tx *boltdb.Tx, balancerName string) (dax.Workers, error) {
-	c := tx.Bucket(bucketNaiveBalancer).Cursor()
-
-	// Deserialize rows into Worker objects.
-	workers := make(dax.Workers, 0)
-
-	prefix := []byte(fmt.Sprintf(prefixFmtWorkers, balancerName))
-	for k, v := c.Seek(prefix); k != nil && bytes.HasPrefix(k, prefix); k, v = c.Next() {
-		if v == nil {
-			w.logger.Printf("nil value for key: %s", k)
-			continue
-		}
-
-		worker, err := keyWorker(k)
-		if err != nil {
-			return nil, errors.Wrapf(err, "getting worker from key: %s", k)
-		}
-
-		workers = append(workers, worker)
-	}
-
-	return workers, nil
-}
-
-func getWorkerInfos(ctx context.Context, tx *boltdb.Tx, balancerName string) (dax.WorkerInfos, error) {
-	c := tx.Bucket(bucketNaiveBalancer).Cursor()
-
-	// Deserialize rows into WorkerInfo objects.
-	workerInfos := make(dax.WorkerInfos, 0)
-
-	prefix := []byte(fmt.Sprintf(prefixFmtWorkers, balancerName))
-	for k, v := c.Seek(prefix); k != nil && bytes.HasPrefix(k, prefix); k, v = c.Next() {
-		worker, err := keyWorker(k)
-		if err != nil {
-			return nil, errors.Wrapf(err, "getting worker from key: %s", k)
-		}
-
-		jobs := dax.NewSet[dax.Job]()
-		if v != nil {
-			jobs, err = decodeJobSet(v)
-			if err != nil {
-				return nil, errors.Wrap(err, "decoding job set")
-			}
-		}
-
-		workerInfo := dax.WorkerInfo{
-			ID:   worker,
-			Jobs: jobs.Sorted(),
-		}
-
-		workerInfos = append(workerInfos, workerInfo)
-	}
-
-	return workerInfos, nil
-}
-
-func (w *workerJobService) WorkerExists(ctx context.Context, balancerName string, worker dax.Worker) (bool, error) {
-	tx, err := w.db.BeginTx(ctx, false)
-	if err != nil {
-		return false, errors.Wrapf(err, "getting tx: %s", balancerName)
-	}
-	defer tx.Rollback()
-
-	bkt := tx.Bucket(bucketNaiveBalancer)
-	if bkt == nil {
-		return false, errors.Errorf(boltdb.ErrFmtBucketNotFound, bucketNaiveBalancer)
-	}
-
-	wrkr := bkt.Get(workerKey(balancerName, worker))
-
-	return wrkr != nil, nil
-}
-
-func (w *workerJobService) CreateWorker(ctx context.Context, balancerName string, worker dax.Worker) error {
-	tx, err := w.db.BeginTx(ctx, true)
-	if err != nil {
-		return errors.Wrap(err, "getting transaction")
-	}
-	defer tx.Rollback()
-
-	bkt := tx.Bucket(bucketNaiveBalancer)
-	if bkt == nil {
-		return errors.Errorf(boltdb.ErrFmtBucketNotFound, bucketNaiveBalancer)
-	}
-
-	// If this worker already exists, don't do anything.
-	wrkr := bkt.Get(workerKey(balancerName, worker))
-	if wrkr != nil {
-		return nil
-	}
-
-	val := []byte("[]")
-	if err := bkt.Put(workerKey(balancerName, worker), val); err != nil {
-		return errors.Wrap(err, "putting worker")
-	}
-
-	return tx.Commit()
-}
-
-func (w *workerJobService) DeleteWorker(ctx context.Context, balancerName string, worker dax.Worker) error {
-	tx, err := w.db.BeginTx(ctx, true)
-	if err != nil {
-		return errors.Wrap(err, "beginning tx")
-	}
-	defer tx.Rollback()
-
-	bkt := tx.Bucket(bucketNaiveBalancer)
-	if bkt == nil {
-		return errors.Errorf(boltdb.ErrFmtBucketNotFound, bucketNaiveBalancer)
-	}
-
-	if err := bkt.Delete(workerKey(balancerName, worker)); err != nil {
-		return errors.Wrapf(err, "deleting node key: %s", workerKey(balancerName, worker))
-	}
-
-	return tx.Commit()
-}
-
-func (w *workerJobService) CreateJobs(ctx context.Context, balancerName string, worker dax.Worker, jobs ...dax.Job) error {
-	tx, err := w.db.BeginTx(ctx, true)
-	if err != nil {
-		return errors.Wrap(err, "beginning tx")
-	}
-	defer tx.Rollback()
-
-	bkt := tx.Bucket(bucketNaiveBalancer)
-	if bkt == nil {
-		return errors.Errorf(boltdb.ErrFmtBucketNotFound, bucketNaiveBalancer)
-	}
-
-	jobset := dax.NewSet[dax.Job]()
-
-	// get worker
-	wrkr := bkt.Get(workerKey(balancerName, worker))
-	if wrkr != nil {
-		jobset, err = decodeJobSet(wrkr)
-		if err != nil {
-			return errors.Wrap(err, "decoding job set")
-		}
-	}
-
-	for _, job := range jobs {
-		jobset.Add(job)
-	}
-	val, err := encodeJobSet(jobset)
-	if err != nil {
-		return errors.Wrap(err, "encoding job set")
-	}
-
-	if err := bkt.Put(workerKey(balancerName, worker), val); err != nil {
-		return errors.Wrap(err, "putting worker")
-	}
-
-	return tx.Commit()
-}
-
-func (w *workerJobService) DeleteJob(ctx context.Context, balancerName string, worker dax.Worker, job dax.Job) error {
-	tx, err := w.db.BeginTx(ctx, true)
-	if err != nil {
-		return errors.Wrap(err, "beginning tx")
-	}
-	defer tx.Rollback()
-
-	bkt := tx.Bucket(bucketNaiveBalancer)
-	if bkt == nil {
-		return errors.Errorf(boltdb.ErrFmtBucketNotFound, bucketNaiveBalancer)
-	}
-
-	// get worker
-	wrkr := bkt.Get(workerKey(balancerName, worker))
-	if wrkr == nil {
-		return nil
-	}
-
-	jobset, err := decodeJobSet(wrkr)
-	if err != nil {
-		return errors.Wrap(err, "decoding job set")
-	}
-	if !jobset.Contains(job) {
-		return nil
-	}
-
-	jobset.Remove(job)
-	val, err := encodeJobSet(jobset)
-	if err != nil {
-		return errors.Wrap(err, "encoding job set")
-	}
-
-	if err := bkt.Put(workerKey(balancerName, worker), val); err != nil {
-		return errors.Wrap(err, "putting worker")
-	}
-
-	return tx.Commit()
-}
-
-func (w *workerJobService) DeleteJobs(ctx context.Context, balancerName, prefix string) (naive.InternalDiffs, error) {
-	tx, err := w.db.BeginTx(ctx, true)
-	if err != nil {
-		return nil, errors.Wrap(err, "beginning tx")
-	}
-	defer tx.Rollback()
-
-	bkt := tx.Bucket(bucketNaiveBalancer)
-	if bkt == nil {
-		return nil, errors.Errorf(boltdb.ErrFmtBucketNotFound, bucketNaiveBalancer)
-	}
-
-	workers, err := w.getWorkers(ctx, tx, balancerName)
-	if err != nil {
-		return nil, errors.Wrap(err, "getting workers")
-	}
-
-	idiffs := naive.NewInternalDiffs()
-	for _, worker := range workers {
-		// get worker
-		wrkr := bkt.Get(workerKey(balancerName, worker))
-		if wrkr == nil {
-			panic("didn't find worker that should... definitely exist")
-		}
-		jobset, err := decodeJobSet(wrkr)
-		if err != nil {
-			return nil, errors.Wrap(err, "decoding job set")
-		}
-
-		jobs := jobset.RemovePrefix(prefix)
-		for _, job := range jobs {
-			idiffs.Removed(worker, job)
-		}
-		val, err := encodeJobSet(jobset)
-		if err != nil {
-			return nil, errors.Wrap(err, "encoding job set")
-		}
-
-		if err := bkt.Put(workerKey(balancerName, worker), val); err != nil {
-			return nil, errors.Wrap(err, "putting worker")
-		}
-
-	}
-
-	return idiffs, tx.Commit()
-}
-
-func (w *workerJobService) ListJobs(ctx context.Context, balancerName string, worker dax.Worker) (dax.Jobs, error) {
-	tx, err := w.db.BeginTx(ctx, false)
-	if err != nil {
-		return nil, errors.Wrap(err, "beginning tx")
-	}
-	defer tx.Rollback()
-
-	bkt := tx.Bucket(bucketNaiveBalancer)
-	if bkt == nil {
-		return nil, errors.Errorf(boltdb.ErrFmtBucketNotFound, bucketNaiveBalancer)
-	}
-
-	jobset := dax.NewSet[dax.Job]()
-
-	// get worker
-	wrkr := bkt.Get(workerKey(balancerName, worker))
-	if wrkr != nil {
-		jobset, err = decodeJobSet(wrkr)
-		if err != nil {
-			return nil, errors.Wrap(err, "decoding job set")
-		}
-	}
-
-	return jobset.Sorted(), nil
-}
-
-func (w *workerJobService) JobCounts(ctx context.Context, balancerName string, workers ...dax.Worker) (map[dax.Worker]int, error) {
-	tx, err := w.db.BeginTx(ctx, false)
-	if err != nil {
-		return nil, errors.Wrapf(err, "getting tx: %s", balancerName)
-	}
-	defer tx.Rollback()
-
-	bkt := tx.Bucket(bucketNaiveBalancer)
-	if bkt == nil {
-		return nil, errors.Errorf(boltdb.ErrFmtBucketNotFound, bucketNaiveBalancer)
-	}
-
-	m := make(map[dax.Worker]int)
-
-	for _, worker := range workers {
-		jobset := dax.NewSet[dax.Job]()
-
-		// get worker
-		wrkr := bkt.Get(workerKey(balancerName, worker))
-		if wrkr != nil {
-			jobset, err = decodeJobSet(wrkr)
-			if err != nil {
-				return nil, errors.Wrap(err, "decoding job set")
-			}
-		}
-
-		m[worker] = len(jobset)
-	}
-
-	return m, nil
-}
-
-// encodeJobSet encode the jobSet into a JSON array of strings.
-func encodeJobSet(jobSet dax.Set[dax.Job]) ([]byte, error) {
-	arr := jobSet.Sorted()
-	b, err := json.Marshal(arr)
-	if err != nil {
-		return nil, errors.Wrap(err, "marshalling json")
-	}
-	return b, nil
-}
-
-// decodeJobSet decode the string (a JSON array of strings) into jobSet.
-func decodeJobSet(v []byte) (dax.Set[dax.Job], error) {
-	var arr []string
-	err := json.Unmarshal(v, &arr)
-	if err != nil {
-		return nil, errors.Wrap(err, "unmarshalling json")
-	}
-
-	js := dax.NewSet[dax.Job]()
-	for _, s := range arr {
-		js.Add(dax.Job(s))
-	}
-
-	return js, nil
-}
-
-// Ensure type implements interface.
-var _ naive.FreeJobService = (*freeJobService)(nil)
-
-type freeJobService struct {
-	db *boltdb.DB
-}
-
-func newFreeJobService(db *boltdb.DB) *freeJobService {
-	return &freeJobService{
-		db: db,
-	}
-}
-
-func (f *freeJobService) CreateFreeJobs(ctx context.Context, balancerName string, jobs ...dax.Job) error {
-	return f.MergeFreeJobs(ctx, balancerName, jobs)
-}
-
-func (f *freeJobService) DeleteFreeJob(ctx context.Context, balancerName string, job dax.Job) error {
-	tx, err := f.db.BeginTx(ctx, true)
-	if err != nil {
-		return errors.Wrap(err, "beginning tx")
-	}
-	defer tx.Rollback()
-
-	bkt := tx.Bucket(bucketNaiveBalancer)
-	if bkt == nil {
-		return errors.Errorf(boltdb.ErrFmtBucketNotFound, bucketNaiveBalancer)
-	}
-
-	// get free jobs
-	fjs := bkt.Get(freeJobKey(balancerName))
-	if fjs == nil {
-		return nil
-	}
-
-	jobset, err := decodeJobSet(fjs)
-	if err != nil {
-		return errors.Wrap(err, "decoding job set")
-	}
-	if !jobset.Contains(job) {
-		return nil
-	}
-
-	jobset.Remove(job)
-	val, err := encodeJobSet(jobset)
-	if err != nil {
-		return errors.Wrap(err, "encoding job set")
-	}
-
-	if err := bkt.Put(freeJobKey(balancerName), val); err != nil {
-		return errors.Wrap(err, "putting free job")
-	}
-
-	return tx.Commit()
-}
-
-func (f *freeJobService) DeleteFreeJobs(ctx context.Context, balancerName, prefix string) error {
-	tx, err := f.db.BeginTx(ctx, true)
-	if err != nil {
-		return errors.Wrap(err, "beginning tx")
-	}
-	defer tx.Rollback()
-
-	bkt := tx.Bucket(bucketNaiveBalancer)
-	if bkt == nil {
-		return errors.Errorf(boltdb.ErrFmtBucketNotFound, bucketNaiveBalancer)
-	}
-
-	// get free jobs
-	fjs := bkt.Get(freeJobKey(balancerName))
-	if fjs == nil {
-		return nil
-	}
-
-	jobset, err := decodeJobSet(fjs)
-	if err != nil {
-		return errors.Wrap(err, "decoding job set")
-	}
-
-	jobset.RemovePrefix(prefix)
-	val, err := encodeJobSet(jobset)
-	if err != nil {
-		return errors.Wrap(err, "encoding job set")
-	}
-
-	if err := bkt.Put(freeJobKey(balancerName), val); err != nil {
-		return errors.Wrap(err, "putting free job")
-	}
-
-	return tx.Commit()
-}
-
-func (f *freeJobService) ListFreeJobs(ctx context.Context, balancerName string) (dax.Jobs, error) {
-	tx, err := f.db.BeginTx(ctx, false)
-	if err != nil {
-		return nil, errors.Wrap(err, "beginning tx")
-	}
-	defer tx.Rollback()
-
-	bkt := tx.Bucket(bucketNaiveBalancer)
-	if bkt == nil {
-		return nil, errors.Errorf(boltdb.ErrFmtBucketNotFound, bucketNaiveBalancer)
-	}
-
-	jobset := dax.NewSet[dax.Job]()
-
-	// get free jobs
-	fjs := bkt.Get(freeJobKey(balancerName))
-	if fjs != nil {
-		jobset, err = decodeJobSet(fjs)
-		if err != nil {
-			return nil, errors.Wrap(err, "decoding job set")
-		}
-	}
-
-	return jobset.Sorted(), nil
-}
-
-func (f *freeJobService) MergeFreeJobs(ctx context.Context, balancerName string, jobs dax.Jobs) error {
-	tx, err := f.db.BeginTx(ctx, true)
-	if err != nil {
-		return err
-	}
-	defer tx.Rollback()
-
-	bkt := tx.Bucket(bucketNaiveBalancer)
-	if bkt == nil {
-		return errors.Errorf(boltdb.ErrFmtBucketNotFound, bucketNaiveBalancer)
-	}
-
-	jobset := dax.NewSet[dax.Job]()
-
-	// get free jobs
-	fjs := bkt.Get(freeJobKey(balancerName))
-	if fjs != nil {
-		jobset, err = decodeJobSet(fjs)
-		if err != nil {
-			return errors.Wrap(err, "decoding job set")
-		}
-	}
-
-	for _, j := range jobs {
-		jobset.Add(j)
-	}
-	val, err := encodeJobSet(jobset)
-	if err != nil {
-		return errors.Wrap(err, "encoding job set")
-	}
-
-	if err := bkt.Put(freeJobKey(balancerName), val); err != nil {
-		return errors.Wrap(err, "putting free job")
-	}
-
-	return tx.Commit()
-}
-
-//////////////////////////////////////////////////////
-
-const (
-	prefixFmtWorkers  = "workers/%s/" // %s - balancerName
-	prefixFmtFreeJobs = "freejobs/%s" // %s - balancerName
-)
-
-// workerKey returns a key based on worker.
-func workerKey(bal string, worker dax.Worker) []byte {
-	key := fmt.Sprintf(prefixFmtWorkers+"%s", bal, worker)
-	return []byte(key)
-}
-
-// keyWorker gets the worker out of the key.
-func keyWorker(key []byte) (dax.Worker, error) {
-	parts := strings.SplitN(string(key), "/", 3)
-	if len(parts) != 3 {
-		return "", errors.New(errors.ErrUncoded, "worker key format expected: `workers/balancer/worker`")
-	}
-
-	return dax.Worker(parts[2]), nil
-}
-
-// freeJobKey returns a key for all freeJobs.
-func freeJobKey(bal string) []byte {
-	key := fmt.Sprintf(prefixFmtFreeJobs, bal)
-	return []byte(key)
-}

dax/mds/controller/naive/boltdb/balancer_test.go

@@ -1,721 +0,0 @@
-package boltdb_test
-
-import (
-	"context"
-	"fmt"
-	"testing"
-
-	"github.com/featurebasedb/featurebase/v3/dax"
-	"github.com/featurebasedb/featurebase/v3/dax/mds/controller/naive/boltdb"
-	testbolt "github.com/featurebasedb/featurebase/v3/dax/test/boltdb"
-	"github.com/featurebasedb/featurebase/v3/logger"
-	"github.com/stretchr/testify/assert"
-)
-
-func TestBalancer(t *testing.T) {
-	db := testbolt.MustOpenDB(t)
-	defer testbolt.MustCloseDB(t, db)
-
-	t.Cleanup(func() {
-		testbolt.CleanupDB(t, db.Path())
-	})
-
-	ctx := context.Background()
-
-	// Initialize the buckets.
-	assert.NoError(t, db.InitializeBuckets(boltdb.NaiveBalancerBuckets...))
-
-	t.Run("SingleWorker", func(t *testing.T) {
-		bal := boltdb.NewBalancer("test-single-worker", db, logger.NopLogger)
-		// addJob is a wrapper around bal.AddJobs() which we added when the
-		// function signature of bal.AddJobs changed to take multiple jobs (and
-		// it therefore no longer satisfied the fn type in this test).
-		addJob := func(ctx context.Context, job fmt.Stringer) ([]dax.WorkerDiff, error) {
-			return bal.AddJobs(ctx, job)
-		}
-		tests := []struct {
-			fn       func(context.Context, fmt.Stringer) ([]dax.WorkerDiff, error)
-			input    string
-			expDiff  []dax.WorkerDiff
-			expState []dax.WorkerInfo
-		}{
-			{
-				// Add job.
-				fn:       addJob,
-				input:    "p2",
-				expDiff:  []dax.WorkerDiff{},
-				expState: []dax.WorkerInfo{},
-			},
-			{
-				// Add worker.
-				fn:    bal.AddWorker,
-				input: "n1",
-				expDiff: []dax.WorkerDiff{
-					{
-						WorkerID:    "n1",
-						AddedJobs:   []dax.Job{"p2"},
-						RemovedJobs: []dax.Job{},
-					},
-				},
-				expState: []dax.WorkerInfo{
-					{
-						ID:   "n1",
-						Jobs: []dax.Job{"p2"},
-					},
-				},
-			},
-			{
-				// Add another job out of order.
-				fn:    addJob,
-				input: "p1",
-				expDiff: []dax.WorkerDiff{
-					{
-						WorkerID:    "n1",
-						AddedJobs:   []dax.Job{"p1"},
-						RemovedJobs: []dax.Job{},
-					},
-				},
-				expState: []dax.WorkerInfo{
-					{
-						ID:   "n1",
-						Jobs: []dax.Job{"p1", "p2"},
-					},
-				},
-			},
-			{
-				// Add another job.
-				fn:    addJob,
-				input: "p3",
-				expDiff: []dax.WorkerDiff{
-					{
-						WorkerID:    "n1",
-						AddedJobs:   []dax.Job{"p3"},
-						RemovedJobs: []dax.Job{},
-					},
-				},
-				expState: []dax.WorkerInfo{
-					{
-						ID:   "n1",
-						Jobs: []dax.Job{"p1", "p2", "p3"},
-					},
-				},
-			},
-			{
-				// Add a duplicate job.
-				fn:      addJob,
-				input:   "p2",
-				expDiff: []dax.WorkerDiff{},
-				expState: []dax.WorkerInfo{
-					{
-						ID:   "n1",
-						Jobs: []dax.Job{"p1", "p2", "p3"},
-					},
-				},
-			},
-		}
-		for i, test := range tests {
-			t.Run(fmt.Sprintf("test-%d", i), func(t *testing.T) {
-				diff, err := test.fn(ctx, newStringWrapper(test.input))
-				assert.NoError(t, err)
-				assert.Equal(t, test.expDiff, diff)
-
-				cs, err := bal.CurrentState(ctx)
-				assert.NoError(t, err)
-				assert.Equal(t, test.expState, cs)
-			})
-		}
-	})
-
-	t.Run("MultipleWorkers", func(t *testing.T) {
-		bal := boltdb.NewBalancer("test-multiple-workers", db, logger.NopLogger)
-		// addJob is a wrapper around bal.AddJobs() which we added when the
-		// function signature of bal.AddJobs changed to take multiple jobs (and
-		// it therefore no longer satisfied the fn type in this test).
-		addJob := func(ctx context.Context, job fmt.Stringer) ([]dax.WorkerDiff, error) {
-			return bal.AddJobs(ctx, job)
-		}
-		tests := []struct {
-			fn       func(context.Context, fmt.Stringer) ([]dax.WorkerDiff, error)
-			input    string
-			balance  bool
-			expDiff  []dax.WorkerDiff
-			expState []dax.WorkerInfo
-		}{
-			{
-				// Balance when empty.
-				balance:  true,
-				expDiff:  []dax.WorkerDiff{},
-				expState: []dax.WorkerInfo{},
-			},
-			{
-				// Add worker.
-				fn:      bal.AddWorker,
-				input:   "n2",
-				expDiff: []dax.WorkerDiff{},
-				expState: []dax.WorkerInfo{
-					{
-						ID:   "n2",
-						Jobs: []dax.Job{},
-					},
-				},
-			},
-			{
-				// Add worker again.
-				fn:      bal.AddWorker,
-				input:   "n2",
-				expDiff: []dax.WorkerDiff{},
-				expState: []dax.WorkerInfo{
-					{
-						ID:   "n2",
-						Jobs: []dax.Job{},
-					},
-				},
-			},
-			{
-				// Add a second worker.
-				fn:      bal.AddWorker,
-				input:   "n1",
-				expDiff: []dax.WorkerDiff{},
-				expState: []dax.WorkerInfo{
-					{
-						ID:   "n1",
-						Jobs: []dax.Job{},
-					},
-					{
-						ID:   "n2",
-						Jobs: []dax.Job{},
-					},
-				},
-			},
-			{
-				// Add job.
-				fn:    addJob,
-				input: "p2",
-				expDiff: []dax.WorkerDiff{
-					{
-						WorkerID:    "n1",
-						AddedJobs:   []dax.Job{"p2"},
-						RemovedJobs: []dax.Job{},
-					},
-				},
-				expState: []dax.WorkerInfo{
-					{
-						ID:   "n1",
-						Jobs: []dax.Job{"p2"},
-					},
-					{
-						ID:   "n2",
-						Jobs: []dax.Job{},
-					},
-				},
-			},
-			{
-				// Add job.
-				fn:    addJob,
-				input: "p3",
-				expDiff: []dax.WorkerDiff{
-					{
-						WorkerID:    "n2",
-						AddedJobs:   []dax.Job{"p3"},
-						RemovedJobs: []dax.Job{},
-					},
-				},
-				expState: []dax.WorkerInfo{
-					{
-						ID:   "n1",
-						Jobs: []dax.Job{"p2"},
-					},
-					{
-						ID:   "n2",
-						Jobs: []dax.Job{"p3"},
-					},
-				},
-			},
-			{
-				// Add job.
-				fn:    addJob,
-				input: "p1",
-				expDiff: []dax.WorkerDiff{
-					{
-						WorkerID:    "n1",
-						AddedJobs:   []dax.Job{"p1"},
-						RemovedJobs: []dax.Job{},
-					},
-				},
-				expState: []dax.WorkerInfo{
-					{
-						ID:   "n1",
-						Jobs: []dax.Job{"p1", "p2"},
-					},
-					{
-						ID:   "n2",
-						Jobs: []dax.Job{"p3"},
-					},
-				},
-			},
-			{
-				// Add a third worker.
-				fn:      bal.AddWorker,
-				input:   "n0",
-				expDiff: []dax.WorkerDiff{},
-				expState: []dax.WorkerInfo{
-					{
-						ID:   "n0",
-						Jobs: []dax.Job{},
-					},
-					{
-						ID:   "n1",
-						Jobs: []dax.Job{"p1", "p2"},
-					},
-					{
-						ID:   "n2",
-						Jobs: []dax.Job{"p3"},
-					},
-				},
-			},
-			{
-				// Add job.
-				fn:    addJob,
-				input: "p4",
-				expDiff: []dax.WorkerDiff{
-					{
-						WorkerID:    "n0",
-						AddedJobs:   []dax.Job{"p4"},
-						RemovedJobs: []dax.Job{},
-					},
-				},
-				expState: []dax.WorkerInfo{
-					{
-						ID:   "n0",
-						Jobs: []dax.Job{"p4"},
-					},
-					{
-						ID:   "n1",
-						Jobs: []dax.Job{"p1", "p2"},
-					},
-					{
-						ID:   "n2",
-						Jobs: []dax.Job{"p3"},
-					},
-				},
-			},
-			{
-				// Add job.
-				fn:    addJob,
-				input: "p5",
-				expDiff: []dax.WorkerDiff{
-					{
-						WorkerID:    "n0",
-						AddedJobs:   []dax.Job{"p5"},
-						RemovedJobs: []dax.Job{},
-					},
-				},
-				expState: []dax.WorkerInfo{
-					{
-						ID:   "n0",
-						Jobs: []dax.Job{"p4", "p5"},
-					},
-					{
-						ID:   "n1",
-						Jobs: []dax.Job{"p1", "p2"},
-					},
-					{
-						ID:   "n2",
-						Jobs: []dax.Job{"p3"},
-					},
-				},
-			},
-			{
-				// Add job.
-				fn:    addJob,
-				input: "p0",
-				expDiff: []dax.WorkerDiff{
-					{
-						WorkerID:    "n2",
-						AddedJobs:   []dax.Job{"p0"},
-						RemovedJobs: []dax.Job{},
-					},
-				},
-				expState: []dax.WorkerInfo{
-					{
-						ID:   "n0",
-						Jobs: []dax.Job{"p4", "p5"},
-					},
-					{
-						ID:   "n1",
-						Jobs: []dax.Job{"p1", "p2"},
-					},
-					{
-						ID:   "n2",
-						Jobs: []dax.Job{"p0", "p3"},
-					},
-				},
-			},
-			{
-				// Add job.
-				fn:    addJob,
-				input: "p6",
-				expDiff: []dax.WorkerDiff{
-					{
-						WorkerID:    "n0",
-						AddedJobs:   []dax.Job{"p6"},
-						RemovedJobs: []dax.Job{},
-					},
-				},
-				expState: []dax.WorkerInfo{
-					{
-						ID:   "n0",
-						Jobs: []dax.Job{"p4", "p5", "p6"},
-					},
-					{
-						ID:   "n1",
-						Jobs: []dax.Job{"p1", "p2"},
-					},
-					{
-						ID:   "n2",
-						Jobs: []dax.Job{"p0", "p3"},
-					},
-				},
-			},
-			{
-				// Add job.
-				fn:    addJob,
-				input: "p7",
-				expDiff: []dax.WorkerDiff{
-					{
-						WorkerID:    "n1",
-						AddedJobs:   []dax.Job{"p7"},
-						RemovedJobs: []dax.Job{},
-					},
-				},
-				expState: []dax.WorkerInfo{
-					{
-						ID:   "n0",
-						Jobs: []dax.Job{"p4", "p5", "p6"},
-					},
-					{
-						ID:   "n1",
-						Jobs: []dax.Job{"p1", "p2", "p7"},
-					},
-					{
-						ID:   "n2",
-						Jobs: []dax.Job{"p0", "p3"},
-					},
-				},
-			},
-
-			//////////////////// Remove /////////////////////////
-
-			{
-				// Remove nonexistent worker.
-				fn:      bal.RemoveWorker,
-				input:   "nonexistent",
-				expDiff: []dax.WorkerDiff{},
-				expState: []dax.WorkerInfo{
-					{
-						ID:   "n0",
-						Jobs: []dax.Job{"p4", "p5", "p6"},
-					},
-					{
-						ID:   "n1",
-						Jobs: []dax.Job{"p1", "p2", "p7"},
-					},
-					{
-						ID:   "n2",
-						Jobs: []dax.Job{"p0", "p3"},
-					},
-				},
-			},
-			{
-				// Remove worker.
-				fn:    bal.RemoveWorker,
-				input: "n1",
-				expDiff: []dax.WorkerDiff{
-					{
-						WorkerID:    "n1",
-						AddedJobs:   []dax.Job{},
-						RemovedJobs: []dax.Job{"p1", "p2", "p7"},
-					},
-				},
-				expState: []dax.WorkerInfo{
-					{
-						ID:   "n0",
-						Jobs: []dax.Job{"p4", "p5", "p6"},
-					},
-					{
-						ID:   "n2",
-						Jobs: []dax.Job{"p0", "p3"},
-					},
-				},
-			},
-
-			{
-				// Remove job (from free list).
-				fn:      bal.RemoveJob,
-				input:   "p2",
-				expDiff: []dax.WorkerDiff{},
-				expState: []dax.WorkerInfo{
-					{
-						ID:   "n0",
-						Jobs: []dax.Job{"p4", "p5", "p6"},
-					},
-					{
-						ID:   "n2",
-						Jobs: []dax.Job{"p0", "p3"},
-					},
-				},
-			},
-
-			{
-				// Balance after remove.
-				balance: true,
-				expDiff: []dax.WorkerDiff{
-					{
-						WorkerID:    "n0",
-						AddedJobs:   []dax.Job{"p7"},
-						RemovedJobs: []dax.Job{},
-					},
-					{
-						WorkerID:    "n2",
-						AddedJobs:   []dax.Job{"p1"},
-						RemovedJobs: []dax.Job{},
-					},
-				},
-				expState: []dax.WorkerInfo{
-					{
-						ID:   "n0",
-						Jobs: []dax.Job{"p4", "p5", "p6", "p7"},
-					},
-					{
-						ID:   "n2",
-						Jobs: []dax.Job{"p0", "p1", "p3"},
-					},
-				},
-			},
-
-			{
-				// Remove job.
-				fn:    bal.RemoveJob,
-				input: "p1",
-				expDiff: []dax.WorkerDiff{
-					{
-						WorkerID:    "n2",
-						AddedJobs:   []dax.Job{},
-						RemovedJobs: []dax.Job{"p1"},
-					},
-				},
-				expState: []dax.WorkerInfo{
-					{
-						ID:   "n0",
-						Jobs: []dax.Job{"p4", "p5", "p6", "p7"},
-					},
-					{
-						ID:   "n2",
-						Jobs: []dax.Job{"p0", "p3"},
-					},
-				},
-			},
-		}
-		for i, test := range tests {
-			t.Run(fmt.Sprintf("test-%d", i), func(t *testing.T) {
-				var diff []dax.WorkerDiff
-				var err error
-				if test.balance {
-					diff, err = bal.Balance(ctx)
-				} else {
-					diff, err = test.fn(ctx, newStringWrapper(test.input))
-				}
-				assert.NoError(t, err)
-				assert.Equal(t, test.expDiff, diff)
-
-				cs, err := bal.CurrentState(ctx)
-				assert.NoError(t, err)
-				assert.Equal(t, test.expState, cs)
-			})
-		}
-	})
-
-	t.Run("WorkerState", func(t *testing.T) {
-		bal := boltdb.NewBalancer("test-worker-state", db, logger.NopLogger)
-
-		_, err := bal.AddWorker(ctx, newStringWrapper("n1"))
-		assert.NoError(t, err)
-		_, err = bal.AddJobs(ctx, newStringWrapper("p1"))
-		assert.NoError(t, err)
-
-		exp := dax.WorkerInfo{
-			ID:   "n1",
-			Jobs: []dax.Job{"p1"},
-		}
-		ws, err := bal.WorkerState(ctx, "n1")
-		assert.NoError(t, err)
-		assert.Equal(t, exp, ws)
-
-		// Worker doesn't exist.
-		exp = dax.WorkerInfo{
-			ID: "x1",
-		}
-		ws, err = bal.WorkerState(ctx, "x1")
-		assert.NoError(t, err)
-		assert.Equal(t, exp, ws)
-	})
-
-	t.Run("WorkersForJobs", func(t *testing.T) {
-		bal := boltdb.NewBalancer("test-workers-for-jobs", db, logger.NopLogger)
-
-		_, err := bal.AddWorker(ctx, newStringWrapper("n1"))
-		assert.NoError(t, err)
-		_, err = bal.AddWorker(ctx, newStringWrapper("n2"))
-		assert.NoError(t, err)
-		for i := 0; i < 12; i++ {
-			_, err = bal.AddJobs(ctx, newStringWrapper(fmt.Sprintf("p%d", i)))
-			assert.NoError(t, err)
-		}
-
-		exp := dax.WorkerInfo{
-			ID:   "n1",
-			Jobs: []dax.Job{"p0", "p10", "p2", "p4", "p6", "p8"},
-		}
-		ws, err := bal.WorkerState(ctx, "n1")
-		assert.NoError(t, err)
-		assert.Equal(t, exp, ws)
-
-		exp = dax.WorkerInfo{
-			ID:   "n2",
-			Jobs: []dax.Job{"p1", "p11", "p3", "p5", "p7", "p9"},
-		}
-		ws, err = bal.WorkerState(ctx, "n2")
-		assert.NoError(t, err)
-		assert.Equal(t, exp, ws)
-
-		tests := []struct {
-			jobs []dax.Job
-			exp  []dax.WorkerInfo
-		}{
-			{
-				jobs: []dax.Job{"p0"},
-				exp: []dax.WorkerInfo{
-					{ID: "n1", Jobs: []dax.Job{"p0"}},
-				},
-			},
-			{
-				jobs: []dax.Job{"p0", "p4"},
-				exp: []dax.WorkerInfo{
-					{ID: "n1", Jobs: []dax.Job{"p0", "p4"}},
-				},
-			},
-			{
-				jobs: []dax.Job{"p0", "p4", "p999"},
-				exp: []dax.WorkerInfo{
-					{ID: "n1", Jobs: []dax.Job{"p0", "p4"}},
-				},
-			},
-			{
-				jobs: []dax.Job{"p0", "p1"},
-				exp: []dax.WorkerInfo{
-					{ID: "n1", Jobs: []dax.Job{"p0"}},
-					{ID: "n2", Jobs: []dax.Job{"p1"}},
-				},
-			},
-			{
-				jobs: []dax.Job{"p5", "p0", "p1", "p8"},
-				exp: []dax.WorkerInfo{
-					{ID: "n1", Jobs: []dax.Job{"p0", "p8"}},
-					{ID: "n2", Jobs: []dax.Job{"p1", "p5"}},
-				},
-			},
-		}
-		for i, test := range tests {
-			t.Run(fmt.Sprintf("test-%d", i), func(t *testing.T) {
-				workers, err := bal.WorkersForJobs(ctx, test.jobs)
-				assert.NoError(t, err)
-				assert.Equal(t, test.exp, workers)
-			})
-		}
-	})
-
-	t.Run("Balance", func(t *testing.T) {
-		bal := boltdb.NewBalancer("test-balance", db, logger.NopLogger)
-
-		// Add two workers with some jobs evenly spread across them.
-		_, err := bal.AddWorker(ctx, newStringWrapper("n1"))
-		assert.NoError(t, err)
-		_, err = bal.AddWorker(ctx, newStringWrapper("n2"))
-		assert.NoError(t, err)
-		for i := 0; i < 13; i++ {
-			_, err = bal.AddJobs(ctx, newStringWrapper(fmt.Sprintf("p%d", i)))
-			assert.NoError(t, err)
-		}
-
-		exp := dax.WorkerInfo{
-			ID:   "n1",
-			Jobs: []dax.Job{"p0", "p10", "p12", "p2", "p4", "p6", "p8"},
-		}
-		ws, err := bal.WorkerState(ctx, "n1")
-		assert.NoError(t, err)
-		assert.Equal(t, exp, ws)
-
-		exp = dax.WorkerInfo{
-			ID:   "n2",
-			Jobs: []dax.Job{"p1", "p11", "p3", "p5", "p7", "p9"},
-		}
-		ws, err = bal.WorkerState(ctx, "n2")
-		assert.NoError(t, err)
-		assert.Equal(t, exp, ws)
-
-		// Now, add a worker and confirm that it currently has no jobs assigned
-		// to it.
-		_, err = bal.AddWorker(ctx, newStringWrapper("n3"))
-		assert.NoError(t, err)
-		exp = dax.WorkerInfo{
-			ID:   "n3",
-			Jobs: []dax.Job{},
-		}
-		ws, err = bal.WorkerState(ctx, "n3")
-		assert.NoError(t, err)
-		assert.Equal(t, exp, ws)
-
-		// Finally, call Balance() and confirm that the appropriate jobs got
-		// reassigned.
-		_, err = bal.Balance(ctx)
-		assert.NoError(t, err)
-
-		exp = dax.WorkerInfo{
-			ID:   "n1",
-			Jobs: []dax.Job{"p0", "p10", "p12", "p2", "p4"},
-		}
-		ws, err = bal.WorkerState(ctx, "n1")
-		assert.NoError(t, err)
-		assert.Equal(t, exp, ws)
-
-		exp = dax.WorkerInfo{
-			ID:   "n2",
-			Jobs: []dax.Job{"p1", "p11", "p3", "p5"},
-		}
-		ws, err = bal.WorkerState(ctx, "n2")
-		assert.NoError(t, err)
-		assert.Equal(t, exp, ws)
-
-		exp = dax.WorkerInfo{
-			ID:   "n3",
-			Jobs: []dax.Job{"p6", "p7", "p8", "p9"},
-		}
-		ws, err = bal.WorkerState(ctx, "n3")
-		assert.NoError(t, err)
-		assert.Equal(t, exp, ws)
-	})
-}
-
-type stringWrapper struct {
-	s string
-}
-
-func newStringWrapper(s string) *stringWrapper {
-	return &stringWrapper{
-		s: s,
-	}
-}
-
-func (s *stringWrapper) String() string {
-	return s.s
-}

dax/mds/controller/sets.go

@@ -85,11 +85,11 @@ func (s TableSet) SortedSlice() dax.TableKeys {
 	return ps
 }
 
-func (s TableSet) QualifiedSortedSlice() map[dax.TableQualifier]dax.TableIDs {
-	m := make(map[dax.TableQualifier]dax.TableIDs)
+func (s TableSet) QualifiedSortedSlice() map[dax.QualifiedDatabaseID]dax.TableIDs {
+	m := make(map[dax.QualifiedDatabaseID]dax.TableIDs)
 	for p := range s {
 		qtid := p.QualifiedTableID()
-		m[qtid.TableQualifier] = append(m[qtid.TableQualifier], qtid.ID)
+		m[qtid.QualifiedDatabaseID] = append(m[qtid.QualifiedDatabaseID], qtid.ID)
 	}
 
 	// Sort the slices in the map.

dax/mds/controller/snapping_turtle.go

@@ -30,7 +30,29 @@ func (c *Controller) snappingTurtleRoutine(period time.Duration, control chan st
 func (c *Controller) snapAll() {
 	c.logger.Debugf("TURTLE: snapAll")
 	ctx := context.Background()
-	computeNodes, err := c.ComputeBalancer.CurrentState(ctx)
+
+	tx, err := c.boltDB.BeginTx(ctx, false)
+	if err != nil {
+		c.logger.Printf("Error getting transaction for snapping turtle: %v", err)
+		return
+	}
+	defer tx.Rollback()
+
+	qdbs, err := c.Schemar.Databases(tx, "")
+	if err != nil {
+		c.logger.Printf("couldn't get databases: %v", err)
+	}
+
+	for _, qdb := range qdbs {
+		c.snapAllForDatabase(tx, qdb.QualifiedID())
+	}
+
+	c.logger.Debugf("TURTLE: snapAll complete")
+}
+
+func (c *Controller) snapAllForDatabase(tx dax.Transaction, qdbid dax.QualifiedDatabaseID) {
+	c.logger.Debugf("TURTLE: snapAllForDatabase: %s", qdbid)
+	computeNodes, err := c.Balancer.CurrentState(tx, dax.RoleTypeCompute, qdbid)
 	if err != nil {
 		c.logger.Printf("Error getting compute balancer state for snapping turtle: %v", err)
 	}
@@ -52,15 +74,17 @@ func (c *Controller) snapAll() {
 			if err != nil {
 				c.logger.Printf("couldn't decode a shard out of the job: '%s', err: %v", workerInfo.Jobs[i], err)
 			}
-			c.SnapshotShardData(ctx, j.t.QualifiedTableID(), j.shardNum())
+			if err := c.snapshotShardData(tx, j.t.QualifiedTableID(), j.shardNum()); err != nil {
+				c.logger.Printf("Couldn't snapshot table: %s, shard: %d, error: %v", j.t, j.shardNum(), err)
+			}
 		}
 		i++
 	}
 
 	// Get all tables across all orgs/dbs so we can snapshot all keyed
-	// fields and look up whether a table is keyed to snapshot it's
+	// fields and look up whether a table is keyed to snapshot its
 	// partitions.
-	tables, err := c.Schemar.Tables(ctx, dax.TableQualifier{})
+	tables, err := c.Schemar.Tables(tx, dax.QualifiedDatabaseID{})
 	if err != nil {
 		c.logger.Printf("Couldn't get schema for snapshotting keys: %v", err)
 		return
@@ -71,8 +95,7 @@ func (c *Controller) snapAll() {
 		tableMap[table.Key()] = table
 		for _, f := range table.Fields {
 			if f.StringKeys() && !f.IsPrimaryKey() {
-				err := c.SnapshotFieldKeys(ctx, table.QualifiedID(), f.Name)
-				if err != nil {
+				if err := c.snapshotFieldKeys(tx, table.QualifiedID(), f.Name); err != nil {
 					c.logger.Printf("Couldn't snapshot table: %s, field: %s, error: %v", table, f.Name, err)
 				}
 			}
@@ -83,7 +106,7 @@ func (c *Controller) snapAll() {
 	// for any partition that goes with a keyed table. Doing the same
 	// weird nested loop thing to avoid doing all jobs on one node
 	// back to back.
-	translateNodes, err := c.TranslateBalancer.CurrentState(ctx)
+	translateNodes, err := c.Balancer.CurrentState(tx, dax.RoleTypeTranslate, qdbid)
 	if err != nil {
 		c.logger.Printf("Error getting translate balancer state for snapping turtle: %v", err)
 	}
@@ -101,12 +124,14 @@ func (c *Controller) snapAll() {
 			if err != nil {
 				table := tableMap[j.table()]
 				if table.StringKeys() {
-					c.SnapshotTableKeys(ctx, table.QualifiedID(), j.partitionNum())
+					if err := c.snapshotTableKeys(tx, table.QualifiedID(), j.partitionNum()); err != nil {
+						c.logger.Printf("Couldn't snapshot table: %s, partition: %d, error: %v", table, j.partitionNum(), err)
+					}
 				}
 				c.logger.Printf("couldn't decode a partition out of the job: '%s', err: %v", workerInfo.Jobs[i], err)
 			}
 		}
 		i++
 	}
-	c.logger.Debugf("TURTLE: snapAll complete")
+	c.logger.Debugf("TURTLE: snapAllForDatabase complete: %s", qdbid)
 }

dax/mds/controller/stringers.go

@@ -20,6 +20,10 @@ func (p pUnit) String() string {
 	return fmt.Sprintf("%s|part_%d", p.t, p.p)
 }
 
+func (p pUnit) Job() dax.Job {
+	return dax.Job(fmt.Sprintf("%s|part_%d", p.t, p.p))
+}
+
 func (p pUnit) table() dax.TableKey {
 	return p.t
 }
@@ -72,6 +76,10 @@ func (s sUnit) String() string {
 	return fmt.Sprintf("%s|shard_%s", s.t, s.s)
 }
 
+func (s sUnit) Job() dax.Job {
+	return dax.Job(fmt.Sprintf("%s|shard_%s", s.t, s.s))
+}
+
 func (s sUnit) table() dax.TableKey {
 	return s.t
 }

dax/mds/http/handler.go

@@ -18,6 +18,9 @@ func Handler(mds *mds.MDS) http.Handler {
 	router.HandleFunc("/health", server.getHealth).Methods("GET").Name("GetHealth")
 
 	// mds endpoints.
+	router.HandleFunc("/create-database", server.postCreateDatabase).Methods("POST").Name("PostCreateDatabase")
+	router.HandleFunc("/database-by-id", server.postDatabaseByID).Methods("POST").Name("PostDatabaseByID")
+
 	router.HandleFunc("/create-table", server.postCreateTable).Methods("POST").Name("PostCreateTable")
 	router.HandleFunc("/drop-table", server.postDropTable).Methods("POST").Name("PostDropTable")
 	router.HandleFunc("/create-field", server.postCreateField).Methods("POST").Name("PostCreateField")
@@ -57,6 +60,55 @@ func (s *server) getHealth(w http.ResponseWriter, r *http.Request) {
 	w.WriteHeader(http.StatusOK)
 }
 
+// POST /create-database
+func (s *server) postCreateDatabase(w http.ResponseWriter, r *http.Request) {
+	body := r.Body
+	defer body.Close()
+
+	ctx := r.Context()
+
+	req := &dax.QualifiedDatabase{}
+	if err := json.NewDecoder(body).Decode(req); err != nil {
+		http.Error(w, err.Error(), http.StatusBadRequest)
+		return
+	}
+
+	err := s.mds.CreateDatabase(ctx, req)
+	if err != nil {
+		http.Error(w, err.Error(), http.StatusBadRequest)
+		return
+	}
+
+	if err := json.NewEncoder(w).Encode(req); err != nil {
+		http.Error(w, err.Error(), http.StatusBadRequest)
+		return
+	}
+}
+
+// POST /database
+func (s *server) postDatabaseByID(w http.ResponseWriter, r *http.Request) {
+	body := r.Body
+	defer body.Close()
+
+	ctx := r.Context()
+
+	qdbid := dax.QualifiedDatabaseID{}
+	if err := json.NewDecoder(body).Decode(&qdbid); err != nil {
+		http.Error(w, err.Error(), http.StatusBadRequest)
+		return
+	}
+	resp, err := s.mds.DatabaseByID(ctx, qdbid)
+	if err != nil {
+		http.Error(w, err.Error(), http.StatusBadRequest)
+		return
+	}
+
+	if err := json.NewEncoder(w).Encode(resp); err != nil {
+		http.Error(w, err.Error(), http.StatusBadRequest)
+		return
+	}
+}
+
 // POST /create-table
 func (s *server) postCreateTable(w http.ResponseWriter, r *http.Request) {
 	body := r.Body
@@ -76,15 +128,12 @@ func (s *server) postCreateTable(w http.ResponseWriter, r *http.Request) {
 		return
 	}
 
-	resp := CreateTableResponse(*req)
-	if err := json.NewEncoder(w).Encode(resp); err != nil {
+	if err := json.NewEncoder(w).Encode(req); err != nil {
 		http.Error(w, err.Error(), http.StatusBadRequest)
 		return
 	}
 }
 
-type CreateTableResponse dax.QualifiedTable
-
 // POST /table
 func (s *server) postTable(w http.ResponseWriter, r *http.Request) {
 	body := r.Body
@@ -122,7 +171,7 @@ func (s *server) postTableID(w http.ResponseWriter, r *http.Request) {
 		return
 	}
 
-	qtid, err := s.mds.TableID(ctx, req.TableQualifier, req.Name)
+	qtid, err := s.mds.TableID(ctx, req.QualifiedDatabaseID, req.Name)
 	if err != nil {
 		http.Error(w, err.Error(), http.StatusBadRequest)
 		return
@@ -221,10 +270,10 @@ func (s *server) postTables(w http.ResponseWriter, r *http.Request) {
 		return
 	}
 
-	qual := dax.NewTableQualifier(req.OrganizationID, req.DatabaseID)
+	qdbid := dax.NewQualifiedDatabaseID(req.OrganizationID, req.DatabaseID)
 	ids := req.TableIDs
 
-	resp, err := s.mds.Tables(ctx, qual, ids...)
+	resp, err := s.mds.Tables(ctx, qdbid, ids...)
 	if err != nil {
 		http.Error(w, err.Error(), http.StatusBadRequest)
 		return

dax/mds/mds.go

@@ -3,7 +3,6 @@ package mds
 
 import (
 	"context"
-	"fmt"
 	"os"
 	"sync"
 	"time"
@@ -57,7 +56,6 @@ type MDS struct {
 	// we need to be sure to close the boltDBs that are created in mds.New()
 	// whenever mds.Close() is called. These are pointers to those DBs so we can
 	// close them.
-	schemarDB    *boltdb.DB
 	controllerDB *boltdb.DB
 
 	logger logger.Logger
@@ -88,25 +86,20 @@ func New(cfg Config) *MDS {
 		logr.Warnf("no DataDir given (like '/path/to/directory') using temp dir at '%s'", cfg.DataDir)
 	}
 
-	schemarDB, err := boltdb.NewSvcBolt(cfg.DataDir, "schemar", schemarboltdb.SchemarBuckets...)
+	buckets := append(schemarboltdb.SchemarBuckets, balancerboltdb.BalancerBuckets...)
+	controllerDB, err := boltdb.NewSvcBolt(cfg.DataDir, "controller", buckets...)
 	if err != nil {
-		logr.Printf("Error creating schemar db: %v", err)
+		logr.Printf(errors.Wrap(err, "creating controller bolt").Error())
 		os.Exit(1)
 	}
 
-	schemar := schemarboltdb.NewSchemar(schemarDB, logr)
-
-	controllerDB, err := boltdb.NewSvcBolt(cfg.DataDir, "balancer", naiveboltdb.NaiveBalancerBuckets...)
-	if err != nil {
-		logr.Printf(errors.Wrap(err, "creating balancer bolt").Error())
-		os.Exit(1)
-	}
+	schemar := schemarboltdb.NewSchemar(controllerDB, logr)
 
 	controllerCfg := controller.Config{
-		Director:          cfg.Director,
-		Schemar:           schemar,
-		ComputeBalancer:   naiveboltdb.NewBalancer("compute", controllerDB, logr),
-		TranslateBalancer: naiveboltdb.NewBalancer("translate", controllerDB, logr),
+		Director: cfg.Director,
+		Schemar:  schemar,
+
+		Balancer: balancerboltdb.NewBalancer(controllerDB, schemar, logr),
 
 		RegistrationBatchTimeout: cfg.RegistrationBatchTimeout,
 		SnappingTurtleTimeout:    cfg.SnappingTurtleTimeout,
@@ -140,7 +133,6 @@ func New(cfg Config) *MDS {
 		poller:     poller,
 		schemar:    schemar,
 
-		schemarDB:    schemarDB,
 		controllerDB: controllerDB,
 
 		logger: logr,
@@ -169,9 +161,6 @@ func (m *MDS) Stop() error {
 	m.poller.Stop()
 	m.controller.Stop()
 
-	if m.schemarDB != nil {
-		m.schemarDB.Close()
-	}
 	if m.controllerDB != nil {
 		m.controllerDB.Close()
 	}
@@ -179,38 +168,17 @@ func (m *MDS) Stop() error {
 	return nil
 }
 
-// sanitizeQTID populates Table.ID (by looking up the table, by name, in
-// schemar) for a given table having only a Name value, but no ID.
-func (m *MDS) sanitizeQTID(ctx context.Context, qtid *dax.QualifiedTableID) error {
-	if qtid.ID == "" {
-		nqtid, err := m.schemar.TableID(ctx, qtid.TableQualifier, qtid.Name)
-		if err != nil {
-			return errors.Wrap(err, "getting table ID")
-		}
-		qtid.ID = nqtid.ID
-	}
-	return nil
+// CreateDatabase handles a create table request.
+func (m *MDS) CreateDatabase(ctx context.Context, qdb *dax.QualifiedDatabase) error {
+	return m.controller.CreateDatabase(ctx, qdb)
+}
+
+func (m *MDS) DatabaseByID(ctx context.Context, qdbid dax.QualifiedDatabaseID) (*dax.QualifiedDatabase, error) {
+	return m.controller.DatabaseByID(ctx, qdbid)
 }
 
 // CreateTable handles a create table request.
 func (m *MDS) CreateTable(ctx context.Context, qtbl *dax.QualifiedTable) error {
-	m.mu.Lock()
-	defer m.mu.Unlock()
-
-	// Create Table ID.
-	if _, err := qtbl.CreateID(); err != nil {
-		return errors.Wrap(err, "creating table ID")
-	}
-
-	// Create the table in schemar.
-	if err := m.schemar.CreateTable(ctx, qtbl); err != nil {
-		return errors.Wrapf(err, "creating table: %s", qtbl)
-	}
-
-	// TODO: if error here, we should probably roll-back the
-	// schemar.CreateTable() request.
-
-	// Add the table to the controller.
 	return m.controller.CreateTable(ctx, qtbl)
 }
 
@@ -218,216 +186,61 @@ func (m *MDS) CreateTable(ctx context.Context, qtbl *dax.QualifiedTable) error {
 // reason about consistency here? What if controller DropTable
 // succeeds, but schemar fails?
 func (m *MDS) DropTable(ctx context.Context, qtid dax.QualifiedTableID) error {
-	m.mu.Lock()
-	defer m.mu.Unlock()
-
-	if err := m.sanitizeQTID(ctx, &qtid); err != nil {
-		return errors.Wrap(err, "sanitizing")
-	}
-
-	if err := m.controller.DropTable(ctx, qtid); err != nil {
-		return errors.Wrapf(err, "dropping table: %s", qtid)
-	}
-
-	return m.schemar.DropTable(ctx, qtid)
-}
-
-type CreateFieldRequest struct {
-	Table dax.TableName
-	Field *dax.Field
+	return m.controller.DropTable(ctx, qtid)
 }
 
 // CreateField handles a create Field request.
 func (m *MDS) CreateField(ctx context.Context, qtid dax.QualifiedTableID, fld *dax.Field) error {
-	m.mu.Lock()
-	defer m.mu.Unlock()
-
-	if err := m.sanitizeQTID(ctx, &qtid); err != nil {
-		return errors.Wrap(err, "sanitizing")
-	}
-
-	// Create the field in schemar.
-	if err := m.schemar.CreateField(ctx, qtid, fld); err != nil {
-		return errors.Wrapf(err, "creating field: %s, %s", qtid, fld)
-	}
-
-	// Add the table to the controller.
 	return m.controller.CreateField(ctx, qtid, fld)
 }
 
 // DropField handles a drop Field request.
 func (m *MDS) DropField(ctx context.Context, qtid dax.QualifiedTableID, fldName dax.FieldName) error {
-	m.mu.Lock()
-	defer m.mu.Unlock()
-
-	if err := m.sanitizeQTID(ctx, &qtid); err != nil {
-		return errors.Wrap(err, "sanitizing")
-	}
-
-	// Drop the field from schemar.
-	if err := m.schemar.DropField(ctx, qtid, fldName); err != nil {
-		return errors.Wrapf(err, "dropping field: %s, %s", qtid, fldName)
-	}
-
-	// Drop the field from the controller.
 	return m.controller.DropField(ctx, qtid, fldName)
 }
 
-type DropFieldResponse struct{}
-
 // Table handles a table request.
 func (m *MDS) Table(ctx context.Context, qtid dax.QualifiedTableID) (*dax.QualifiedTable, error) {
-	m.mu.RLock()
-	defer m.mu.RUnlock()
-
-	if err := m.sanitizeQTID(ctx, &qtid); err != nil {
-		return nil, errors.Wrap(err, "sanitizing")
-	}
-
-	return m.schemar.Table(ctx, qtid)
+	return m.controller.Table(ctx, qtid)
 }
 
 // Tables handles a tables request.
-func (m *MDS) Tables(ctx context.Context, qual dax.TableQualifier, ids ...dax.TableID) ([]*dax.QualifiedTable, error) {
-	m.mu.RLock()
-	defer m.mu.RUnlock()
-
-	return m.schemar.Tables(ctx, qual, ids...)
+func (m *MDS) Tables(ctx context.Context, qdbid dax.QualifiedDatabaseID, ids ...dax.TableID) ([]*dax.QualifiedTable, error) {
+	return m.controller.Tables(ctx, qdbid, ids...)
 }
 
 // TableID handles a table id (i.e. by name) request.
-func (m *MDS) TableID(ctx context.Context, qual dax.TableQualifier, name dax.TableName) (dax.QualifiedTableID, error) {
-	m.mu.RLock()
-	defer m.mu.RUnlock()
-
-	return m.schemar.TableID(ctx, qual, name)
+func (m *MDS) TableID(ctx context.Context, qdbid dax.QualifiedDatabaseID, name dax.TableName) (dax.QualifiedTableID, error) {
+	return m.controller.TableID(ctx, qdbid, name)
 }
 
 // IngestPartition handles an ingest partition request.
 func (m *MDS) IngestPartition(ctx context.Context, qtid dax.QualifiedTableID, partnNum dax.PartitionNum) (dax.Address, error) {
-	m.mu.RLock()
-	defer m.mu.RUnlock()
-
-	if err := m.sanitizeQTID(ctx, &qtid); err != nil {
-		return "", errors.Wrap(err, "sanitizing")
-	}
-
-	// Verify that the table exists.
-	if _, err := m.schemar.Table(ctx, qtid); err != nil {
-		return "", err
-	}
-
-	partitions := dax.PartitionNums{partnNum}
-
-	nodes, err := m.controller.TranslateNodes(ctx, qtid, partitions, true)
-	if err != nil {
-		return "", err
-	}
-
-	if l := len(nodes); l == 0 {
-		return "", controller.NewErrNoAvailableNode()
-	} else if l > 1 {
-		return "", controller.NewErrInternal(
-			fmt.Sprintf("unexpected number of nodes: %d", l))
-	}
-
-	node := nodes[0]
-
-	// Verify that the node returned is actually responsible for the partition
-	// requested.
-	if node.Table != qtid.Key() {
-		return "", controller.NewErrInternal(
-			fmt.Sprintf("table returned (%s) does not match requested (%s)", node.Table, qtid))
-	} else if l := len(node.Partitions); l != 1 {
-		return "", controller.NewErrInternal(
-			fmt.Sprintf("unexpected number of partitions returned: %d", l))
-	} else if p := node.Partitions[0]; p != partnNum {
-		return "", controller.NewErrInternal(
-			fmt.Sprintf("partition returned (%d) does not match requested (%d)", p, partnNum))
-	}
-
-	return node.Address, nil
+	return m.controller.IngestPartition(ctx, qtid, partnNum)
 }
 
 // IngestShard handles an ingest shard request.
 func (m *MDS) IngestShard(ctx context.Context, qtid dax.QualifiedTableID, shrdNum dax.ShardNum) (dax.Address, error) {
-	m.mu.RLock()
-	defer m.mu.RUnlock()
-
-	if err := m.sanitizeQTID(ctx, &qtid); err != nil {
-		return "", errors.Wrap(err, "sanitizing")
-	}
-
-	// Verify that the table exists.
-	if _, err := m.schemar.Table(ctx, qtid); err != nil {
-		return "", err
-	}
-
-	shards := dax.ShardNums{shrdNum}
-
-	nodes, err := m.controller.ComputeNodes(ctx, qtid, shards, true)
-	if err != nil {
-		return "", err
-	}
-
-	if l := len(nodes); l == 0 {
-		return "", controller.NewErrNoAvailableNode()
-	} else if l > 1 {
-		return "", controller.NewErrInternal(
-			fmt.Sprintf("unexpected number of nodes: %d", l))
-	}
-
-	node := nodes[0]
-
-	// Verify that the node returned is actually responsible for the shard
-	// requested.
-	if node.Table != qtid.Key() {
-		return "", controller.NewErrInternal(
-			fmt.Sprintf("table returned (%s) does not match requested (%s)", node.Table, qtid))
-	} else if l := len(node.Shards); l != 1 {
-		return "", controller.NewErrInternal(
-			fmt.Sprintf("unexpected number of shards returned: %d", l))
-	} else if s := node.Shards[0]; s != shrdNum {
-		return "", controller.NewErrInternal(
-			fmt.Sprintf("shard returned (%d) does not match requested (%d)", s, shrdNum))
-	}
-
-	return node.Address, nil
+	return m.controller.IngestShard(ctx, qtid, shrdNum)
 }
 
 // SnapshotTable handles a snapshot table request.
 func (m *MDS) SnapshotTable(ctx context.Context, qtid dax.QualifiedTableID) error {
-	if err := m.sanitizeQTID(ctx, &qtid); err != nil {
-		return errors.Wrap(err, "sanitizing")
-	}
-
 	return m.controller.SnapshotTable(ctx, qtid)
 }
 
 // SnapshotShardData handles a snapshot shard request.
 func (m *MDS) SnapshotShardData(ctx context.Context, qtid dax.QualifiedTableID, shardNum dax.ShardNum) error {
-	if err := m.sanitizeQTID(ctx, &qtid); err != nil {
-		return errors.Wrap(err, "sanitizing")
-	}
-
 	return m.controller.SnapshotShardData(ctx, qtid, shardNum)
 }
 
 // SnapshotTableKeys handles a snapshot table/keys request.
 func (m *MDS) SnapshotTableKeys(ctx context.Context, qtid dax.QualifiedTableID, partitionNum dax.PartitionNum) error {
-	if err := m.sanitizeQTID(ctx, &qtid); err != nil {
-		return errors.Wrap(err, "sanitizing")
-	}
-
 	return m.controller.SnapshotTableKeys(ctx, qtid, partitionNum)
 }
 
 // SnapshotFieldKeys handles a snapshot field/keys request.
 func (m *MDS) SnapshotFieldKeys(ctx context.Context, qtid dax.QualifiedTableID, fldName dax.FieldName) error {
-	if err := m.sanitizeQTID(ctx, &qtid); err != nil {
-		return errors.Wrap(err, "sanitizing")
-	}
-
 	return m.controller.SnapshotFieldKeys(ctx, qtid, fldName)
 }
 
@@ -466,11 +279,7 @@ func (m *MDS) DeregisterNodes(ctx context.Context, addrs ...dax.Address) error {
 // ComputeNodes gets the compute nodes responsible for the table/shards
 // specified in the ComputeNodeRequest.
 func (m *MDS) ComputeNodes(ctx context.Context, qtid dax.QualifiedTableID, shardNums ...dax.ShardNum) ([]dax.ComputeNode, error) {
-	if err := m.sanitizeQTID(ctx, &qtid); err != nil {
-		return nil, errors.Wrap(err, "sanitizing")
-	}
-
-	return m.controller.ComputeNodes(ctx, qtid, shardNums, false)
+	return m.controller.ComputeNodes(ctx, qtid, shardNums)
 }
 
 func (m *MDS) DebugNodes(ctx context.Context) ([]*dax.Node, error) {
@@ -480,9 +289,5 @@ func (m *MDS) DebugNodes(ctx context.Context) ([]*dax.Node, error) {
 // TranslateNodes gets the translate nodes responsible for the table/partitions
 // specified in the TranslateNodeRequest.
 func (m *MDS) TranslateNodes(ctx context.Context, qtid dax.QualifiedTableID, partitionNums ...dax.PartitionNum) ([]dax.TranslateNode, error) {
-	if err := m.sanitizeQTID(ctx, &qtid); err != nil {
-		return nil, errors.Wrap(err, "sanitizing")
-	}
-
-	return m.controller.TranslateNodes(ctx, qtid, partitionNums, false)
+	return m.controller.TranslateNodes(ctx, qtid, partitionNums)
 }

dax/mds/schemar/boltdb/schemar.go

@@ -4,7 +4,6 @@ package boltdb
 
 import (
 	"bytes"
-	"context"
 	"encoding/json"
 	"fmt"
 	"strings"
@@ -44,9 +43,224 @@ func NewSchemar(db *boltdb.DB, logger logger.Logger) *Schemar {
 	}
 }
 
+// CreateDatabase creates the database provided. If a database with the same
+// name already exists then an error is returned. For now, we are not going to
+// store the tables in the schemar Database struct.
+func (s *Schemar) CreateDatabase(tx dax.Transaction, qdb *dax.QualifiedDatabase) error {
+	// Ensure the database id is not blank.
+	if qdb.ID == "" {
+		return schemar.NewErrDatabaseIDInvalid(qdb.ID)
+	}
+
+	// Ensure the database name is not blank.
+	if qdb.Name == "" {
+		return schemar.NewErrDatabaseNameInvalid(qdb.Name)
+	}
+
+	// Set the CreateAt value for the database.
+	// TODO(tlt): We may want to consider erroring here if the value is != 0.
+	if qdb.CreatedAt == 0 {
+		now := timestamp()
+		qdb.CreatedAt = now
+	}
+
+	//////////// end validation
+
+	txx, ok := tx.(*boltdb.Tx)
+	if !ok {
+		return dax.NewErrInvalidTransaction()
+	}
+
+	// Ensure a database with that ID doesn't already exist.
+	if db, _ := s.databaseByID(txx, qdb.OrganizationID, qdb.ID); db != nil {
+		return dax.NewErrDatabaseIDExists(qdb.QualifiedID())
+	}
+
+	if err := s.putDatabase(txx, qdb); err != nil {
+		return errors.Wrap(err, "putting database")
+	}
+
+	// In addition to storing the database in databaseKey, we want to store a
+	// reverse-lookup (i.e. index) on database name to the databaseKey.
+	if err := s.putDatabaseName(txx, qdb); err != nil {
+		return errors.Wrap(err, "putting database name")
+	}
+
+	return nil
+}
+
+func (s *Schemar) DatabaseByID(tx dax.Transaction, qdbid dax.QualifiedDatabaseID) (*dax.QualifiedDatabase, error) {
+	txx, ok := tx.(*boltdb.Tx)
+	if !ok {
+		return nil, dax.NewErrInvalidTransaction()
+	}
+
+	return s.databaseByID(txx, qdbid.OrganizationID, qdbid.DatabaseID)
+}
+
+func (s *Schemar) databaseByID(tx *boltdb.Tx, orgID dax.OrganizationID, id dax.DatabaseID) (*dax.QualifiedDatabase, error) {
+	bkt := tx.Bucket(bucketSchemar)
+	if bkt == nil {
+		return nil, errors.Errorf(boltdb.ErrFmtBucketNotFound, bucketSchemar)
+	}
+
+	b := bkt.Get(databaseKey(orgID, id))
+	if b == nil {
+		return nil, dax.NewErrDatabaseIDDoesNotExist(dax.QualifiedDatabaseID{OrganizationID: orgID, DatabaseID: id})
+	}
+
+	database := &dax.QualifiedDatabase{}
+	if err := json.Unmarshal(b, database); err != nil {
+		return nil, errors.Wrap(err, "unmarshalling database json")
+	}
+
+	return database, nil
+}
+
+func (s *Schemar) putDatabase(tx *boltdb.Tx, qdb *dax.QualifiedDatabase) error {
+	bkt := tx.Bucket(bucketSchemar)
+	if bkt == nil {
+		return errors.Errorf(boltdb.ErrFmtBucketNotFound, bucketSchemar)
+	}
+
+	val, err := json.Marshal(qdb)
+	if err != nil {
+		return errors.Wrap(err, "marshalling database to json")
+	}
+
+	return bkt.Put(databaseKey(qdb.OrganizationID, qdb.ID), val)
+}
+
+func (s *Schemar) putDatabaseName(tx *boltdb.Tx, qdb *dax.QualifiedDatabase) error {
+	bkt := tx.Bucket(bucketSchemar)
+	if bkt == nil {
+		return errors.Errorf(boltdb.ErrFmtBucketNotFound, bucketSchemar)
+	}
+
+	return bkt.Put(databaseNameKey(qdb.OrganizationID, qdb.Name), databaseKey(qdb.OrganizationID, qdb.ID))
+}
+
+// DropDatabase drops the given database. If the named/IDed database does not
+// exist then an error is returned.
+func (s *Schemar) DropDatabase(tx dax.Transaction, qdbid dax.QualifiedDatabaseID) error {
+	txx, ok := tx.(*boltdb.Tx)
+	if !ok {
+		return dax.NewErrInvalidTransaction()
+	}
+
+	// Ensure the database exists.
+	qdb, err := s.databaseByID(txx, qdbid.OrganizationID, qdbid.DatabaseID)
+	if err != nil {
+		return errors.Wrap(err, "getting database by id")
+	}
+
+	bkt := txx.Bucket(bucketSchemar)
+	if bkt == nil {
+		return errors.Errorf(boltdb.ErrFmtBucketNotFound, bucketSchemar)
+	}
+
+	// Delete the database by ID.
+	if err := bkt.Delete(databaseKey(qdb.OrganizationID, qdb.ID)); err != nil {
+		return errors.Wrap(err, "deleting database by id")
+	}
+
+	// Delete the reverse-lookup database by Name.
+	if err := bkt.Delete(databaseNameKey(qdb.OrganizationID, qdb.Name)); err != nil {
+		return errors.Wrap(err, "deleting database by name")
+	}
+
+	return nil
+}
+
+// SetDatabaseOptions overwrites the existing database options with those
+// provided for the given database.
+func (s *Schemar) SetDatabaseOptions(tx dax.Transaction, qdbid dax.QualifiedDatabaseID, opts dax.DatabaseOptions) error {
+	txx, ok := tx.(*boltdb.Tx)
+	if !ok {
+		return dax.NewErrInvalidTransaction()
+	}
+
+	// Get the database.
+	qdb, err := s.databaseByID(txx, qdbid.OrganizationID, qdbid.DatabaseID)
+	if err != nil {
+		return errors.Wrapf(err, "getting database: %s", qdbid)
+	}
+
+	// Set the new options.
+	qdb.Options = opts
+
+	// Put the database.
+	if err := s.putDatabase(txx, qdb); err != nil {
+		return errors.Wrap(err, "putting database")
+	}
+
+	return nil
+}
+
+func (s *Schemar) Databases(tx dax.Transaction, orgID dax.OrganizationID, ids ...dax.DatabaseID) ([]*dax.QualifiedDatabase, error) {
+	txx, ok := tx.(*boltdb.Tx)
+	if !ok {
+		return nil, dax.NewErrInvalidTransaction()
+	}
+
+	return s.getDatabases(txx, orgID, ids...)
+}
+
+func (s *Schemar) getDatabases(tx *boltdb.Tx, orgID dax.OrganizationID, ids ...dax.DatabaseID) (dax.QualifiedDatabases, error) {
+	c := tx.Bucket(bucketSchemar).Cursor()
+
+	// Deserialize rows into Database objects.
+	databases := make(dax.QualifiedDatabases, 0)
+
+	var filterByID bool
+	if len(ids) > 0 {
+		filterByID = true
+	}
+
+	prefix := []byte(fmt.Sprintf(prefixFmtDatabases, orgID))
+	if orgID == "" {
+		prefix = []byte(prefixDatabases)
+	}
+
+	for k, v := c.Seek(prefix); k != nil && bytes.HasPrefix(k, prefix); k, v = c.Next() {
+		if v == nil {
+			s.logger.Printf("nil value for key: %s", k)
+			continue
+		}
+
+		dbID, err := keyDatabaseID(k)
+		if err != nil {
+			return nil, errors.Wrap(err, "getting database from key")
+		}
+
+		// Only include databases provided in the ids filter.
+		if filterByID && !containsDatabaseID(ids, dbID) {
+			continue
+		}
+
+		database := &dax.QualifiedDatabase{}
+		if err := json.Unmarshal(v, database); err != nil {
+			return nil, errors.Wrap(err, "unmarshalling database json")
+		}
+
+		databases = append(databases, database)
+	}
+
+	return databases, nil
+}
+
+func containsDatabaseID(s []dax.DatabaseID, e dax.DatabaseID) bool {
+	for _, a := range s {
+		if a == e {
+			return true
+		}
+	}
+	return false
+}
+
 // CreateTable creates the table provided. If a table with the same name already
 // exists then an error is returned.
-func (s *Schemar) CreateTable(ctx context.Context, qtbl *dax.QualifiedTable) error {
+func (s *Schemar) CreateTable(tx dax.Transaction, qtbl *dax.QualifiedTable) error {
 	// Ensure the table id is not blank.
 	if qtbl.ID == "" {
 		return schemar.NewErrTableIDInvalid(qtbl.ID)
@@ -76,33 +290,37 @@ func (s *Schemar) CreateTable(ctx context.Context, qtbl *dax.QualifiedTable) err
 
 	//////////// end validation
 
-	tx, err := s.db.BeginTx(ctx, true)
-	if err != nil {
-		return errors.Wrap(err, "getting transaction")
+	txx, ok := tx.(*boltdb.Tx)
+	if !ok {
+		return dax.NewErrInvalidTransaction()
+	}
+
+	// Ensure the database, defined in the table's QualifiedDatabaseID, exists.
+	if _, err := s.databaseByID(txx, qtbl.OrganizationID, qtbl.DatabaseID); err != nil {
+		return errors.Wrap(err, "validating database")
 	}
-	defer tx.Rollback()
 
 	// Ensure a table with that ID doesn't already exist.
-	if t, _ := s.tableByID(tx, qtbl.TableQualifier, qtbl.ID); t != nil {
+	if t, _ := s.tableByID(txx, qtbl.QualifiedDatabaseID, qtbl.ID); t != nil {
 		return dax.NewErrTableIDExists(qtbl.QualifiedID())
 	}
 
-	if err := s.putTable(tx, qtbl); err != nil {
+	if err := s.putTable(txx, qtbl); err != nil {
 		return errors.Wrap(err, "putting table")
 	}
 
 	// In addition to storing the table in tableKey, we want to store a reverse-lookup
 	// (i.e. index) on table name to the tableKey.
-	if err := s.putTableName(tx, qtbl); err != nil {
+	if err := s.putTableName(txx, qtbl); err != nil {
 		return errors.Wrap(err, "putting table name")
 	}
 
-	return tx.Commit()
+	return nil
 }
 
 // CreateField creates the field provided in the given table. If a field with
 // the same name already exists then an error is returned.
-func (s *Schemar) CreateField(ctx context.Context, qtid dax.QualifiedTableID, fld *dax.Field) error {
+func (s *Schemar) CreateField(tx dax.Transaction, qtid dax.QualifiedTableID, fld *dax.Field) error {
 	// Ensure the field name is not blank.
 	if fld.Name == "" {
 		return schemar.NewErrFieldNameInvalid(fld.Name)
@@ -110,14 +328,13 @@ func (s *Schemar) CreateField(ctx context.Context, qtid dax.QualifiedTableID, fl
 
 	//////////// end validation
 
-	tx, err := s.db.BeginTx(ctx, true)
-	if err != nil {
-		return errors.Wrap(err, "getting transaction")
+	txx, ok := tx.(*boltdb.Tx)
+	if !ok {
+		return dax.NewErrInvalidTransaction()
 	}
-	defer tx.Rollback()
 
 	// Get the table.
-	qtbl, err := s.tableByQTID(tx, qtid)
+	qtbl, err := s.tableByQTID(txx, qtid)
 	if err != nil {
 		return errors.Wrap(err, "getting table by id")
 	}
@@ -130,23 +347,22 @@ func (s *Schemar) CreateField(ctx context.Context, qtid dax.QualifiedTableID, fl
 	qtbl.Fields = append(qtbl.Fields, fld)
 
 	// Write table back to database.
-	if err := s.putTable(tx, qtbl); err != nil {
+	if err := s.putTable(txx, qtbl); err != nil {
 		return errors.Wrap(err, "putting table")
 	}
 
-	return tx.Commit()
+	return nil
 }
 
 // DropField removes the field from the table.
-func (s *Schemar) DropField(ctx context.Context, qtid dax.QualifiedTableID, fldName dax.FieldName) error {
-	tx, err := s.db.BeginTx(ctx, true)
-	if err != nil {
-		return errors.Wrap(err, "getting transaction")
+func (s *Schemar) DropField(tx dax.Transaction, qtid dax.QualifiedTableID, fldName dax.FieldName) error {
+	txx, ok := tx.(*boltdb.Tx)
+	if !ok {
+		return dax.NewErrInvalidTransaction()
 	}
-	defer tx.Rollback()
 
 	// Get the table.
-	qtbl, err := s.tableByQTID(tx, qtid)
+	qtbl, err := s.tableByQTID(txx, qtid)
 	if err != nil {
 		return errors.Wrap(err, "getting table by id")
 	}
@@ -159,11 +375,11 @@ func (s *Schemar) DropField(ctx context.Context, qtid dax.QualifiedTableID, fldN
 	_ = qtbl.RemoveField(fldName)
 
 	// Write table back to database.
-	if err := s.putTable(tx, qtbl); err != nil {
+	if err := s.putTable(txx, qtbl); err != nil {
 		return errors.Wrap(err, "putting table")
 	}
 
-	return tx.Commit()
+	return nil
 }
 
 func (s *Schemar) putTable(tx *boltdb.Tx, qtbl *dax.QualifiedTable) error {
@@ -191,43 +407,43 @@ func (s *Schemar) putTableName(tx *boltdb.Tx, qtbl *dax.QualifiedTable) error {
 
 // Table returns the TableInfo for the given table. An error is returned if the
 // table does not exist.
-func (s *Schemar) Table(ctx context.Context, qtid dax.QualifiedTableID) (*dax.QualifiedTable, error) {
-	tx, err := s.db.BeginTx(ctx, false)
-	if err != nil {
-		return nil, errors.Wrap(err, "beginning tx")
+func (s *Schemar) Table(tx dax.Transaction, qtid dax.QualifiedTableID) (*dax.QualifiedTable, error) {
+	txx, ok := tx.(*boltdb.Tx)
+	if !ok {
+		return nil, dax.NewErrInvalidTransaction()
 	}
-	defer tx.Rollback()
 
-	return s.tableByQTID(tx, qtid)
+	return s.tableByQTID(txx, qtid)
 }
 
-// tableByQTID gets the full qualified table by the QualifiedTableID whether it has Name or ID set.
+// tableByQTID gets the full qualified table by the QualifiedTableID whether it
+// has Name or ID set.
 func (s *Schemar) tableByQTID(tx *boltdb.Tx, qtid dax.QualifiedTableID) (*dax.QualifiedTable, error) {
 	if qtid.ID == "" {
-		return s.tableByName(tx, qtid.TableQualifier, qtid.Name)
+		return s.tableByName(tx, qtid.QualifiedDatabaseID, qtid.Name)
 	}
 
-	return s.tableByID(tx, qtid.TableQualifier, qtid.ID)
+	return s.tableByID(tx, qtid.QualifiedDatabaseID, qtid.ID)
 }
 
-func (s *Schemar) tableByName(tx *boltdb.Tx, qual dax.TableQualifier, name dax.TableName) (*dax.QualifiedTable, error) {
-	qtid, err := s.tableIDByName(tx, qual, name)
+func (s *Schemar) tableByName(tx *boltdb.Tx, qdbid dax.QualifiedDatabaseID, name dax.TableName) (*dax.QualifiedTable, error) {
+	qtid, err := s.tableIDByName(tx, qdbid, name)
 	if err != nil {
 		return nil, errors.Wrap(err, "getting table ID")
 	}
 
-	return s.tableByID(tx, qtid.TableQualifier, qtid.ID) // TODO remove?
+	return s.tableByID(tx, qtid.QualifiedDatabaseID, qtid.ID) // TODO remove?
 }
 
-func (s *Schemar) tableByID(tx *boltdb.Tx, qual dax.TableQualifier, id dax.TableID) (*dax.QualifiedTable, error) {
+func (s *Schemar) tableByID(tx *boltdb.Tx, qdbid dax.QualifiedDatabaseID, id dax.TableID) (*dax.QualifiedTable, error) {
 	bkt := tx.Bucket(bucketSchemar)
 	if bkt == nil {
 		return nil, errors.Errorf(boltdb.ErrFmtBucketNotFound, bucketSchemar)
 	}
 
-	b := bkt.Get(tableKey(qual.OrganizationID, qual.DatabaseID, id))
+	b := bkt.Get(tableKey(qdbid.OrganizationID, qdbid.DatabaseID, id))
 	if b == nil {
-		return nil, dax.NewErrTableIDDoesNotExist(dax.QualifiedTableID{TableQualifier: qual, ID: id})
+		return nil, dax.NewErrTableIDDoesNotExist(dax.QualifiedTableID{QualifiedDatabaseID: qdbid, ID: id})
 	}
 
 	table := &dax.QualifiedTable{}
@@ -238,13 +454,13 @@ func (s *Schemar) tableByID(tx *boltdb.Tx, qual dax.TableQualifier, id dax.Table
 	return table, nil
 }
 
-func (s *Schemar) tableIDByName(tx *boltdb.Tx, qual dax.TableQualifier, name dax.TableName) (dax.QualifiedTableID, error) {
+func (s *Schemar) tableIDByName(tx *boltdb.Tx, qdbid dax.QualifiedDatabaseID, name dax.TableName) (dax.QualifiedTableID, error) {
 	bkt := tx.Bucket(bucketSchemar)
 	if bkt == nil {
 		return dax.QualifiedTableID{}, errors.Errorf(boltdb.ErrFmtBucketNotFound, bucketSchemar)
 	}
 
-	b := bkt.Get(tableNameKey(qual.OrganizationID, qual.DatabaseID, name))
+	b := bkt.Get(tableNameKey(qdbid.OrganizationID, qdbid.DatabaseID, name))
 	if b == nil {
 		return dax.QualifiedTableID{}, dax.NewErrTableNameDoesNotExist(name)
 	}
@@ -254,17 +470,16 @@ func (s *Schemar) tableIDByName(tx *boltdb.Tx, qual dax.TableQualifier, name dax
 
 // Tables returns a list of Table for all existing tables. If one or more table
 // IDs is provided, then only those will be included in the output.
-func (s *Schemar) Tables(ctx context.Context, qual dax.TableQualifier, ids ...dax.TableID) ([]*dax.QualifiedTable, error) {
-	tx, err := s.db.BeginTx(ctx, false)
-	if err != nil {
-		return nil, errors.Wrap(err, "beginning tx")
+func (s *Schemar) Tables(tx dax.Transaction, qdbid dax.QualifiedDatabaseID, ids ...dax.TableID) ([]*dax.QualifiedTable, error) {
+	txx, ok := tx.(*boltdb.Tx)
+	if !ok {
+		return nil, dax.NewErrInvalidTransaction()
 	}
-	defer tx.Rollback()
 
-	return s.getTables(ctx, tx, qual, ids...)
+	return s.getTables(txx, qdbid, ids...)
 }
 
-func (s *Schemar) getTables(ctx context.Context, tx *boltdb.Tx, qual dax.TableQualifier, ids ...dax.TableID) (dax.QualifiedTables, error) {
+func (s *Schemar) getTables(tx *boltdb.Tx, qdbid dax.QualifiedDatabaseID, ids ...dax.TableID) (dax.QualifiedTables, error) {
 	c := tx.Bucket(bucketSchemar).Cursor()
 
 	// Deserialize rows into Table objects.
@@ -275,11 +490,11 @@ func (s *Schemar) getTables(ctx context.Context, tx *boltdb.Tx, qual dax.TableQu
 		filterByID = true
 	}
 
-	prefix := []byte(fmt.Sprintf(prefixFmtTables, qual.OrganizationID, qual.DatabaseID))
-	if qual.OrganizationID == "" && qual.DatabaseID == "" {
+	prefix := []byte(fmt.Sprintf(prefixFmtTables, qdbid.OrganizationID, qdbid.DatabaseID))
+	if qdbid.OrganizationID == "" && qdbid.DatabaseID == "" {
 		prefix = []byte(prefixTables)
-	} else if qual.DatabaseID == "" {
-		prefix = []byte(fmt.Sprintf(prefixFmtTablesOrg, qual.OrganizationID))
+	} else if qdbid.DatabaseID == "" {
+		prefix = []byte(fmt.Sprintf(prefixFmtTablesOrg, qdbid.OrganizationID))
 	}
 
 	for k, v := c.Seek(prefix); k != nil && bytes.HasPrefix(k, prefix); k, v = c.Next() {
@@ -320,20 +535,19 @@ func containsTableID(s []dax.TableID, e dax.TableID) bool {
 
 // DropTable drops the given table. If the named/IDed table does not exist
 // then an error is returned.
-func (s *Schemar) DropTable(ctx context.Context, qtid dax.QualifiedTableID) error {
-	tx, err := s.db.BeginTx(ctx, true)
-	if err != nil {
-		return errors.Wrap(err, "getting transaction")
+func (s *Schemar) DropTable(tx dax.Transaction, qtid dax.QualifiedTableID) error {
+	txx, ok := tx.(*boltdb.Tx)
+	if !ok {
+		return dax.NewErrInvalidTransaction()
 	}
-	defer tx.Rollback()
 
 	// Ensure the table exists.
-	qtbl, err := s.tableByQTID(tx, qtid)
+	qtbl, err := s.tableByQTID(txx, qtid)
 	if err != nil {
 		return errors.Wrap(err, "getting table by id")
 	}
 
-	bkt := tx.Bucket(bucketSchemar)
+	bkt := txx.Bucket(bucketSchemar)
 	if bkt == nil {
 		return errors.Errorf(boltdb.ErrFmtBucketNotFound, bucketSchemar)
 	}
@@ -348,16 +562,52 @@ func (s *Schemar) DropTable(ctx context.Context, qtid dax.QualifiedTableID) erro
 		return errors.Wrap(err, "deleting table by name")
 	}
 
-	return tx.Commit()
+	return nil
+}
+
+func (s *Schemar) TableID(tx dax.Transaction, qdbid dax.QualifiedDatabaseID, name dax.TableName) (dax.QualifiedTableID, error) {
+	txx, ok := tx.(*boltdb.Tx)
+	if !ok {
+		return dax.QualifiedTableID{}, dax.NewErrInvalidTransaction()
+	}
+
+	return s.tableIDByName(txx, qdbid, name)
 }
 
 const (
 	prefixTables        = "tables/"
-	prefixFmtTablesOrg  = prefixTables + "%s/"
-	prefixFmtTables     = prefixFmtTablesOrg + "%s/"
-	prefixFmtTableNames = "tablenames/%s/%s/"
+	prefixFmtTablesOrg  = prefixTables + "%s/"       // org-id
+	prefixFmtTables     = prefixFmtTablesOrg + "%s/" // db-id
+	prefixFmtTableNames = "tablenames/%s/%s/"        // org-id, db-id
+
+	prefixDatabases        = "databases/"
+	prefixFmtDatabases     = prefixDatabases + "%s/"   // org-id
+	prefixFmtDatabase      = prefixFmtDatabases + "%s" // db-id
+	prefixFmtDatabaseNames = "databasenames/%s/"       // org-id
 )
 
+// databaseKey returns a key based on a qualified database ID.
+func databaseKey(orgID dax.OrganizationID, dbID dax.DatabaseID) []byte {
+	key := fmt.Sprintf(prefixFmtDatabase, orgID, dbID)
+	return []byte(key)
+}
+
+// databaseNameKey returns a key based on a qualified database name.
+func databaseNameKey(orgID dax.OrganizationID, name dax.DatabaseName) []byte {
+	key := fmt.Sprintf(prefixFmtDatabaseNames+"%s", orgID, name)
+	return []byte(key)
+}
+
+// keyDatabaseID gets the DatabaseID out of the key.
+func keyDatabaseID(key []byte) (dax.DatabaseID, error) {
+	parts := strings.Split(string(key), "/")
+	if len(parts) != 3 {
+		return "", errors.New(errors.ErrUncoded, "database key format expected: `databases/orgID/dbID`")
+	}
+
+	return dax.DatabaseID(parts[2]), nil
+}
+
 // tableKey returns a key based on a qualified table ID.
 func tableKey(orgID dax.OrganizationID, dbID dax.DatabaseID, tblID dax.TableID) []byte {
 	key := fmt.Sprintf(prefixFmtTables+"%s", orgID, dbID, tblID)
@@ -388,7 +638,7 @@ func keyQualifiedTableID(key []byte) (dax.QualifiedTableID, error) {
 	}
 
 	return dax.NewQualifiedTableID(
-		dax.NewTableQualifier(
+		dax.NewQualifiedDatabaseID(
 			dax.OrganizationID(parts[1]),
 			dax.DatabaseID(parts[2]),
 		),
@@ -396,16 +646,6 @@ func keyQualifiedTableID(key []byte) (dax.QualifiedTableID, error) {
 	), nil
 }
 
-func (s *Schemar) TableID(ctx context.Context, qual dax.TableQualifier, name dax.TableName) (dax.QualifiedTableID, error) {
-	tx, err := s.db.BeginTx(ctx, false)
-	if err != nil {
-		return dax.QualifiedTableID{}, err
-	}
-	defer tx.Rollback()
-
-	return s.tableIDByName(tx, qual, name)
-}
-
 func timestamp() int64 {
 	return time.Now().UnixNano()
 }

dax/mds/schemar/boltdb/schemar_test.go

@@ -16,6 +16,7 @@ import (
 func TestSchemar(t *testing.T) {
 	orgID := dax.OrganizationID("acme")
 	dbID := dax.DatabaseID("db1")
+	dbName := dax.DatabaseName("dbname1")
 	invalidTableID := dax.TableID("invalidID")
 	tableName := dax.TableName("foo")
 	tableName0 := dax.TableName("foo")
@@ -25,21 +26,36 @@ func TestSchemar(t *testing.T) {
 	partitionN := 12
 
 	ctx := context.Background()
-	qual := dax.NewTableQualifier(orgID, dbID)
+	qdbid := dax.NewQualifiedDatabaseID(orgID, dbID)
 
-	db := testbolt.MustOpenDB(t)
-	defer testbolt.MustCloseDB(t, db)
-
-	t.Cleanup(func() {
-		testbolt.CleanupDB(t, db.Path())
-	})
-
-	// Initialize the buckets.
-	assert.NoError(t, db.InitializeBuckets(boltdb.SchemarBuckets...))
+	qdb := &dax.QualifiedDatabase{
+		OrganizationID: orgID,
+		Database: dax.Database{
+			ID:   dbID,
+			Name: dbName,
+		},
+	}
 
 	t.Run("NewSchemar", func(t *testing.T) {
+		db := testbolt.MustOpenDB(t)
+		defer testbolt.MustCloseDB(t, db)
+
+		t.Cleanup(func() {
+			testbolt.CleanupDB(t, db.Path())
+		})
+
+		// Initialize the buckets.
+		assert.NoError(t, db.InitializeBuckets(boltdb.SchemarBuckets...))
+
 		s := boltdb.NewSchemar(db, logger.NopLogger)
 
+		tx, err := db.BeginTx(ctx, true)
+		assert.NoError(t, err)
+		defer tx.Rollback()
+
+		// Create database.
+		assert.NoError(t, s.CreateDatabase(tx, qdb))
+
 		// Add new table.
 		tbl := dax.NewTable(tableName)
 		tbl.CreateID()
@@ -53,11 +69,11 @@ func TestSchemar(t *testing.T) {
 				Type: dax.BaseTypeInt,
 			},
 		}
-		qtbl := dax.NewQualifiedTable(qual, tbl)
-		assert.NoError(t, s.CreateTable(ctx, qtbl))
+		qtbl := dax.NewQualifiedTable(qdbid, tbl)
+		assert.NoError(t, s.CreateTable(tx, qtbl))
 
 		// Try adding the table again.
-		err := s.CreateTable(ctx, qtbl)
+		err = s.CreateTable(tx, qtbl)
 		if assert.Error(t, err) {
 			assert.True(t, errors.Is(err, dax.ErrTableIDExists))
 		}
@@ -66,32 +82,32 @@ func TestSchemar(t *testing.T) {
 
 		// Get the table.
 		{
-			tbl, err := s.Table(ctx, qtid)
+			tbl, err := s.Table(tx, qtid)
 			assert.NoError(t, err)
 			assert.Equal(t, tableName, tbl.Name)
 		}
 
 		// Drop the table.
-		assert.NoError(t, s.DropTable(ctx, qtid))
+		assert.NoError(t, s.DropTable(tx, qtid))
 
 		// Make sure the reverse-lookup (table by name) was dropped as well.
 		{
-			_, err := s.TableID(ctx, qual, tableName)
+			_, err := s.TableID(tx, qdbid, tableName)
 			if assert.Error(t, err) {
 				assert.True(t, errors.Is(err, dax.ErrTableNameDoesNotExist))
 			}
 		}
 
 		// Try adding the table (i.e. the same table name) again.
-		assert.NoError(t, s.CreateTable(ctx, qtbl))
+		assert.NoError(t, s.CreateTable(tx, qtbl))
 
 		// Drop the table again.
-		assert.NoError(t, s.DropTable(ctx, qtid))
+		assert.NoError(t, s.DropTable(tx, qtid))
 
 		// Drop invalid table.
 		{
-			iqtid := dax.NewQualifiedTableID(qual, invalidTableID)
-			err := s.DropTable(ctx, iqtid)
+			iqtid := dax.NewQualifiedTableID(qdbid, invalidTableID)
+			err := s.DropTable(tx, iqtid)
 			if assert.Error(t, err) {
 				assert.True(t, errors.Is(err, dax.ErrTableIDDoesNotExist))
 			}
@@ -99,19 +115,36 @@ func TestSchemar(t *testing.T) {
 	})
 
 	t.Run("GetTables", func(t *testing.T) {
+		db := testbolt.MustOpenDB(t)
+		defer testbolt.MustCloseDB(t, db)
+
+		t.Cleanup(func() {
+			testbolt.CleanupDB(t, db.Path())
+		})
+
+		// Initialize the buckets.
+		assert.NoError(t, db.InitializeBuckets(boltdb.SchemarBuckets...))
+
 		s := boltdb.NewSchemar(db, logger.NopLogger)
 
+		tx, err := db.BeginTx(ctx, true)
+		assert.NoError(t, err)
+		defer tx.Rollback()
+
+		// Create database.
+		assert.NoError(t, s.CreateDatabase(tx, qdb))
+
 		exp := []*dax.QualifiedTable{}
-		tables, err := s.Tables(ctx, qual)
+		tables, err := s.Tables(tx, qdbid)
 		assert.NoError(t, err)
 		assert.Equal(t, exp, tables)
 
-		qtbl0 := daxtest.TestQualifiedTableWithID(t, qual, tableID0, tableName0, partitionN, false)
-		qtbl1 := daxtest.TestQualifiedTableWithID(t, qual, tableID1, tableName1, partitionN, false)
+		qtbl0 := daxtest.TestQualifiedTableWithID(t, qdbid, tableID0, tableName0, partitionN, false)
+		qtbl1 := daxtest.TestQualifiedTableWithID(t, qdbid, tableID1, tableName1, partitionN, false)
 
 		// Add a couple of tables.
-		assert.NoError(t, s.CreateTable(ctx, qtbl0))
-		assert.NoError(t, s.CreateTable(ctx, qtbl1))
+		assert.NoError(t, s.CreateTable(tx, qtbl0))
+		assert.NoError(t, s.CreateTable(tx, qtbl1))
 
 		exp = []*dax.QualifiedTable{
 			qtbl1,
@@ -119,62 +152,76 @@ func TestSchemar(t *testing.T) {
 		}
 
 		// All tables.
-		tables, err = s.Tables(ctx, qual)
+		tables, err = s.Tables(tx, qdbid)
 		assert.NoError(t, err)
 		assert.Equal(t, exp, tables)
 
 		// With a valid filter.
-		tables, err = s.Tables(ctx, qual, qtbl0.ID)
+		tables, err = s.Tables(tx, qdbid, qtbl0.ID)
 		assert.NoError(t, err)
 		assert.Equal(t, exp[1:], tables)
 
 		// With an invalid filter.
-		tables, err = s.Tables(ctx, qual, invalidTableID)
+		tables, err = s.Tables(tx, qdbid, invalidTableID)
 		assert.NoError(t, err)
 		assert.Equal(t, exp[0:0], tables)
 
 		// With both valid and invalid filters.
-		tables, err = s.Tables(ctx, qual, qtbl0.ID, invalidTableID)
+		tables, err = s.Tables(tx, qdbid, qtbl0.ID, invalidTableID)
 		assert.NoError(t, err)
 		assert.Equal(t, exp[1:], tables)
 
 		// With all valid filters.
-		tables, err = s.Tables(ctx, qual, qtbl0.ID, qtbl1.ID)
+		tables, err = s.Tables(tx, qdbid, qtbl0.ID, qtbl1.ID)
 		assert.NoError(t, err)
 		assert.Equal(t, exp, tables)
 	})
 
 	t.Run("GetTablesAll", func(t *testing.T) {
-		// get a fresh DB
 		db := testbolt.MustOpenDB(t)
 		defer testbolt.MustCloseDB(t, db)
 
 		t.Cleanup(func() {
 			testbolt.CleanupDB(t, db.Path())
 		})
+
 		// Initialize the buckets.
 		assert.NoError(t, db.InitializeBuckets(boltdb.SchemarBuckets...))
 
 		s := boltdb.NewSchemar(db, logger.NopLogger)
 
-		qtbl0 := daxtest.TestQualifiedTableWithID(t, qual, tableID0, tableName0, partitionN, false)
-		orgID2 := dax.OrganizationID("acme2")
-		qual2 := dax.NewTableQualifier(orgID2, dbID)
-		tableID2 := "3"
-		qtbl2 := daxtest.TestQualifiedTableWithID(t, qual2, tableID2, dax.TableName("two"), partitionN, false)
+		tx, err := db.BeginTx(ctx, true)
+		assert.NoError(t, err)
+		defer tx.Rollback()
 
-		assert.NoError(t, s.CreateTable(ctx, qtbl0))
-		assert.NoError(t, s.CreateTable(ctx, qtbl2))
+		qtbl0 := daxtest.TestQualifiedTableWithID(t, qdbid, tableID0, tableName0, partitionN, false)
+		orgID2 := dax.OrganizationID("acme2")
+		qdbid2 := dax.NewQualifiedDatabaseID(orgID2, dbID)
+		tableID2 := "3"
+		qtbl2 := daxtest.TestQualifiedTableWithID(t, qdbid2, tableID2, dax.TableName("two"), partitionN, false)
+
+		// Create databases.
+		assert.NoError(t, s.CreateDatabase(tx, qdb))
+		qdb2 := &dax.QualifiedDatabase{
+			OrganizationID: orgID2,
+			Database: dax.Database{
+				ID:   dbID,
+				Name: dbName,
+			},
+		}
+		assert.NoError(t, s.CreateDatabase(tx, qdb2))
+
+		assert.NoError(t, s.CreateTable(tx, qtbl0))
+		assert.NoError(t, s.CreateTable(tx, qtbl2))
 
 		exp := []*dax.QualifiedTable{qtbl0, qtbl2}
 
-		tables, err := s.Tables(ctx, dax.TableQualifier{})
+		tables, err := s.Tables(tx, dax.QualifiedDatabaseID{})
 		assert.NoError(t, err)
 		assert.Equal(t, exp, tables)
 
-		tables, err = s.Tables(ctx, dax.TableQualifier{OrganizationID: orgID2})
+		tables, err = s.Tables(tx, dax.QualifiedDatabaseID{OrganizationID: orgID2})
 		assert.NoError(t, err)
 		assert.Equal(t, []*dax.QualifiedTable{qtbl2}, tables)
-
 	})
 }

dax/mds/schemar/errors.go

@@ -8,6 +8,9 @@ import (
 )
 
 const (
+	ErrCodeDatabaseIDInvalid   errors.Code = "DatabaseIDInvalid"
+	ErrCodeDatabaseNameInvalid errors.Code = "DatabaseNameInvalid"
+
 	ErrCodeTableIDInvalid    errors.Code = "TableIDInvalid"
 	ErrCodeTableNameInvalid  errors.Code = "TableNameInvalid"
 	ErrCodeInvalidPrimaryKey errors.Code = "InvalidPrimaryKey"
@@ -15,6 +18,20 @@ const (
 	ErrCodeFieldNameInvalid errors.Code = "FieldNameInvalid"
 )
 
+func NewErrDatabaseIDInvalid(databaseID dax.DatabaseID) error {
+	return errors.New(
+		ErrCodeDatabaseIDInvalid,
+		fmt.Sprintf("database ID '%s' is invalid", databaseID),
+	)
+}
+
+func NewErrDatabaseNameInvalid(databaseName dax.DatabaseName) error {
+	return errors.New(
+		ErrCodeDatabaseNameInvalid,
+		fmt.Sprintf("database name '%s' is invalid", databaseName),
+	)
+}
+
 func NewErrTableIDInvalid(tableID dax.TableID) error {
 	return errors.New(
 		ErrCodeTableIDInvalid,

dax/mds/schemar/http/handler.go

@@ -1,199 +0,0 @@
-package http
-
-import (
-	"encoding/json"
-	"net/http"
-
-	"github.com/gorilla/mux"
-	"github.com/featurebasedb/featurebase/v3/dax"
-	"github.com/featurebasedb/featurebase/v3/dax/mds/schemar"
-)
-
-func Handler(s schemar.Schemar) http.Handler {
-	svr := &server{
-		schemar: s,
-	}
-
-	router := mux.NewRouter()
-	router.HandleFunc("/health", svr.getHealth).Methods("GET").Name("GetHealth")
-	router.HandleFunc("/create-table", svr.postCreateTable).Methods("POST").Name("PostCreateTable")
-	router.HandleFunc("/drop-table", svr.postDropTable).Methods("POST").Name("PostDropTable")
-	router.HandleFunc("/table", svr.postTable).Methods("POST").Name("PostTable")
-	router.HandleFunc("/tables", svr.postTables).Methods("POST").Name("PostTables")
-	return router
-}
-
-type server struct {
-	schemar schemar.Schemar
-}
-
-// GET /health
-func (s *server) getHealth(w http.ResponseWriter, r *http.Request) {
-	w.WriteHeader(http.StatusOK)
-}
-
-// POST /create-table
-func (s *server) postCreateTable(w http.ResponseWriter, r *http.Request) {
-	body := r.Body
-	defer body.Close()
-
-	ctx := r.Context()
-
-	req := &dax.QualifiedTable{}
-	if err := json.NewDecoder(body).Decode(req); err != nil {
-		http.Error(w, err.Error(), http.StatusBadRequest)
-		return
-	}
-
-	err := s.schemar.CreateTable(ctx, req)
-	if err != nil {
-		http.Error(w, err.Error(), http.StatusBadRequest)
-		return
-	}
-
-	resp := struct{}{}
-	if err := json.NewEncoder(w).Encode(resp); err != nil {
-		http.Error(w, err.Error(), http.StatusBadRequest)
-		return
-	}
-}
-
-// POST /drop-table
-func (s *server) postDropTable(w http.ResponseWriter, r *http.Request) {
-	body := r.Body
-	defer body.Close()
-
-	ctx := r.Context()
-
-	req := DropTableRequest{}
-	if err := json.NewDecoder(body).Decode(&req); err != nil {
-		http.Error(w, err.Error(), http.StatusBadRequest)
-		return
-	}
-
-	qtid := req.TableKey.QualifiedTableID()
-
-	err := s.schemar.DropTable(ctx, qtid)
-	if err != nil {
-		http.Error(w, err.Error(), http.StatusBadRequest)
-		return
-	}
-
-	resp := struct{}{}
-	if err := json.NewEncoder(w).Encode(resp); err != nil {
-		http.Error(w, err.Error(), http.StatusBadRequest)
-		return
-	}
-}
-
-// // POST /create-field
-// func (s *server) postCreateField(w http.ResponseWriter, r *http.Request) {
-// 	body := r.Body
-// 	defer body.Close()
-
-// 	req := mds.CreateFieldRequest{}
-// 	if err := json.NewDecoder(body).Decode(&req); err != nil {
-// 		http.Error(w, err.Error(), http.StatusBadRequest)
-// 		return
-// 	}
-
-// 	resp, err := s.mds.CreateField(req)
-// 	if err != nil {
-// 		http.Error(w, err.Error(), http.StatusBadRequest)
-// 		return
-// 	}
-
-// 	if err := json.NewEncoder(w).Encode(resp); err != nil {
-// 		http.Error(w, err.Error(), http.StatusBadRequest)
-// 		return
-// 	}
-// }
-
-// // POST /drop-field
-// func (s *server) postDropField(w http.ResponseWriter, r *http.Request) {
-// 	body := r.Body
-// 	defer body.Close()
-
-// 	req := mds.DropFieldRequest{}
-// 	if err := json.NewDecoder(body).Decode(&req); err != nil {
-// 		http.Error(w, err.Error(), http.StatusBadRequest)
-// 		return
-// 	}
-
-// 	resp, err := s.mds.DropField(req)
-// 	if err != nil {
-// 		http.Error(w, err.Error(), http.StatusBadRequest)
-// 		return
-// 	}
-
-// 	if err := json.NewEncoder(w).Encode(resp); err != nil {
-// 		http.Error(w, err.Error(), http.StatusBadRequest)
-// 		return
-// 	}
-// }
-
-// POST /table
-func (s *server) postTable(w http.ResponseWriter, r *http.Request) {
-	body := r.Body
-	defer body.Close()
-
-	ctx := r.Context()
-
-	req := TableRequest{}
-	if err := json.NewDecoder(body).Decode(&req); err != nil {
-		http.Error(w, err.Error(), http.StatusBadRequest)
-		return
-	}
-
-	qtid := req.TableKey.QualifiedTableID()
-
-	resp, err := s.schemar.Table(ctx, qtid)
-	if err != nil {
-		http.Error(w, err.Error(), http.StatusBadRequest)
-		return
-	}
-
-	if err := json.NewEncoder(w).Encode(resp); err != nil {
-		http.Error(w, err.Error(), http.StatusBadRequest)
-		return
-	}
-}
-
-// POST /tables
-func (s *server) postTables(w http.ResponseWriter, r *http.Request) {
-	body := r.Body
-	defer body.Close()
-
-	ctx := r.Context()
-
-	req := TablesRequest{}
-	if err := json.NewDecoder(body).Decode(&req); err != nil {
-		http.Error(w, err.Error(), http.StatusBadRequest)
-		return
-	}
-	qual := dax.NewTableQualifier(req.OrganizationID, req.DatabaseID)
-	resp, err := s.schemar.Tables(ctx, qual, req.TableIDs...)
-	if err != nil {
-		http.Error(w, err.Error(), http.StatusBadRequest)
-		return
-	}
-
-	if err := json.NewEncoder(w).Encode(resp); err != nil {
-		http.Error(w, err.Error(), http.StatusBadRequest)
-		return
-	}
-}
-
-type DropTableRequest struct {
-	TableKey dax.TableKey `json:"table-key"`
-}
-
-type TableRequest struct {
-	TableKey dax.TableKey `json:"table-key"`
-}
-
-type TablesRequest struct {
-	OrganizationID dax.OrganizationID `json:"org-id"`
-	DatabaseID     dax.DatabaseID     `json:"db-id"`
-	TableIDs       dax.TableIDs       `json:"table-ids"`
-}

dax/mds/schemar/schemar.go

@@ -8,23 +8,34 @@ import (
 )
 
 type Schemar interface {
-	CreateTable(context.Context, *dax.QualifiedTable) error
-	DropTable(context.Context, dax.QualifiedTableID) error
-	CreateField(context.Context, dax.QualifiedTableID, *dax.Field) error
-	DropField(context.Context, dax.QualifiedTableID, dax.FieldName) error
-	Table(context.Context, dax.QualifiedTableID) (*dax.QualifiedTable, error)
+	CreateDatabase(dax.Transaction, *dax.QualifiedDatabase) error
+	DropDatabase(dax.Transaction, dax.QualifiedDatabaseID) error
+	DatabaseByID(dax.Transaction, dax.QualifiedDatabaseID) (*dax.QualifiedDatabase, error)
 
-	// Tables returns a list of tables. If the qualifiers DatabaseID
-	// is empty, all tables in the org will be returned. If the
-	// OrganizationID is empty, all tables will be returned. If both
-	// are populated, only tables in that databse will be returned. If
-	// greater than zero table IDs are passed in the third argument,
-	// only tables matching those IDs will be returned.
-	Tables(context.Context, dax.TableQualifier, ...dax.TableID) ([]*dax.QualifiedTable, error)
+	SetDatabaseOptions(dax.Transaction, dax.QualifiedDatabaseID, dax.DatabaseOptions) error
+
+	// Databases returns a list of databases. If the OrganizationID is empty,
+	// all databases will be returned. If greater than zero database IDs are
+	// passed in the second argument, only databases matching those IDs will be
+	// returned.
+	Databases(dax.Transaction, dax.OrganizationID, ...dax.DatabaseID) ([]*dax.QualifiedDatabase, error)
+
+	CreateTable(dax.Transaction, *dax.QualifiedTable) error
+	DropTable(dax.Transaction, dax.QualifiedTableID) error
+	CreateField(dax.Transaction, dax.QualifiedTableID, *dax.Field) error
+	DropField(dax.Transaction, dax.QualifiedTableID, dax.FieldName) error
+	Table(dax.Transaction, dax.QualifiedTableID) (*dax.QualifiedTable, error)
+
+	// Tables returns a list of tables. If the qualifiers DatabaseID is empty,
+	// all tables in the org will be returned. If the OrganizationID is empty,
+	// all tables will be returned. If both are populated, only tables in that
+	// database will be returned. If greater than zero table IDs are passed in
+	// the third argument, only tables matching those IDs will be returned.
+	Tables(dax.Transaction, dax.QualifiedDatabaseID, ...dax.TableID) ([]*dax.QualifiedTable, error)
 
 	// TableID is a reverse-lookup method to get the TableID for a given
 	// qualified TableName.
-	TableID(context.Context, dax.TableQualifier, dax.TableName) (dax.QualifiedTableID, error)
+	TableID(dax.Transaction, dax.QualifiedDatabaseID, dax.TableName) (dax.QualifiedTableID, error)
 }
 
 //////////////////////////////////////////////
@@ -39,23 +50,50 @@ func NewNopSchemar() *NopSchemar {
 	return &NopSchemar{}
 }
 
-func (s *NopSchemar) CreateTable(ctx context.Context, qtbl *dax.QualifiedTable) error { return nil }
-func (s *NopSchemar) DropTable(ctx context.Context, qtid dax.QualifiedTableID) error {
+func (s *NopSchemar) CreateDatabase(tx dax.Transaction, qtbl *dax.QualifiedDatabase) error {
 	return nil
 }
-func (s *NopSchemar) CreateField(ctx context.Context, qtid dax.QualifiedTableID, fld *dax.Field) error {
+
+func (s *NopSchemar) DropDatabase(tx dax.Transaction, qdbid dax.QualifiedDatabaseID) error {
 	return nil
 }
-func (s *NopSchemar) DropField(ctx context.Context, qtid dax.QualifiedTableID, fld dax.FieldName) error {
-	return nil
-}
-func (s *NopSchemar) Table(ctx context.Context, qtid dax.QualifiedTableID) (*dax.QualifiedTable, error) {
+
+func (s *NopSchemar) DatabaseByID(dax.Transaction, dax.QualifiedDatabaseID) (*dax.QualifiedDatabase, error) {
 	return nil, nil
 }
-func (s *NopSchemar) Tables(ctx context.Context, qual dax.TableQualifier, ids ...dax.TableID) ([]*dax.QualifiedTable, error) {
+
+func (s *NopSchemar) SetDatabaseOptions(tx dax.Transaction, qdbid dax.QualifiedDatabaseID, opts dax.DatabaseOptions) error {
+	return nil
+}
+
+func (s *NopSchemar) Databases(dax.Transaction, dax.OrganizationID, ...dax.DatabaseID) ([]*dax.QualifiedDatabase, error) {
+	return nil, nil
+}
+
+func (s *NopSchemar) CreateTable(tx dax.Transaction, qtbl *dax.QualifiedTable) error {
+	return nil
+}
+
+func (s *NopSchemar) DropTable(tx dax.Transaction, qtid dax.QualifiedTableID) error {
+	return nil
+}
+
+func (s *NopSchemar) CreateField(tx dax.Transaction, qtid dax.QualifiedTableID, fld *dax.Field) error {
+	return nil
+}
+
+func (s *NopSchemar) DropField(tx dax.Transaction, qtid dax.QualifiedTableID, fld dax.FieldName) error {
+	return nil
+}
+
+func (s *NopSchemar) Table(tx dax.Transaction, qtid dax.QualifiedTableID) (*dax.QualifiedTable, error) {
+	return nil, nil
+}
+
+func (s *NopSchemar) Tables(tx dax.Transaction, qdbid dax.QualifiedDatabaseID, ids ...dax.TableID) ([]*dax.QualifiedTable, error) {
 	return []*dax.QualifiedTable{}, nil
 }
 
-func (s *NopSchemar) TableID(context.Context, dax.TableQualifier, dax.TableName) (dax.QualifiedTableID, error) {
+func (s *NopSchemar) TableID(dax.Transaction, dax.QualifiedDatabaseID, dax.TableName) (dax.QualifiedTableID, error) {
 	return dax.QualifiedTableID{}, nil
 }

dax/mds/schemar/schemar_test.go

@@ -1,124 +0,0 @@
-package schemar_test
-
-import (
-	"context"
-	"testing"
-
-	"github.com/featurebasedb/featurebase/v3/dax"
-	daxtest "github.com/featurebasedb/featurebase/v3/dax/test"
-	"github.com/featurebasedb/featurebase/v3/errors"
-	"github.com/stretchr/testify/assert"
-)
-
-func TestSchemar(t *testing.T) {
-	orgID := dax.OrganizationID("acme")
-	dbID := dax.DatabaseID("db1")
-	invalidTableID := dax.TableID("invalidID")
-	tableName := dax.TableName("foo")
-	tableName0 := dax.TableName("foo")
-	tableName1 := dax.TableName("bar")
-	tableID0 := "2"
-	tableID1 := "1"
-	partitionN := 12
-
-	ctx := context.Background()
-	qual := dax.NewTableQualifier(orgID, dbID)
-
-	t.Run("NewSchemar", func(t *testing.T) {
-		s, cleanup := daxtest.NewSchemar(t)
-		defer cleanup()
-
-		// Add new table.
-		tbl := dax.NewTable(tableName)
-		tbl.Fields = []*dax.Field{
-			{
-				Name: dax.PrimaryKeyFieldName,
-				Type: dax.BaseTypeString,
-			},
-			{
-				Name: "intField",
-				Type: dax.BaseTypeInt,
-			},
-		}
-		qtbl := dax.NewQualifiedTable(qual, tbl)
-		qtbl.CreateID()
-		assert.NoError(t, s.CreateTable(ctx, qtbl))
-
-		// Try adding the table again.
-		err := s.CreateTable(ctx, qtbl)
-		if assert.Error(t, err) {
-			assert.True(t, errors.Is(err, dax.ErrTableIDExists))
-		}
-
-		qtid := qtbl.QualifiedID()
-
-		// Get the table.
-		{
-			tbl, err := s.Table(ctx, qtid)
-			assert.NoError(t, err)
-			assert.Equal(t, tableName, tbl.Name)
-		}
-
-		// Drop the table.
-		{
-			err := s.DropTable(ctx, qtid)
-			assert.NoError(t, err)
-		}
-
-		// Drop invalid table.
-		{
-			iqtid := dax.NewQualifiedTableID(qual, invalidTableID)
-			err := s.DropTable(ctx, iqtid)
-			if assert.Error(t, err) {
-				assert.True(t, errors.Is(err, dax.ErrTableIDDoesNotExist))
-			}
-		}
-	})
-
-	t.Run("GetTables", func(t *testing.T) {
-		s, cleanup := daxtest.NewSchemar(t)
-		defer cleanup()
-
-		exp := []*dax.QualifiedTable{}
-		tables, err := s.Tables(ctx, qual)
-		assert.NoError(t, err)
-		assert.Equal(t, exp, tables)
-
-		qtbl0 := daxtest.TestQualifiedTableWithID(t, qual, tableID0, tableName0, partitionN, false)
-		qtbl1 := daxtest.TestQualifiedTableWithID(t, qual, tableID1, tableName1, partitionN, false)
-
-		// Add a couple of tables.
-		assert.NoError(t, s.CreateTable(ctx, qtbl0))
-		assert.NoError(t, s.CreateTable(ctx, qtbl1))
-
-		exp = []*dax.QualifiedTable{
-			qtbl1,
-			qtbl0,
-		}
-
-		// All tables.
-		tables, err = s.Tables(ctx, qual)
-		assert.NoError(t, err)
-		assert.Equal(t, exp, tables)
-
-		// With a valid filter.
-		tables, err = s.Tables(ctx, qual, qtbl0.ID)
-		assert.NoError(t, err)
-		assert.Equal(t, exp[1:], tables)
-
-		// With an invalid filter.
-		tables, err = s.Tables(ctx, qual, invalidTableID)
-		assert.NoError(t, err)
-		assert.Equal(t, exp[0:0], tables)
-
-		// With both valid and invalid filters.
-		tables, err = s.Tables(ctx, qual, qtbl0.ID, invalidTableID)
-		assert.NoError(t, err)
-		assert.Equal(t, exp[1:], tables)
-
-		// With all valid filters.
-		tables, err = s.Tables(ctx, qual, qtbl0.ID, qtbl1.ID)
-		assert.NoError(t, err)
-		assert.Equal(t, exp, tables)
-	})
-}

dax/node.go

@@ -3,6 +3,7 @@ package dax
 import (
 	"context"
 	"fmt"
+	"strings"
 
 	"github.com/featurebasedb/featurebase/v3/errors"
 )
@@ -15,7 +16,23 @@ type Node struct {
 	RoleTypes []RoleType `json:"role-types"`
 }
 
-// AssignedNode is used in API responses.
+// Nodes is a slice of *Node. It's useful for printing the nodes as a list of
+// node.Addresses via its String() method.
+type Nodes []*Node
+
+// String prints the slice of node addresses in Nodes.
+func (n Nodes) String() string {
+	out := make([]string, 0, len(n))
+	for i := range n {
+		out = append(out, string(n[i].Address))
+	}
+	return "[" + strings.Join(out, ",") + "]"
+}
+
+// AssignedNode represents a Worker which has been assigned a role. Note that
+// the worker which it represents might be responsible for multiple roles, but
+// AssignedNode only ever represents one of those roles at a time. This is
+// because it is always the response of a RoleType-specific request.
 type AssignedNode struct {
 	Address Address `json:"address"`
 	Role    Role    `json:"role"`
@@ -23,10 +40,10 @@ type AssignedNode struct {
 
 // NodeService represents a service for managing Nodes.
 type NodeService interface {
-	CreateNode(context.Context, Address, *Node) error
-	ReadNode(context.Context, Address) (*Node, error)
-	DeleteNode(context.Context, Address) error
-	Nodes(context.Context) ([]*Node, error)
+	CreateNode(Transaction, Address, *Node) error
+	ReadNode(Transaction, Address) (*Node, error)
+	DeleteNode(Transaction, Address) error
+	Nodes(Transaction) ([]*Node, error)
 }
 
 // ComputeNode represents a compute node and the table/shards for which it is

dax/queryer/client/client.go

@@ -49,12 +49,12 @@ func (c *Client) Health() bool {
 	return true
 }
 
-func (c *Client) QuerySQL(ctx context.Context, qual dax.TableQualifier, sql string) (*featurebase.WireQueryResponse, error) {
+func (c *Client) QuerySQL(ctx context.Context, qdbid dax.QualifiedDatabaseID, sql string) (*featurebase.WireQueryResponse, error) {
 	url := fmt.Sprintf("%s/sql", c.address.WithScheme(defaultScheme))
 
 	req := &queryerhttp.SQLRequest{
-		OrganizationID: qual.OrganizationID,
-		DatabaseID:     qual.DatabaseID,
+		OrganizationID: qdbid.OrganizationID,
+		DatabaseID:     qdbid.DatabaseID,
 		SQL:            sql,
 	}
 
@@ -85,12 +85,13 @@ func (c *Client) QuerySQL(ctx context.Context, qual dax.TableQualifier, sql stri
 
 	return wireResp, nil
 }
-func (c *Client) QueryPQL(ctx context.Context, qual dax.TableQualifier, table dax.TableName, pql string) (*featurebase.WireQueryResponse, error) {
+
+func (c *Client) QueryPQL(ctx context.Context, qdbid dax.QualifiedDatabaseID, table dax.TableName, pql string) (*featurebase.WireQueryResponse, error) {
 	url := fmt.Sprintf("%s/query", c.address.WithScheme(defaultScheme))
 
 	req := &queryerhttp.QueryRequest{
-		OrganizationID: qual.OrganizationID,
-		DatabaseID:     qual.DatabaseID,
+		OrganizationID: qdbid.OrganizationID,
+		DatabaseID:     qdbid.DatabaseID,
 		Table:          table,
 		PQL:            pql,
 	}

dax/queryer/http/handler.go

@@ -49,15 +49,15 @@ func (s *server) postQuery(w http.ResponseWriter, r *http.Request) {
 
 	var resp interface{}
 	var err error
-	qual := dax.NewTableQualifier(req.OrganizationID, req.DatabaseID)
+	qdbid := dax.NewQualifiedDatabaseID(req.OrganizationID, req.DatabaseID)
 	if req.SQL != "" {
-		resp, err = s.queryer.QuerySQL(ctx, qual, req.SQL)
+		resp, err = s.queryer.QuerySQL(ctx, qdbid, req.SQL)
 		if err != nil {
 			http.Error(w, err.Error(), http.StatusBadRequest)
 			return
 		}
 	} else {
-		resp, err = s.queryer.QueryPQL(ctx, qual, req.Table, req.PQL)
+		resp, err = s.queryer.QueryPQL(ctx, qdbid, req.Table, req.PQL)
 		if err != nil {
 			http.Error(w, err.Error(), http.StatusBadRequest)
 			return
@@ -83,8 +83,8 @@ func (s *server) postSQL(w http.ResponseWriter, r *http.Request) {
 
 	ctx := r.Context()
 
-	qual := dax.NewTableQualifier(req.OrganizationID, req.DatabaseID)
-	resp, err := s.queryer.QuerySQL(ctx, qual, req.SQL)
+	qdbid := dax.NewQualifiedDatabaseID(req.OrganizationID, req.DatabaseID)
+	resp, err := s.queryer.QuerySQL(ctx, qdbid, req.SQL)
 	if err != nil {
 		http.Error(w, err.Error(), http.StatusBadRequest)
 		return

dax/queryer/orchestrator.go

@@ -1840,6 +1840,7 @@ func (o *orchestrator) executeExtract(ctx context.Context, tableKeyer dax.TableK
 		}
 		fields[i] = fieldName
 	}
+	// TODO(tlt): is `fields` used?
 
 	// Merge returned results at coordinating node.
 	reduceFn := func(ctx context.Context, prev, v interface{}) interface{} {
@@ -3511,13 +3512,13 @@ func callArgString(call *pql.Call, key string) string {
 
 type qualifiedOrchestrator struct {
 	*orchestrator
-	qual dax.TableQualifier
+	qdbid dax.QualifiedDatabaseID
 }
 
-func newQualifiedOrchestrator(orch *orchestrator, qual dax.TableQualifier) *qualifiedOrchestrator {
+func newQualifiedOrchestrator(orch *orchestrator, qdbid dax.QualifiedDatabaseID) *qualifiedOrchestrator {
 	return &qualifiedOrchestrator{
 		orchestrator: orch,
-		qual:         qual,
+		qdbid:        qdbid,
 	}
 }
 
@@ -3528,7 +3529,7 @@ func (o *qualifiedOrchestrator) Execute(ctx context.Context, tableKeyer dax.Tabl
 
 	switch keyer := tableKeyer.(type) {
 	case *dax.Table:
-		qtbl = dax.NewQualifiedTable(o.qual, keyer)
+		qtbl = dax.NewQualifiedTable(o.qdbid, keyer)
 	case *dax.QualifiedTable:
 		qtbl = keyer
 	default:

dax/queryer/queryer.go

@@ -32,7 +32,7 @@ import (
 // of "Queryer" nodes, which handle incoming query requests.
 type Queryer struct {
 	mu            sync.RWMutex
-	orchestrators map[dax.TableQualifier]*qualifiedOrchestrator
+	orchestrators map[dax.QualifiedDatabaseID]*qualifiedOrchestrator
 
 	fbClient *featurebase.InternalClient
 
@@ -47,7 +47,7 @@ func New(cfg Config) *Queryer {
 	q := &Queryer{
 		noder:         dax.NewNopNoder(),
 		schemar:       dax.NewNopSchemar(),
-		orchestrators: make(map[dax.TableQualifier]*qualifiedOrchestrator),
+		orchestrators: make(map[dax.QualifiedDatabaseID]*qualifiedOrchestrator),
 		logger:        logger.NopLogger,
 	}
 
@@ -59,13 +59,13 @@ func New(cfg Config) *Queryer {
 }
 
 // Orchestrator gets (or creates) an instance of qualifiedOrchestrator based on
-// the provided dax.TableQualifier.
-func (q *Queryer) Orchestrator(qual dax.TableQualifier) *qualifiedOrchestrator {
+// the provided dax.QualifiedDatabaseID.
+func (q *Queryer) Orchestrator(qdbid dax.QualifiedDatabaseID) *qualifiedOrchestrator {
 	// Try to get orchestrator under a read lock first.
 	if orch := func() *qualifiedOrchestrator {
 		q.mu.RLock()
 		defer q.mu.RUnlock()
-		if orch, ok := q.orchestrators[qual]; ok {
+		if orch, ok := q.orchestrators[qdbid]; ok {
 			return orch
 		}
 		return nil
@@ -77,11 +77,11 @@ func (q *Queryer) Orchestrator(qual dax.TableQualifier) *qualifiedOrchestrator {
 	// lock and try a read/write.
 	q.mu.Lock()
 	defer q.mu.Unlock()
-	if orch, ok := q.orchestrators[qual]; ok {
+	if orch, ok := q.orchestrators[qdbid]; ok {
 		return orch
 	}
 
-	sapi := newQualifiedSchemaAPI(qual, q.schemar)
+	sapi := newQualifiedSchemaAPI(qdbid, q.schemar)
 
 	orch := &orchestrator{
 		schema:   sapi,
@@ -92,8 +92,8 @@ func (q *Queryer) Orchestrator(qual dax.TableQualifier) *qualifiedOrchestrator {
 		logger: q.logger,
 	}
 
-	qorch := newQualifiedOrchestrator(orch, qual)
-	q.orchestrators[qual] = qorch
+	qorch := newQualifiedOrchestrator(orch, qdbid)
+	q.orchestrators[qdbid] = qorch
 
 	return qorch
 }
@@ -132,12 +132,9 @@ func (q *Queryer) Start() error {
 	return nil
 }
 
-func (q *Queryer) QuerySQL(ctx context.Context, qual dax.TableQualifier, sql string) (*featurebase.WireQueryResponse, error) {
+func (q *Queryer) QuerySQL(ctx context.Context, qdbid dax.QualifiedDatabaseID, sql string) (*featurebase.WireQueryResponse, error) {
 	start := time.Now()
 
-	if len(sql) > 0 && sql[0] == '[' {
-		return q.parseAndQueryPQL(ctx, qual, sql)
-	}
 	ret := &featurebase.WireQueryResponse{}
 
 	applyExecutionTime := func() {
@@ -149,6 +146,19 @@ func (q *Queryer) QuerySQL(ctx context.Context, qual dax.TableQualifier, sql str
 		applyExecutionTime()
 	}
 
+	// If PQL, run that instead.
+	if len(sql) > 0 && sql[0] == '[' {
+		if pqlResp, err := q.parseAndQueryPQL(ctx, qdbid, sql); err != nil {
+			applyError(errors.Wrap(err, "querying pql"))
+			return ret, nil
+		} else {
+			ret = pqlResp
+		}
+		applyExecutionTime()
+
+		return ret, nil
+	}
+
 	// Create a requestID and add it to the context.
 	requestID, err := uuid.NewV4()
 	if err != nil {
@@ -165,19 +175,17 @@ func (q *Queryer) QuerySQL(ctx context.Context, qual dax.TableQualifier, sql str
 	}
 
 	// SchemaAPI
-	sapi := newQualifiedSchemaAPI(qual, q.schemar)
+	sapi := newQualifiedSchemaAPI(qdbid, q.schemar)
 
 	// Importer
-	imp := idkmds.NewImporter(q.noder, q.schemar, qual, nil)
+	imp := idkmds.NewImporter(q.noder, q.schemar, qdbid, nil)
 
-	// TODO(tlt): this obviously doesn't work; we don't have an API here. We
-	// need a dax-compatible implementation of the SystemAPI (or at least a
-	// no-op implementation).
-	sysapi := &featurebase.FeatureBaseSystemAPI{API: nil}
+	// TODO(tlt): We need a dax-compatible implementation of the SystemAPI.
+	sysapi := &featurebase.NopSystemAPI{}
 
 	systemLayer := systemlayer.NewSystemLayer()
 
-	pl := planner.NewExecutionPlanner(q.Orchestrator(qual), sapi, sysapi, systemLayer, imp, q.logger, sql)
+	pl := planner.NewExecutionPlanner(q.Orchestrator(qdbid), sapi, sysapi, systemLayer, imp, q.logger, sql)
 
 	planOp, err := pl.CompilePlan(ctx, st)
 	if err != nil {
@@ -229,7 +237,7 @@ func (q *Queryer) QuerySQL(ctx context.Context, qual dax.TableQualifier, sql str
 	return ret, nil
 }
 
-func (q *Queryer) parseAndQueryPQL(ctx context.Context, qual dax.TableQualifier, sql string) (*featurebase.WireQueryResponse, error) {
+func (q *Queryer) parseAndQueryPQL(ctx context.Context, qdbid dax.QualifiedDatabaseID, sql string) (*featurebase.WireQueryResponse, error) {
 	var i int
 	for i = 1; sql[i] != ']'; i++ {
 		if i == len(sql)-1 {
@@ -240,7 +248,7 @@ func (q *Queryer) parseAndQueryPQL(ctx context.Context, qual dax.TableQualifier,
 	query := sql[i+1:]
 	fmt.Println("got table/query", table, query)
 
-	return q.QueryPQL(ctx, qual, dax.TableName(table), query)
+	return q.queryPQL(ctx, qdbid, dax.TableName(table), query)
 }
 
 // convertIndex tries to covert any "index" specified in the call.Args map to a
@@ -250,9 +258,9 @@ func (q *Queryer) parseAndQueryPQL(ctx context.Context, qual dax.TableQualifier,
 // modify Call.CallIndex() to be TableKeyer aware. I didn't do that along with
 // these changes because I'm not sure if we want to introduce dax types into the
 // pql package.
-func (q *Queryer) convertIndex(ctx context.Context, qual dax.TableQualifier, call *featurebase_pql.Call) {
+func (q *Queryer) convertIndex(ctx context.Context, qdbid dax.QualifiedDatabaseID, call *featurebase_pql.Call) {
 	if index := call.CallIndex(); index != "" {
-		qtbl, err := q.schemar.TableByName(ctx, qual, dax.TableName(index))
+		qtbl, err := q.schemar.TableByName(ctx, qdbid, dax.TableName(index))
 		if err != nil {
 			return
 		}
@@ -261,11 +269,37 @@ func (q *Queryer) convertIndex(ctx context.Context, qual dax.TableQualifier, cal
 
 	// Apply to children.
 	for _, child := range call.Children {
-		q.convertIndex(ctx, qual, child)
+		q.convertIndex(ctx, qdbid, child)
 	}
 }
 
-func (q *Queryer) QueryPQL(ctx context.Context, qual dax.TableQualifier, table dax.TableName, pql string) (*featurebase.WireQueryResponse, error) {
+func (q *Queryer) QueryPQL(ctx context.Context, qdbid dax.QualifiedDatabaseID, table dax.TableName, pql string) (*featurebase.WireQueryResponse, error) {
+	start := time.Now()
+
+	ret := &featurebase.WireQueryResponse{}
+
+	applyExecutionTime := func() {
+		ret.ExecutionTime = time.Since(start).Microseconds()
+	}
+
+	applyError := func(e error) {
+		ret.Error = e.Error()
+		applyExecutionTime()
+	}
+
+	if pqlResp, err := q.queryPQL(ctx, qdbid, table, pql); err != nil {
+		applyError(errors.Wrap(err, "querying pql"))
+		return ret, nil
+	} else {
+		ret = pqlResp
+	}
+
+	applyExecutionTime()
+
+	return ret, nil
+}
+
+func (q *Queryer) queryPQL(ctx context.Context, qdbid dax.QualifiedDatabaseID, table dax.TableName, pql string) (*featurebase.WireQueryResponse, error) {
 	// Parse the pql into a pql.Query containing []pql.Call.
 	qry, err := featurebase_pql.NewParser(strings.NewReader(pql)).Parse()
 	if err != nil {
@@ -276,14 +310,14 @@ func (q *Queryer) QueryPQL(ctx context.Context, qual dax.TableQualifier, table d
 	}
 
 	// Replace any "index" arguments within the PQL with a TableKey.
-	q.convertIndex(ctx, qual, qry.Calls[0])
+	q.convertIndex(ctx, qdbid, qry.Calls[0])
 
-	qtbl, err := q.schemar.TableByName(ctx, qual, dax.TableName(table))
+	qtbl, err := q.schemar.TableByName(ctx, qdbid, dax.TableName(table))
 	if err != nil {
 		return nil, errors.Wrap(err, "converting index to qualified table")
 	}
 
-	results, err := q.Orchestrator(qual).Execute(ctx, qtbl, qry, nil, &featurebase.ExecOptions{})
+	results, err := q.Orchestrator(qdbid).Execute(ctx, qtbl, qry, nil, &featurebase.ExecOptions{})
 	if err != nil {
 		return nil, errors.Wrap(err, "orchestrator.Execute")
 	}

dax/queryer/schema_api.go

@@ -12,31 +12,31 @@ import (
 var _ pilosa.SchemaAPI = (*qualifiedSchemaAPI)(nil)
 
 // qualifiedSchemaAPI is a wrapper around schemaAPI. It is initialized with a
-// TableQualifier, and it uses this qualifer to convert between, for example,
-// FeatureBase index name (a string) and TableKey. It requires a Schemar to do
-// that lookup/conversion.
+// QualifiedDatabaseID, and it uses this qualifer to convert between, for
+// example, FeatureBase index name (a string) and TableKey. It requires a
+// Schemar to do that lookup/conversion.
 type qualifiedSchemaAPI struct {
-	qual    dax.TableQualifier
+	qdbid   dax.QualifiedDatabaseID
 	schemar dax.Schemar
 }
 
-func newQualifiedSchemaAPI(qual dax.TableQualifier, schema dax.Schemar) *qualifiedSchemaAPI {
+func newQualifiedSchemaAPI(qdbid dax.QualifiedDatabaseID, schema dax.Schemar) *qualifiedSchemaAPI {
 	return &qualifiedSchemaAPI{
-		qual:    qual,
+		qdbid:   qdbid,
 		schemar: schema,
 	}
 }
 
 func (s *qualifiedSchemaAPI) TableByName(ctx context.Context, tname dax.TableName) (*dax.Table, error) {
-	qtbl, err := s.schemar.TableByName(ctx, s.qual, tname)
+	qtbl, err := s.schemar.TableByName(ctx, s.qdbid, tname)
 	if err != nil {
-		return nil, errors.Wrapf(err, "getting table id: (%s) %s", s.qual, tname)
+		return nil, errors.Wrapf(err, "getting table id: (%s) %s", s.qdbid, tname)
 	}
 	return &qtbl.Table, nil
 }
 
 func (s *qualifiedSchemaAPI) TableByID(ctx context.Context, tid dax.TableID) (*dax.Table, error) {
-	qtid := dax.NewQualifiedTableID(s.qual, tid)
+	qtid := dax.NewQualifiedTableID(s.qdbid, tid)
 
 	qtbl, err := s.schemar.TableByID(ctx, qtid)
 	if err != nil {
@@ -47,7 +47,7 @@ func (s *qualifiedSchemaAPI) TableByID(ctx context.Context, tid dax.TableID) (*d
 }
 
 func (s *qualifiedSchemaAPI) Tables(ctx context.Context) ([]*dax.Table, error) {
-	qtbls, err := s.schemar.Tables(ctx, s.qual)
+	qtbls, err := s.schemar.Tables(ctx, s.qdbid)
 	if err != nil {
 		return nil, errors.Wrap(err, "getting tables")
 	}
@@ -61,32 +61,32 @@ func (s *qualifiedSchemaAPI) Tables(ctx context.Context) ([]*dax.Table, error) {
 }
 
 func (s *qualifiedSchemaAPI) CreateTable(ctx context.Context, tbl *dax.Table) error {
-	qtbl := dax.NewQualifiedTable(s.qual, tbl)
+	qtbl := dax.NewQualifiedTable(s.qdbid, tbl)
 	return s.schemar.CreateTable(ctx, qtbl)
 }
 
 func (s *qualifiedSchemaAPI) CreateField(ctx context.Context, tname dax.TableName, fld *dax.Field) error {
-	qtbl, err := s.schemar.TableByName(ctx, s.qual, tname)
+	qtbl, err := s.schemar.TableByName(ctx, s.qdbid, tname)
 	if err != nil {
-		return errors.Wrapf(err, "getting table by name: (%s) %s", s.qual, tname)
+		return errors.Wrapf(err, "getting table by name: (%s) %s", s.qdbid, tname)
 	}
 
 	return s.schemar.CreateField(ctx, qtbl.QualifiedID(), fld)
 }
 
 func (s *qualifiedSchemaAPI) DeleteTable(ctx context.Context, tname dax.TableName) error {
-	qtbl, err := s.schemar.TableByName(ctx, s.qual, tname)
+	qtbl, err := s.schemar.TableByName(ctx, s.qdbid, tname)
 	if err != nil {
-		return errors.Wrapf(err, "getting table by name: (%s) %s", s.qual, tname)
+		return errors.Wrapf(err, "getting table by name: (%s) %s", s.qdbid, tname)
 	}
 
 	return s.schemar.DropTable(ctx, qtbl.QualifiedID())
 }
 
 func (s *qualifiedSchemaAPI) DeleteField(ctx context.Context, tname dax.TableName, fname dax.FieldName) error {
-	qtid, err := s.schemar.TableByName(ctx, s.qual, tname)
+	qtid, err := s.schemar.TableByName(ctx, s.qdbid, tname)
 	if err != nil {
-		return errors.Wrapf(err, "getting table by name: (%s) %s", s.qual, tname)
+		return errors.Wrapf(err, "getting table by name: (%s) %s", s.qdbid, tname)
 	}
 
 	return s.schemar.DropField(ctx, qtid.Key().QualifiedTableID(), fname)

dax/role.go

@@ -28,8 +28,8 @@ type Role interface {
 }
 
 // Ensure type implements interface.
-var _ Role = &ComputeRole{}
-var _ Role = &TranslateRole{}
+var _ Role = (*ComputeRole)(nil)
+var _ Role = (*TranslateRole)(nil)
 
 // ComputeRole is a role specific to compute nodes.
 type ComputeRole struct {
@@ -37,7 +37,7 @@ type ComputeRole struct {
 	Shards   ShardNums `json:"shards"`
 }
 
-// Type returns the type for ComputeRole. This is mainly to impolement the Role
+// Type returns the type for ComputeRole. This is mainly to implement the Role
 // interface.
 func (cr *ComputeRole) Type() RoleType {
 	return RoleTypeCompute
@@ -50,7 +50,7 @@ type TranslateRole struct {
 	Fields     []FieldName   `json:"fields"`
 }
 
-// Type returns the type for TransteRole. This is mainly to impolement the Role
+// Type returns the type for TranslateRole. This is mainly to implement the Role
 // interface.
 func (cr *TranslateRole) Type() RoleType {
 	return RoleTypeTranslate

dax/schema.go

@@ -3,16 +3,52 @@ package dax
 import "context"
 
 // Schemar is similar to the pilosa.SchemaAPI interface, but it takes
-// TableQualifiers into account.
+// QualifiedDatabaseIDs into account. Note that it is also similar to the
+// schemar.Schemar interface, but that is used internally, typically within the
+// Controller, and it takes Transactions rather than a Context, because its
+// methods are assumed to be used as part of a larger request.
+// TODO(tlt): clean up the mds/controller/schemar Schemar interface confusion.
 type Schemar interface {
-	TableByName(ctx context.Context, qual TableQualifier, tname TableName) (*QualifiedTable, error)
-	TableByID(ctx context.Context, qtid QualifiedTableID) (*QualifiedTable, error)
-	Tables(ctx context.Context, qual TableQualifier, tids ...TableID) ([]*QualifiedTable, error)
+	//////////////////////////////////////////////////////////////////////////
+	// Database methods
+	//////////////////////////////////////////////////////////////////////////
+
+	CreateDatabase(context.Context, *QualifiedDatabase) error
+	// DropDatabase(context.Context, QualifiedDatabaseID) error
+
+	// DatabaseByName(ctx context.Context, orgID OrganizationID, dbname DatabaseName) (*QualifiedDatabase, error)
+	DatabaseByID(ctx context.Context, qdbid QualifiedDatabaseID) (*QualifiedDatabase, error)
+
+	// // Databases returns a list of databases. If the OrganizationID is empty,
+	// // all databases will be returned. If greater than zero database IDs are
+	// // passed in the second argument, only databases matching those IDs will be
+	// // returned.
+	// Databases(context.Context, OrganizationID, ...DatabaseID) ([]*QualifiedDatabase, error)
+
+	// SetDatabaseOptions(context.Context, QualifiedDatabaseID, DatabaseOptions) error
+
+	//////////////////////////////////////////////////////////////////////////
+	// Table methods
+	//////////////////////////////////////////////////////////////////////////
 
 	CreateTable(ctx context.Context, qtbl *QualifiedTable) error
-	CreateField(ctx context.Context, qtid QualifiedTableID, fld *Field) error
-
 	DropTable(ctx context.Context, qtid QualifiedTableID) error
+
+	TableByName(ctx context.Context, qdbid QualifiedDatabaseID, tname TableName) (*QualifiedTable, error)
+	TableByID(ctx context.Context, qtid QualifiedTableID) (*QualifiedTable, error)
+
+	// Tables returns a list of tables. If the qualifiers DatabaseID is empty,
+	// all tables in the org will be returned. If the OrganizationID is empty,
+	// all tables will be returned. If both are populated, only tables in that
+	// database will be returned. If greater than zero table IDs are passed in
+	// the third argument, only tables matching those IDs will be returned.
+	Tables(ctx context.Context, qdbid QualifiedDatabaseID, tids ...TableID) ([]*QualifiedTable, error)
+
+	//////////////////////////////////////////////////////////////////////////
+	// Field methods
+	//////////////////////////////////////////////////////////////////////////
+
+	CreateField(ctx context.Context, qtid QualifiedTableID, fld *Field) error
 	DropField(ctx context.Context, qtid QualifiedTableID, fname FieldName) error
 }
 
@@ -28,13 +64,19 @@ func NewNopSchemar() *NopSchemar {
 	return &NopSchemar{}
 }
 
-func (s *NopSchemar) TableByName(context.Context, TableQualifier, TableName) (*QualifiedTable, error) {
+func (s *NopSchemar) CreateDatabase(context.Context, *QualifiedDatabase) error {
+	return nil
+}
+func (s *NopSchemar) DatabaseByID(ctx context.Context, qdbid QualifiedDatabaseID) (*QualifiedDatabase, error) {
+	return nil, nil
+}
+func (s *NopSchemar) TableByName(context.Context, QualifiedDatabaseID, TableName) (*QualifiedTable, error) {
 	return nil, nil
 }
 func (s *NopSchemar) TableByID(ctx context.Context, qtid QualifiedTableID) (*QualifiedTable, error) {
 	return nil, nil
 }
-func (s *NopSchemar) Tables(ctx context.Context, qual TableQualifier, tids ...TableID) ([]*QualifiedTable, error) {
+func (s *NopSchemar) Tables(ctx context.Context, qdbid QualifiedDatabaseID, tids ...TableID) ([]*QualifiedTable, error) {
 	return nil, nil
 }
 func (s *NopSchemar) CreateTable(ctx context.Context, qtbl *QualifiedTable) error {

dax/storage/storage_test.go

@@ -34,10 +34,10 @@ func TestResourceManager(t *testing.T) {
 	mm := NewResourceManager(sn, wl, logger.NewStandardLogger(os.Stderr))
 
 	qtid := dax.QualifiedTableID{
-		TableQualifier: dax.TableQualifier{
-			OrganizationID: dax.OrganizationID("org1"),
-			DatabaseID:     dax.DatabaseID("db1"),
-		},
+		QualifiedDatabaseID: dax.NewQualifiedDatabaseID(
+			dax.OrganizationID("org1"),
+			dax.DatabaseID("db1"),
+		),
 		ID:   dax.TableID("blah"),
 		Name: "blah",
 	}

dax/table.go

@@ -22,14 +22,17 @@ import (
 //
 // OrganizationID - carried over from ControlPlane; currently uuid
 // DatabaseID - carried over from ControlPlane; currently uuid
+// Database - base Database struct
+// DatabaseKey - a string representation of OrganizationID and DatabaseID
 // TableID - internally stored as a uint64; presented as a hex string.
 // TableName - human-friendly string table name
 // Table - base Table struct; includes a TableID and a TableName
-// TableQualifier - combination of OrganizationID and DatabaseID
-// QualifiedTable - TableQualifier plus a Table
-// QualifiedTableID - TableQualifier plus a TableID
+// QualifiedDatabase - OrganizationID plus a Database
+// QualifiedDatabaseID - combination of OrganizationID and DatabaseID
+// QualifiedTable - QualifiedDatabaseID plus a Table
+// QualifiedTableID - QualifiedDatabaseID plus a TableID
 // TableKey - a string representation of OrganizationID, DatabaseID, and
-// TableID, which is safe to use as a FeatureBase index name.
+//            TableID, which is safe to use as a FeatureBase index name.
 //
 // Example:
 // OrganizationID - "29-ae44-41"
@@ -37,7 +40,7 @@ import (
 // TableID - 123456789 (hex string: "499602d2")
 // TableName - foo
 // Table - {ID:"499602d2", Name: "foo", Fields: ... }
-// TableQualifier - {Org: "29-ae44-41", DB: "75-d1a2-4f"}
+// QualifierDatabaseID - {Org: "29-ae44-41", DB: "75-d1a2-4f"}
 // QualifiedTable - {Org: "29-ae44-41", DB: "75-d1a2-4f", Table: *tbl}
 // QualifiedTableID - {Org: "29-ae44-41", DB: "75-d1a2-4f", TableID: "499602d2"}
 // TableKey - "tbl__29-ae44-41__75-d1a2-4f__499602d2"
@@ -52,6 +55,13 @@ import (
 // underscore.
 const TableKeyDelimiter = "__"
 
+// PrefixDatabase is used as a prefix to DatabaseKey strings because FeatureBase
+// indexes must start with an alpha (a-z) character. Because the string
+// representation of a uuid (i.e. the OrganizationID value) can start with a
+// numeric value, we can't have OrganizationId (or any of the other ID values
+// which make up the DatabaseKey) be at the beginning of the DatabaseKey.
+const PrefixDatabase = "db"
+
 // PrefixTable is used as a prefix to TableKey strings because FeatureBase
 // indexes must start with an alpha (a-z) character. Because the string
 // representation of a uuid (i.e. the OrganizationID value) can start with a
@@ -103,6 +113,124 @@ type OrganizationID string
 // value could be any string.
 type DatabaseID string
 
+// DatabaseKey is a globally unique identifier for a database; it is effectively the
+// compound key: (org, database). This is (hopefully) the value that will
+// be used when interfacing with services which are unaware of qualifiers.
+type DatabaseKey string
+
+// QualifiedDatabaseID returns the QualifiedDatabaseID based on the key. If
+// DatabaseKey can't be parsed into a valid (i.e. complete) QualifiedDatabaseID,
+// then blank values are used where necessary.
+func (dk DatabaseKey) QualifiedDatabaseID() QualifiedDatabaseID {
+	qdbid, err := QualifiedDatabaseIDFromKey(string(dk))
+	if err != nil {
+		return NewQualifiedDatabaseID("", DatabaseID(dk))
+	}
+	return qdbid
+}
+
+// DatabaseName is a human-friendly string.
+type DatabaseName string
+
+// Database represents a database and its configuration.
+type Database struct {
+	ID      DatabaseID      `json:"id"`
+	Name    DatabaseName    `json:"name"`
+	Options DatabaseOptions `json:"options"`
+	// Tables  []*Table        `json:"tables"`
+
+	Description string `json:"description,omitempty"`
+	Owner       string `json:"owner,omitempty"`
+	CreatedAt   int64  `json:"createdAt,omitempty"`
+	UpdatedAt   int64  `json:"updatedAt,omitempty"`
+	UpdatedBy   string `json:"updatedBy,omitempty"`
+}
+
+// DatabaseOptions are used to configure a database.
+type DatabaseOptions struct {
+	WorkersMin int `json:"workers-min"`
+	WorkersMax int `json:"workers-max"`
+}
+
+// QualifiedDatabase is a Database along with its OrganizationID.
+type QualifiedDatabase struct {
+	OrganizationID OrganizationID `json:"org-id"`
+	Database
+}
+
+type QualifiedDatabases []*QualifiedDatabase
+
+// Key returns the string-encoded (delimited by DatabaseKeyDelimiter) globally
+// unique DatabaseKey.
+func (qdb QualifiedDatabase) Key() DatabaseKey {
+	return qdb.QualifiedID().Key()
+}
+
+// String returns a human-friendly version of the QualifiedDatabase. It is only
+// used for display purposes; it is not used as any kind of key.
+func (qdb QualifiedDatabase) String() string {
+	return fmt.Sprintf("%s (%s)", qdb.QualifiedID(), qdb.Name)
+}
+
+// QualifiedID returns the QualifiedDatabaseID for the database.
+func (qdb *QualifiedDatabase) QualifiedID() QualifiedDatabaseID {
+	return QualifiedDatabaseID{
+		OrganizationID: qdb.OrganizationID,
+		DatabaseID:     qdb.ID,
+	}
+}
+
+// QualifiedDatabaseID is a DatabaseID along with its OrganizationID.
+type QualifiedDatabaseID struct {
+	OrganizationID OrganizationID `json:"org-id"`
+	DatabaseID     DatabaseID     `json:"db-id"`
+}
+
+// NewQualifiedDatabaseID is a helper function used to create a
+// QualifiedDatabaseID from the provided arguments.
+func NewQualifiedDatabaseID(orgID OrganizationID, dbID DatabaseID) QualifiedDatabaseID {
+	return QualifiedDatabaseID{
+		OrganizationID: orgID,
+		DatabaseID:     dbID,
+	}
+}
+
+// String returns a human-friendly version of the QualifiedDatabaseID. It is only
+// used for display purposes; it is not used as any kind of key. For that, see
+// the QualifiedDatabaseID.Key() method and the DatabaseKey type.
+func (qdbid QualifiedDatabaseID) String() string {
+	return fmt.Sprintf("[%s:%s]", qdbid.OrganizationID, qdbid.DatabaseID)
+}
+
+// Key returns the string-encoded (delimited by TableKeyDelimiter) globally
+// unique DatabaseKey. The key has a prefix because FeatureBase index name
+// restrictions require the name to start with a non-numeric value, and since a
+// uuid can contain a number as its first character, we have to prefix it with
+// something.
+func (qdbid QualifiedDatabaseID) Key() DatabaseKey {
+	if qdbid.DatabaseID == "" {
+		panic("QualifiedDatabaseID.Key called without an ID set")
+	}
+	return DatabaseKey(fmt.Sprintf("%s%s%s%s%s",
+		PrefixDatabase,
+		TableKeyDelimiter,
+		qdbid.OrganizationID,
+		TableKeyDelimiter,
+		qdbid.DatabaseID))
+}
+
+// QualifiedDatabaseIDs is a list of QualifiedDatabaseID.
+type QualifiedDatabaseIDs []QualifiedDatabaseID
+
+func (s QualifiedDatabaseIDs) Len() int { return len(s) }
+func (s QualifiedDatabaseIDs) Less(i, j int) bool {
+	if s[i].OrganizationID != s[j].OrganizationID {
+		return s[i].OrganizationID < s[j].OrganizationID
+	}
+	return s[i].DatabaseID < s[j].DatabaseID
+}
+func (s QualifiedDatabaseIDs) Swap(i, j int) { s[i], s[j] = s[j], s[i] }
+
 // TableKeyer is an interface implemented by any type which can produce, and be
 // represented by, a TableKey. In the case of a QualifiedTable, its TableKey
 // might be something like `tbl__org__db__tableid`, while a general pilosa
@@ -129,7 +257,7 @@ func (s StringTableKeyer) Key() TableKey {
 // TableKey as the value for index.Name.
 type TableKey string
 
-func (t TableKey) Key() TableKey { return t }
+func (tk TableKey) Key() TableKey { return tk }
 
 // QualifiedTableID returns the QualifiedTableID based on the key. If TableKey
 // can't be parsed into a valid (i.e. complete) QualifiedTableID, then blank
@@ -138,7 +266,7 @@ func (tk TableKey) QualifiedTableID() QualifiedTableID {
 	qtid, err := QualifiedTableIDFromKey(string(tk))
 	if err != nil {
 		return NewQualifiedTableID(
-			NewTableQualifier("", ""),
+			NewQualifiedDatabaseID("", ""),
 			TableID(tk),
 		)
 	}
@@ -153,7 +281,7 @@ func (s TableKeys) Less(i, j int) bool { return s[i] < s[j] }
 func (s TableKeys) Swap(i, j int)      { s[i], s[j] = s[j], s[i] }
 
 // TableID is a table identifier. It is unique within the scope of a
-// TableQualifier. Coupled with a TableQualifier, it makes up a
+// QualifiedDatabaseID. Coupled with a QualifiedDatabaseID, it makes up a
 // QualifiedTableID and, when encoded as a string, a TableKey.
 type TableID string
 
@@ -165,7 +293,7 @@ func (s TableIDs) Less(i, j int) bool { return s[i] < s[j] }
 func (s TableIDs) Swap(i, j int)      { s[i], s[j] = s[j], s[i] }
 
 // TableName is a human-friendly string. While it is not used as a primary key,
-// uniqueness is generally enforced within the scope of a TableQualifier.
+// uniqueness is generally enforced within the scope of a QualifiedDatabaseID.
 type TableName string
 
 // TableNames is a sortable slice of TableName.
@@ -224,7 +352,7 @@ func (t *Table) CreateID() (TableID, error) {
 }
 
 // NewTable returns a new instance of table with a pseudo-random ID which is
-// assumed to be unique within the scope of a TableQualifier.
+// assumed to be unique within the scope of a QualifiedDatabaseID.
 func NewTable(name TableName) *Table {
 	return &Table{
 		Name:   name,
@@ -315,84 +443,6 @@ func (o Tables) Len() int           { return len(o) }
 func (o Tables) Less(i, j int) bool { return o[i].Name < o[j].Name }
 func (o Tables) Swap(i, j int)      { o[i], o[j] = o[j], o[i] }
 
-// TableQualifierKey is the unique TableQualifier values encoded as a string. The
-// current encoding is delimited as `prefix|OrganizationID|DatabaseID` (where
-// the pipe may be some other delimiter) by the TableQualifier.Key() method.
-type TableQualifierKey string
-
-// Qualifier returns the Qualifier based on the values encoded into the
-// TableQualifierKey string.
-func (tqk TableQualifierKey) Qualifier() TableQualifier {
-	parts := strings.Split(string(tqk), TableKeyDelimiter)
-
-	if len(parts) < 3 {
-		return NewTableQualifier("", "")
-	}
-
-	return NewTableQualifier(
-		OrganizationID(parts[1]),
-		DatabaseID(parts[2]),
-	)
-}
-
-// OrganizationID returns the OrganizationID value that has been encoded into
-// the TableQualifierKey string.
-func (tqk TableQualifierKey) OrganizationID() OrganizationID {
-	parts := strings.Split(string(tqk), TableKeyDelimiter)
-
-	if len(parts) < 2 {
-		return ""
-	}
-
-	return OrganizationID(parts[1])
-}
-
-// DatabaseID returns the DatabaseID value that has been encoded into the
-// TableQualifierKey string.
-func (tqk TableQualifierKey) DatabaseID() DatabaseID {
-	parts := strings.Split(string(tqk), TableKeyDelimiter)
-
-	if len(parts) < 3 {
-		return ""
-	}
-
-	return DatabaseID(parts[2])
-}
-
-// TableQualifier contains all the elements required to fully qualify a table.
-type TableQualifier struct {
-	OrganizationID OrganizationID `json:"org-id"`
-	DatabaseID     DatabaseID     `json:"db-id"`
-}
-
-// NewTableQualifier is a helper function used to create a TableQualifier from
-// the provided arguments.
-func NewTableQualifier(orgID OrganizationID, dbID DatabaseID) TableQualifier {
-	return TableQualifier{
-		OrganizationID: orgID,
-		DatabaseID:     dbID,
-	}
-}
-
-// String returns a human-friendly version of the TableQualifier. It is only
-// used for display purposes; it is not used as any kind of key. For that, see
-// the TableQualifier.Key() method and the TableQualifierKey type.
-func (tq TableQualifier) String() string {
-	return fmt.Sprintf("[%s:%s]", tq.OrganizationID, tq.DatabaseID)
-}
-
-// Key returns the string-encoded (delimited by TableKeyDelimiter)
-// TableQualifierKey.
-func (tq TableQualifier) Key() TableQualifierKey {
-	return TableQualifierKey(fmt.Sprintf("%s%s%s%s%s",
-		PrefixTable,
-		TableKeyDelimiter,
-		tq.OrganizationID,
-		TableKeyDelimiter,
-		tq.DatabaseID,
-	))
-}
-
 ////////////////////////////////////////////////
 
 // QualifiedTableID is a globally unique table identifier. It is a
@@ -400,17 +450,17 @@ func (tq TableQualifier) Key() TableQualifierKey {
 // portion). Most things will take a Name or an ID and do the right
 // thing™.
 type QualifiedTableID struct {
-	TableQualifier
+	QualifiedDatabaseID
 	ID   TableID   `json:"id"`
 	Name TableName `json:"name"`
 }
 
 // NewQualifiedTableID is a helper function used to create a QualifiedTableID
 // from the provided arguments.
-func NewQualifiedTableID(q TableQualifier, id TableID) QualifiedTableID {
+func NewQualifiedTableID(qdbid QualifiedDatabaseID, tid TableID) QualifiedTableID {
 	return QualifiedTableID{
-		TableQualifier: q,
-		ID:             id,
+		QualifiedDatabaseID: qdbid,
+		ID:                  tid,
 	}
 }
 
@@ -422,7 +472,7 @@ func QualifiedTableIDFromKey(key string) (QualifiedTableID, error) {
 	case 4:
 		// prefix|orgID|dbID|tblID
 		return NewQualifiedTableID(
-			NewTableQualifier(
+			NewQualifiedDatabaseID(
 				OrganizationID(parts[1]),
 				DatabaseID(parts[2]),
 			),
@@ -433,14 +483,31 @@ func QualifiedTableIDFromKey(key string) (QualifiedTableID, error) {
 	}
 }
 
-// String returns a human-friendly version of the TableQualifierID. It is only
-// used for display purposes; it is not used as any kind of key. For that, see
-// the TableQualifierID.Key() method.
+// QualifiedDatabaseIDFromKey decodes a string key into a QualifiedDatabaseID.
+// The key is assumed to have been encoded using the QualifiedDatabaseID.Key()
+// method.
+func QualifiedDatabaseIDFromKey(key string) (QualifiedDatabaseID, error) {
+	parts := strings.Split(key, TableKeyDelimiter)
+	switch len(parts) {
+	case 3:
+		// prefix|orgID|dbID
+		return NewQualifiedDatabaseID(
+			OrganizationID(parts[1]),
+			DatabaseID(parts[2]),
+		), nil
+	default:
+		return QualifiedDatabaseID{}, errors.Errorf("invalid key: %s", key)
+	}
+}
+
+// String returns a human-friendly version of the QualifiedDatabaseID. It is
+// only used for display purposes; it is not used as any kind of key. For that,
+// see the QualifiedDatabaseID.Key() method.
 func (qtid QualifiedTableID) String() string {
 	if qtid.ID == "" {
-		return fmt.Sprintf("%s%s", qtid.TableQualifier, qtid.Name)
+		return fmt.Sprintf("%s%s", qtid.QualifiedDatabaseID, qtid.Name)
 	}
-	return fmt.Sprintf("%s%s", qtid.TableQualifier, qtid.ID)
+	return fmt.Sprintf("%s%s", qtid.QualifiedDatabaseID, qtid.ID)
 }
 
 // Key returns the string-encoded (delimited by TableKeyDelimiter) globally
@@ -452,35 +519,49 @@ func (qtid QualifiedTableID) Key() TableKey {
 	if qtid.ID == "" {
 		panic("QualifiedTableID.Key called without an ID set")
 	}
-	return TableKey(fmt.Sprintf("%s%s%s",
-		qtid.TableQualifier.Key(),
+	return TableKey(fmt.Sprintf("%s%s%s%s%s%s%s",
+		PrefixTable,
+		TableKeyDelimiter,
+		qtid.OrganizationID,
+		TableKeyDelimiter,
+		qtid.DatabaseID,
 		TableKeyDelimiter,
 		qtid.ID))
 }
 
 // Equals returns true if `other` is the same as qtid. Note: the `Name` value is
-// ignored in this comparison; only `TableQualifier` and `ID` are considered.
+// ignored in this comparison; only `QualifiedDatabaseID` and `ID` are
+// considered.
 func (qtid QualifiedTableID) Equals(other QualifiedTableID) bool {
-	if qtid.TableQualifier == other.TableQualifier && qtid.ID == other.ID {
+	if qtid.QualifiedDatabaseID == other.QualifiedDatabaseID && qtid.ID == other.ID {
 		return true
 	}
 	return false
 }
 
+// Qualifier returns the QualifiedDatabaseID (qdbid) portion of the
+// QualifiedTableID (qtid).
+func (qtid QualifiedTableID) Qualifier() QualifiedDatabaseID {
+	return QualifiedDatabaseID{
+		OrganizationID: qtid.OrganizationID,
+		DatabaseID:     qtid.DatabaseID,
+	}
+}
+
 ////////////////////////////////////////////////
 
-// QualifiedTable wraps Table and includes a TableQualifier.
+// QualifiedTable wraps Table and includes a QualifiedDatabaseID.
 type QualifiedTable struct {
+	QualifiedDatabaseID
 	Table
-	TableQualifier
 }
 
 // NewQualifiedTable returns the tbl as a QualifiedTable with the provided
-// TableQualifier.
-func NewQualifiedTable(qual TableQualifier, tbl *Table) *QualifiedTable {
+// QualifiedDatabaseID.
+func NewQualifiedTable(qdbid QualifiedDatabaseID, tbl *Table) *QualifiedTable {
 	return &QualifiedTable{
-		Table:          *tbl,
-		TableQualifier: qual,
+		QualifiedDatabaseID: qdbid,
+		Table:               *tbl,
 	}
 }
 
@@ -496,17 +577,17 @@ func (qt QualifiedTable) String() string {
 	return fmt.Sprintf("%s (%s)", qt.QualifiedID(), qt.Name)
 }
 
-// Qualifier returns the TableQualifier portion of the QualifiedTable.
-func (qt *QualifiedTable) Qualifier() TableQualifier {
-	return qt.TableQualifier
+// Qualifier returns the QualifiedDatabaseID portion of the QualifiedTable.
+func (qt *QualifiedTable) Qualifier() QualifiedDatabaseID {
+	return qt.QualifiedDatabaseID
 }
 
 // QualifiedID returns the QualifiedTableID for the table.
 func (qt *QualifiedTable) QualifiedID() QualifiedTableID {
 	return QualifiedTableID{
-		TableQualifier: qt.TableQualifier,
-		ID:             qt.ID,
-		Name:           qt.Name,
+		QualifiedDatabaseID: qt.QualifiedDatabaseID,
+		ID:                  qt.ID,
+		Name:                qt.Name,
 	}
 }
 

dax/table_test.go

@@ -313,17 +313,17 @@ func TestTable(t *testing.T) {
 		t.Run("New", func(t *testing.T) {
 			tbl := dax.NewTable(tableName)
 			tbl.CreateID()
-			qual := dax.TableQualifier{
+			qdbid := dax.QualifiedDatabaseID{
 				OrganizationID: orgID,
 				DatabaseID:     dbID,
 			}
-			qtbl := dax.NewQualifiedTable(qual, tbl)
+			qtbl := dax.NewQualifiedTable(qdbid, tbl)
 			assert.NotEmpty(t, qtbl.ID)
 			assert.Equal(t, tbl.ID, qtbl.ID)
 
 			tq := qtbl.Qualifier()
-			assert.Equal(t, qual.OrganizationID, tq.OrganizationID)
-			assert.Equal(t, qual.DatabaseID, tq.DatabaseID)
+			assert.Equal(t, qdbid.OrganizationID, tq.OrganizationID)
+			assert.Equal(t, qdbid.DatabaseID, tq.DatabaseID)
 
 			wrappedTable := qtbl.Table
 			assert.Equal(t, tbl.Name, wrappedTable.Name)
@@ -345,11 +345,11 @@ func TestTable(t *testing.T) {
 		t.Run("ToJSON", func(t *testing.T) {
 			tbl := dax.NewTable(tableName)
 			tbl.CreateID()
-			qual := dax.TableQualifier{
+			qdbid := dax.QualifiedDatabaseID{
 				OrganizationID: orgID,
 				DatabaseID:     dbID,
 			}
-			qtbl := dax.NewQualifiedTable(qual, tbl)
+			qtbl := dax.NewQualifiedTable(qdbid, tbl)
 			id := qtbl.ID
 
 			b, err := json.Marshal(qtbl)

dax/test/dax/dax_test.go

@@ -26,7 +26,18 @@ func TestDAXIntegration(t *testing.T) {
 		t.Skip("skipping integration test")
 	}
 
-	qual := dax.NewTableQualifier("acme", "db1")
+	qdbid := dax.NewQualifiedDatabaseID("acme", "db1")
+	qdb := &dax.QualifiedDatabase{
+		OrganizationID: qdbid.OrganizationID,
+		Database: dax.Database{
+			ID:   qdbid.DatabaseID,
+			Name: "dbname1",
+			Options: dax.DatabaseOptions{
+				WorkersMin: 1,
+				WorkersMax: 1,
+			},
+		},
+	}
 
 	t.Run("ServiceStart", func(t *testing.T) {
 		t.Run("AllServicesByDefault", func(t *testing.T) {
@@ -104,6 +115,14 @@ func TestDAXIntegration(t *testing.T) {
 
 		svcmgr := mc.Manage()
 
+		// Set up MDS client.
+		mdsClient := mdsclient.New(svcmgr.MDS.Address(), svcmgr.Logger)
+
+		// Create database.
+		qdb.Options.WorkersMin = 1
+		qdb.Options.WorkersMax = 1
+		assert.NoError(t, mdsClient.CreateDatabase(context.Background(), qdb))
+
 		// skips is a list of tests which are currently not passing in dax. We
 		// need to get these passing before alpha.
 		skips := []string{
@@ -160,7 +179,7 @@ func TestDAXIntegration(t *testing.T) {
 
 		runTableTests(t,
 			svcmgr.Queryer.Address(),
-			basicTableTestConfig(qual, tableTests...)...,
+			basicTableTestConfig(qdbid, tableTests...)...,
 		)
 	})
 
@@ -168,9 +187,19 @@ func TestDAXIntegration(t *testing.T) {
 		mc := test.MustRunManagedCommand(t)
 		defer mc.Close()
 
+		svcmgr := mc.Manage()
+
+		// Set up MDS client.
+		mdsClient := mdsclient.New(svcmgr.MDS.Address(), svcmgr.Logger)
+
+		// Create database.
+		qdb.Options.WorkersMin = 1
+		qdb.Options.WorkersMax = 1
+		assert.NoError(t, mdsClient.CreateDatabase(context.Background(), qdb))
+
 		runTableTests(t,
 			mc.Manage().Queryer.Address(),
-			basicTableTestConfig(qual, defs.Keyed)...,
+			basicTableTestConfig(qdbid, defs.Keyed)...,
 		)
 	})
 
@@ -182,6 +211,14 @@ func TestDAXIntegration(t *testing.T) {
 
 		svcmgr := mc.Manage()
 
+		// Set up MDS client.
+		mdsClient := mdsclient.New(svcmgr.MDS.Address(), svcmgr.Logger)
+
+		// Create database.
+		qdb.Options.WorkersMin = 2
+		qdb.Options.WorkersMax = 2
+		assert.NoError(t, mdsClient.CreateDatabase(context.Background(), qdb))
+
 		computers := svcmgr.Computers()
 		computerKey0 := dax.ServiceKey(dax.ServicePrefixComputer + "0")
 		computerKey1 := dax.ServiceKey(dax.ServicePrefixComputer + "1")
@@ -189,13 +226,10 @@ func TestDAXIntegration(t *testing.T) {
 		// Ingest and query some data.
 		runTableTests(t,
 			svcmgr.Queryer.Address(),
-			basicTableTestConfig(qual, defs.Keyed)...,
+			basicTableTestConfig(qdbid, defs.Keyed)...,
 		)
 
-		// Set up MDS client.
-		mdsClient := mdsclient.New(svcmgr.MDS.Address(), svcmgr.Logger)
-
-		qtid, err := mdsClient.TableID(context.Background(), qual, dax.TableName(defs.Keyed.Name(0)))
+		qtid, err := mdsClient.TableID(context.Background(), qdbid, dax.TableName(defs.Keyed.Name(0)))
 		assert.NoError(t, err)
 
 		// ensure partitions are covered
@@ -238,13 +272,21 @@ func TestDAXIntegration(t *testing.T) {
 
 		svcmgr := mc.Manage()
 
+		// Set up MDS client.
+		mdsClient := mdsclient.New(svcmgr.MDS.Address(), svcmgr.Logger)
+
+		// Create database.
+		qdb.Options.WorkersMin = 1
+		qdb.Options.WorkersMax = 1
+		assert.NoError(t, mdsClient.CreateDatabase(context.Background(), qdb))
+
 		computerKey0 := dax.ServiceKey(dax.ServicePrefixComputer + "0")
 
 		// Ingest and query some data.
 		t.Run("ingest and query some data", func(t *testing.T) {
 			runTableTests(t,
 				svcmgr.Queryer.Address(),
-				basicTableTestConfig(qual, defs.Keyed)...,
+				basicTableTestConfig(qdbid, defs.Keyed)...,
 			)
 		})
 
@@ -268,7 +310,7 @@ func TestDAXIntegration(t *testing.T) {
 			runTableTests(t,
 				svcmgr.Queryer.Address(),
 				tableTestConfig{
-					qual:       qual,
+					qdbid:      qdbid,
 					test:       defs.Keyed,
 					skipCreate: true,
 					skipInsert: true,
@@ -284,6 +326,14 @@ func TestDAXIntegration(t *testing.T) {
 
 		svcmgr := mc.Manage()
 
+		// Set up MDS client.
+		mdsClient := mdsclient.New(svcmgr.MDS.Address(), svcmgr.Logger)
+
+		// Create database.
+		qdb.Options.WorkersMin = 1
+		qdb.Options.WorkersMax = 1
+		assert.NoError(t, mdsClient.CreateDatabase(context.Background(), qdb))
+
 		computerKey0 := dax.ServiceKey(dax.ServicePrefixComputer + "0")
 
 		// Ingest and query some data.
@@ -291,19 +341,16 @@ func TestDAXIntegration(t *testing.T) {
 			runTableTests(t,
 				svcmgr.Queryer.Address(),
 				tableTestConfig{
-					qual:      qual,
+					qdbid:     qdbid,
 					test:      defs.Keyed,
 					insertSet: 0,
 				},
 			)
 		})
 
-		// Set up MDS client.
-		mdsClient := mdsclient.New(svcmgr.MDS.Address(), svcmgr.Logger)
-
 		// Snapshot table
 		ctx := context.Background()
-		qtid, err := mdsClient.TableID(ctx, qual, dax.TableName(defs.Keyed.Name(0)))
+		qtid, err := mdsClient.TableID(ctx, qdbid, dax.TableName(defs.Keyed.Name(0)))
 		assert.NoError(t, err)
 
 		mdsClient.SnapshotTable(ctx, qtid)
@@ -313,7 +360,7 @@ func TestDAXIntegration(t *testing.T) {
 			runTableTests(t,
 				svcmgr.Queryer.Address(),
 				tableTestConfig{
-					qual:       qual,
+					qdbid:      qdbid,
 					test:       defs.Keyed,
 					skipCreate: true,
 					insertSet:  1,
@@ -341,7 +388,7 @@ func TestDAXIntegration(t *testing.T) {
 			runTableTests(t,
 				svcmgr.Queryer.Address(),
 				tableTestConfig{
-					qual:       qual,
+					qdbid:      qdbid,
 					test:       defs.Keyed,
 					skipCreate: true,
 					skipInsert: true,
@@ -357,6 +404,14 @@ func TestDAXIntegration(t *testing.T) {
 
 		svcmgr := mc.Manage()
 
+		// Set up MDS client.
+		mdsClient := mdsclient.New(svcmgr.MDS.Address(), svcmgr.Logger)
+
+		// Create database.
+		qdb.Options.WorkersMin = 1
+		qdb.Options.WorkersMax = 1
+		assert.NoError(t, mdsClient.CreateDatabase(context.Background(), qdb))
+
 		computerKey0 := dax.ServiceKey(dax.ServicePrefixComputer + "0")
 
 		mdsKey := dax.ServiceKey(dax.ServicePrefixMDS)
@@ -365,13 +420,10 @@ func TestDAXIntegration(t *testing.T) {
 		// Ingest and query some data.
 		runTableTests(t,
 			svcmgr.Queryer.Address(),
-			basicTableTestConfig(qual, defs.Keyed)...,
+			basicTableTestConfig(qdbid, defs.Keyed)...,
 		)
 
-		// Set up MDS client.
-		mdsClient := mdsclient.New(svcmgr.MDS.Address(), svcmgr.Logger)
-
-		qtid, err := mdsClient.TableID(context.Background(), qual, dax.TableName(defs.Keyed.Name(0)))
+		qtid, err := mdsClient.TableID(context.Background(), qdbid, dax.TableName(defs.Keyed.Name(0)))
 		assert.NoError(t, err)
 
 		// ensure partitions are covered
@@ -410,13 +462,21 @@ func TestDAXIntegration(t *testing.T) {
 
 		svcmgr := mc.Manage()
 
+		// Set up MDS client.
+		mdsClient := mdsclient.New(svcmgr.MDS.Address(), svcmgr.Logger)
+
+		// Create database.
+		qdb.Options.WorkersMin = 1
+		qdb.Options.WorkersMax = 1
+		assert.NoError(t, mdsClient.CreateDatabase(context.Background(), qdb))
+
 		computerKey0 := dax.ServiceKey(dax.ServicePrefixComputer + "0")
 
 		// Ingest and query some data.
 		t.Run("ingest and query some data", func(t *testing.T) {
 			runTableTests(t,
 				svcmgr.Queryer.Address(),
-				basicTableTestConfig(qual, defs.Keyed)...,
+				basicTableTestConfig(qdbid, defs.Keyed)...,
 			)
 		})
 
@@ -439,7 +499,7 @@ func TestDAXIntegration(t *testing.T) {
 			runTableTests(t,
 				svcmgr.Queryer.Address(),
 				tableTestConfig{
-					qual:       qual,
+					qdbid:      qdbid,
 					test:       defs.Keyed,
 					skipCreate: true,
 					skipInsert: true,
@@ -452,7 +512,7 @@ func TestDAXIntegration(t *testing.T) {
 ///////////////////////////////////////////////////
 
 type tableTestConfig struct {
-	qual       dax.TableQualifier
+	qdbid      dax.QualifiedDatabaseID
 	test       defs.TableTest
 	skipCreate bool
 	skipInsert bool
@@ -461,20 +521,19 @@ type tableTestConfig struct {
 	querySet   int
 }
 
-func basicTableTestConfig(qual dax.TableQualifier, tests ...defs.TableTest) []tableTestConfig {
+func basicTableTestConfig(qdbid dax.QualifiedDatabaseID, tests ...defs.TableTest) []tableTestConfig {
 	ret := make([]tableTestConfig, len(tests))
 
 	for i := range tests {
 		ret[i] = tableTestConfig{
-			qual: qual,
-			test: tests[i],
+			qdbid: qdbid,
+			test:  tests[i],
 		}
 	}
 
 	return ret
 }
 
-// func runTableTests(t *testing.T, queryerAddr dax.Address, qual dax.TableQualifier, doCreate bool, tests ...defs.TableTest) {
 func runTableTests(t *testing.T, queryerAddr dax.Address, cfgs ...tableTestConfig) {
 	emptyWireQueryResponse := &featurebase.WireQueryResponse{
 		Schema: featurebase.WireQuerySchema{
@@ -491,7 +550,7 @@ func runTableTests(t *testing.T, queryerAddr dax.Address, cfgs ...tableTestConfi
 			if !cfg.skipCreate {
 				// Create a table.
 				if cfg.test.HasTable() {
-					resp := runSQL(t, queryerAddr, cfg.qual, cfg.test.CreateTable())
+					resp := runSQL(t, queryerAddr, cfg.qdbid, cfg.test.CreateTable())
 					assertResponseEqual(t, emptyWireQueryResponse, resp)
 				}
 			}
@@ -499,7 +558,7 @@ func runTableTests(t *testing.T, queryerAddr dax.Address, cfgs ...tableTestConfi
 			if !cfg.skipInsert {
 				// Populate table with data.
 				if cfg.test.HasTable() && cfg.test.HasData() {
-					resp := runSQL(t, queryerAddr, cfg.qual, cfg.test.InsertInto(t, cfg.insertSet))
+					resp := runSQL(t, queryerAddr, cfg.qdbid, cfg.test.InsertInto(t, cfg.insertSet))
 					assertResponseEqual(t, emptyWireQueryResponse, resp)
 				}
 			}
@@ -525,7 +584,7 @@ func runTableTests(t *testing.T, queryerAddr dax.Address, cfgs ...tableTestConfi
 								expRows = sqltest.ExpRowsPlus1[cfg.querySet-1]
 							}
 
-							resp := runSQL(t, queryerAddr, cfg.qual, sql)
+							resp := runSQL(t, queryerAddr, cfg.qdbid, sql)
 							headers := resp.Schema.Fields
 							rows := resp.Data
 							var err error
@@ -602,7 +661,7 @@ func runTableTests(t *testing.T, queryerAddr dax.Address, cfgs ...tableTestConfi
 								expRows = pqltest.ExpRowsPlus1[cfg.querySet-1]
 							}
 
-							resp := runPQL(t, queryerAddr, cfg.qual, pqltest.Table, pql)
+							resp := runPQL(t, queryerAddr, cfg.qdbid, pqltest.Table, pql)
 							headers := resp.Schema.Fields
 							rows := resp.Data
 							var err error
@@ -680,23 +739,23 @@ func (c *wireResponseComparer) Equal() bool {
 	return assert.Equal(c.tb, c.exp, c.got)
 }
 
-func runSQL(tb testing.TB, queryerAddr dax.Address, qual dax.TableQualifier, sql string) *featurebase.WireQueryResponse {
+func runSQL(tb testing.TB, queryerAddr dax.Address, qdbid dax.QualifiedDatabaseID, sql string) *featurebase.WireQueryResponse {
 	tb.Helper()
 
 	client := queryerclient.New(queryerAddr, logger.StderrLogger)
 
-	resp, err := client.QuerySQL(context.Background(), qual, sql)
+	resp, err := client.QuerySQL(context.Background(), qdbid, sql)
 	assert.NoError(tb, err)
 
 	return resp
 }
 
-func runPQL(tb testing.TB, queryerAddr dax.Address, qual dax.TableQualifier, table string, pql string) *featurebase.WireQueryResponse {
+func runPQL(tb testing.TB, queryerAddr dax.Address, qdbid dax.QualifiedDatabaseID, table string, pql string) *featurebase.WireQueryResponse {
 	tb.Helper()
 
 	client := queryerclient.New(queryerAddr, logger.StderrLogger)
 
-	resp, err := client.QueryPQL(context.Background(), qual, dax.TableName(table), pql)
+	resp, err := client.QueryPQL(context.Background(), qdbid, dax.TableName(table), pql)
 	assert.NoError(tb, err)
 
 	return resp

dax/test/table.go

@@ -11,7 +11,7 @@ import (
 // general configuration. This function creates a Table with a random TableID.
 // If you need to specify the TableID yourself, use the TestQualifiedTableWithID
 // function.
-func TestQualifiedTable(t *testing.T, qual dax.TableQualifier, name dax.TableName, partitionN int, keyed bool) *dax.QualifiedTable {
+func TestQualifiedTable(t *testing.T, qdbid dax.QualifiedDatabaseID, name dax.TableName, partitionN int, keyed bool) *dax.QualifiedTable {
 	t.Helper()
 
 	var pkFieldType dax.BaseType
@@ -22,7 +22,6 @@ func TestQualifiedTable(t *testing.T, qual dax.TableQualifier, name dax.TableNam
 	}
 
 	tbl := dax.NewTable(name)
-	tbl.CreateID()
 	tbl.PartitionN = partitionN
 	tbl.Fields = []*dax.Field{
 		{
@@ -32,14 +31,33 @@ func TestQualifiedTable(t *testing.T, qual dax.TableQualifier, name dax.TableNam
 	}
 
 	return dax.NewQualifiedTable(
-		qual,
+		qdbid,
 		tbl,
 	)
 }
 
+// TestQualifiedDatabaseWithID is a test helper function for creating a database
+// based on a general configuration, and having the specified DatabaseID.
+func TestQualifiedDatabaseWithID(t *testing.T, orgID dax.OrganizationID, id dax.DatabaseID, name dax.DatabaseName, opts dax.DatabaseOptions) *dax.QualifiedDatabase {
+	t.Helper()
+
+	db := dax.Database{
+		ID:      dax.DatabaseID(id),
+		Name:    name,
+		Options: opts,
+	}
+
+	qdb := &dax.QualifiedDatabase{
+		OrganizationID: orgID,
+		Database:       db,
+	}
+
+	return qdb
+}
+
 // TestQualifiedTableWithID is a test helper function for creating a table based
 // on a general configuration, and having the specified TableID.
-func TestQualifiedTableWithID(t *testing.T, qual dax.TableQualifier, id string, name dax.TableName, partitionN int, keyed bool) *dax.QualifiedTable {
+func TestQualifiedTableWithID(t *testing.T, qdbid dax.QualifiedDatabaseID, id string, name dax.TableName, partitionN int, keyed bool) *dax.QualifiedTable {
 	t.Helper()
 
 	var pkFieldType dax.BaseType
@@ -63,7 +81,7 @@ func TestQualifiedTableWithID(t *testing.T, qual dax.TableQualifier, id string,
 	}
 
 	return dax.NewQualifiedTable(
-		qual,
+		qdbid,
 		tbl,
 	)
 }

dax/transaction.go

@@ -0,0 +1,9 @@
+package dax
+
+import "context"
+
+type Transaction interface {
+	Commit() error
+	Context() context.Context
+	Rollback() error
+}

dax/workerjob.go

@@ -5,44 +5,39 @@ import (
 	"strings"
 )
 
-// Worker is a generic identifier used to represent a service responsible for
-// doing certain jobs. In the case of dax, this is typically the Address of a
-// compute or translate node. Services such as the Balancer use Workers (as
-// opposed to specifically using Address) in order to remain generic, and to
-// keep the business logic between services slightly less coupled.
-type Worker string
-
-// Workers is a sortable slice of Worker.
-type Workers []Worker
-
-func (w Workers) Len() int           { return len(w) }
-func (w Workers) Less(i, j int) bool { return w[i] < w[j] }
-func (w Workers) Swap(i, j int)      { w[i], w[j] = w[j], w[i] }
-
 // Job is a generic identifier used to represent a specific role assigned to a
 // worker.
 type Job string
 
+// Job allows a Job to implement the Jobber interface.
+func (j Job) Job() Job {
+	return j
+}
+
 // Jobs is a slice of Job.
 type Jobs []Job
 
+type Jobber interface {
+	Job() Job
+}
+
 // WorkerInfo represents a Worker and the Jobs to which it has been assigned.
 type WorkerInfo struct {
-	ID   Worker
-	Jobs []Job
+	Address Address
+	Jobs    []Job
 }
 
 // WorkerInfos is a sortable slice of WorkerInfo.
 type WorkerInfos []WorkerInfo
 
 func (w WorkerInfos) Len() int           { return len(w) }
-func (w WorkerInfos) Less(i, j int) bool { return w[i].ID < w[j].ID }
+func (w WorkerInfos) Less(i, j int) bool { return w[i].Address < w[j].Address }
 func (w WorkerInfos) Swap(i, j int)      { w[i], w[j] = w[j], w[i] }
 
 // WorkerDiff represents the changes made to a Worker following the latest
 // event.
 type WorkerDiff struct {
-	WorkerID    Worker
+	Address     Address
 	AddedJobs   []Job
 	RemovedJobs []Job
 }
@@ -51,7 +46,7 @@ type WorkerDiff struct {
 // ID. Any job that is added and then removed or removed and then
 // added cancels out and won't be present after add is called.
 func (w *WorkerDiff) Add(w2 WorkerDiff) {
-	if w.WorkerID != w2.WorkerID {
+	if w.Address != w2.Address {
 		panic("can't add worker diffs from different workers")
 	}
 	a1 := NewSet(w.AddedJobs...)
@@ -74,7 +69,7 @@ func (w *WorkerDiff) Add(w2 WorkerDiff) {
 type WorkerDiffs []WorkerDiff
 
 func (w WorkerDiffs) Len() int           { return len(w) }
-func (w WorkerDiffs) Less(i, j int) bool { return w[i].WorkerID < w[j].WorkerID }
+func (w WorkerDiffs) Less(i, j int) bool { return w[i].Address < w[j].Address }
 func (w WorkerDiffs) Swap(i, j int)      { w[i], w[j] = w[j], w[i] }
 
 // Set is a set of stringy items.
@@ -109,7 +104,8 @@ func (s Set[K]) Remove(k K) {
 	delete(s, k)
 }
 
-func (s Set[K]) RemovePrefix(prefix string) []K {
+// RemoveByPrefix removes all items from Set that have the given prefix.
+func (s Set[K]) RemoveByPrefix(prefix string) []K {
 	ret := make([]K, 0)
 	for k := range s {
 		if strings.HasPrefix(string(k), prefix) {

idk/ingest.go

@@ -1007,10 +1007,10 @@ func (m *Main) setupClient() (*tls.Config, error) {
 		// MDS doesn't auto-create a table based on IDK ingest; the table must
 		// already exist.
 		mdsClient := mdsclient.New(dax.Address(m.MDSAddress), m.log)
-		qual := dax.NewTableQualifier(m.OrganizationID, m.DatabaseID)
-		qtid, err := mdsClient.TableID(ctx, qual, m.TableName)
+		qdbid := dax.NewQualifiedDatabaseID(m.OrganizationID, m.DatabaseID)
+		qtid, err := mdsClient.TableID(ctx, qdbid, m.TableName)
 		if err != nil {
-			return nil, errors.Wrapf(err, "getting table id: qual: %s, table name: %s", qual, m.TableName)
+			return nil, errors.Wrapf(err, "getting table id: qual: %s, table name: %s", qdbid, m.TableName)
 		}
 		qtbl, err := mdsClient.Table(ctx, qtid)
 		if err != nil {
@@ -1019,7 +1019,7 @@ func (m *Main) setupClient() (*tls.Config, error) {
 
 		m.Qtbl = qtbl
 		m.Index = string(qtbl.Key())
-		m.SchemaManager = mds.NewSchemaManager(dax.Address(m.MDSAddress), qual, m.log)
+		m.SchemaManager = mds.NewSchemaManager(dax.Address(m.MDSAddress), qdbid, m.log)
 
 		m.NewImporterFn = func() pilosacore.Importer {
 			return mds.NewImporter(mdsClient, mdsClient, qtbl.Qualifier(), &qtbl.Table)

idk/ingest_test.go

@@ -58,7 +58,7 @@ func configureTestFlagsMDS(main *Main, address dax.Address, qtbl *dax.QualifiedT
 
 	mdsClient := mdsclient.New(dax.Address(address), logger.StderrLogger)
 	main.NewImporterFn = func() pilosa.Importer {
-		return mds.NewImporter(mdsClient, mdsClient, qtbl.TableQualifier, &qtbl.Table)
+		return mds.NewImporter(mdsClient, mdsClient, qtbl.QualifiedDatabaseID, &qtbl.Table)
 	}
 }
 
@@ -1756,6 +1756,24 @@ func TestBatchTargetMDS(t *testing.T) {
 	orgID := dax.OrganizationID("acme")
 	dbID := dax.DatabaseID("db1")
 
+	mdsClient := mdsclient.New(mdsAddress, logger.StderrLogger)
+
+	ctx := context.Background()
+
+	// Create the database.
+	qdb := &dax.QualifiedDatabase{
+		OrganizationID: orgID,
+		Database: dax.Database{
+			ID:   dbID,
+			Name: "dbname1",
+			Options: dax.DatabaseOptions{
+				WorkersMin: 1,
+				WorkersMax: 1,
+			},
+		},
+	}
+	mdsClient.CreateDatabase(ctx, qdb)
+
 	t.Run("FieldTypes", func(t *testing.T) {
 		tests := []struct {
 			fieldType    dax.BaseType
@@ -1859,14 +1877,11 @@ func TestBatchTargetMDS(t *testing.T) {
 				}
 
 				qtbl := dax.NewQualifiedTable(
-					dax.NewTableQualifier(orgID, dbID),
+					dax.NewQualifiedDatabaseID(orgID, dbID),
 					tbl,
 				)
 
-				ctx := context.Background()
-
 				// Create the table in MDS Schemar.
-				mdsClient := mdsclient.New(mdsAddress, logger.StderrLogger)
 				if err := mdsClient.CreateTable(ctx, qtbl); err != nil {
 					t.Fatalf("creating table: %v", err)
 				}

idk/mds/importer.go

@@ -21,16 +21,16 @@ type importer struct {
 	noder   dax.Noder
 	schemar dax.Schemar
 
-	mu   sync.Mutex
-	qual dax.TableQualifier
-	tbl  *dax.Table
+	mu    sync.Mutex
+	qdbid dax.QualifiedDatabaseID
+	tbl   *dax.Table
 }
 
-func NewImporter(noder dax.Noder, schemar dax.Schemar, qual dax.TableQualifier, tbl *dax.Table) *importer {
+func NewImporter(noder dax.Noder, schemar dax.Schemar, qdbid dax.QualifiedDatabaseID, tbl *dax.Table) *importer {
 	return &importer{
 		noder:   noder,
 		schemar: schemar,
-		qual:    qual,
+		qdbid:   qdbid,
 		tbl:     tbl,
 	}
 }
@@ -261,10 +261,10 @@ func (m *importer) getQtbl(ctx context.Context, tid dax.TableID) (*dax.Qualified
 	defer m.mu.Unlock()
 
 	if m.tbl != nil {
-		return dax.NewQualifiedTable(m.qual, m.tbl), nil
+		return dax.NewQualifiedTable(m.qdbid, m.tbl), nil
 	}
 
-	qtid := dax.NewQualifiedTableID(m.qual, tid)
+	qtid := dax.NewQualifiedTableID(m.qdbid, tid)
 
 	qtbl, err := m.schemar.TableByID(ctx, qtid)
 	if err != nil {

idk/mds/schemamanager.go

@@ -18,14 +18,14 @@ import (
 // schemaManager
 type schemaManager struct {
 	client *mdsclient.Client
-	qual   dax.TableQualifier
+	qdbid  dax.QualifiedDatabaseID
 	logger logger.Logger
 }
 
-func NewSchemaManager(mdsAddress dax.Address, qual dax.TableQualifier, logger logger.Logger) *schemaManager {
+func NewSchemaManager(mdsAddress dax.Address, qdbid dax.QualifiedDatabaseID, logger logger.Logger) *schemaManager {
 	return &schemaManager{
 		client: mdsclient.New(mdsAddress, logger),
-		qual:   qual,
+		qdbid:  qdbid,
 		logger: logger,
 	}
 }
@@ -41,7 +41,7 @@ func (s *schemaManager) Schema() (*featurebase_client.Schema, error) {
 	// method returns.
 	schema := featurebase_client.NewSchema()
 
-	tables, err := s.client.Tables(context.Background(), s.qual)
+	tables, err := s.client.Tables(context.Background(), s.qdbid)
 	if err != nil {
 		return nil, err
 	}

sql3/planner/compilecreateview.go

@@ -3,8 +3,8 @@
 package planner
 
 import (
-	"github.com/molecula/featurebase/v3/sql3/parser"
-	"github.com/molecula/featurebase/v3/sql3/planner/types"
+	"github.com/featurebasedb/featurebase/v3/sql3/parser"
+	"github.com/featurebasedb/featurebase/v3/sql3/planner/types"
 )
 
 // compileCreateViewStatement compiles a parser.CreateViewStatement AST into a PlanOperator

sql3/planner/compiledropview.go

@@ -3,9 +3,9 @@
 package planner
 
 import (
-	"github.com/molecula/featurebase/v3/sql3"
-	"github.com/molecula/featurebase/v3/sql3/parser"
-	"github.com/molecula/featurebase/v3/sql3/planner/types"
+	"github.com/featurebasedb/featurebase/v3/sql3"
+	"github.com/featurebasedb/featurebase/v3/sql3/parser"
+	"github.com/featurebasedb/featurebase/v3/sql3/planner/types"
 )
 
 // compileDropViewStatement compiles a DROP VIEW statement into a PlanOperator.

sql3/planner/opalterview.go

@@ -6,8 +6,8 @@ import (
 	"context"
 	"fmt"
 
-	"github.com/molecula/featurebase/v3/sql3"
-	"github.com/molecula/featurebase/v3/sql3/planner/types"
+	"github.com/featurebasedb/featurebase/v3/sql3"
+	"github.com/featurebasedb/featurebase/v3/sql3/planner/types"
 )
 
 // PlanOpAlterView implements the ALTER VIEW operator

sql3/planner/opcreateview.go

@@ -6,9 +6,9 @@ import (
 	"context"
 	"fmt"
 
-	"github.com/molecula/featurebase/v3/dax"
-	"github.com/molecula/featurebase/v3/sql3"
-	"github.com/molecula/featurebase/v3/sql3/planner/types"
+	"github.com/featurebasedb/featurebase/v3/dax"
+	"github.com/featurebasedb/featurebase/v3/sql3"
+	"github.com/featurebasedb/featurebase/v3/sql3/planner/types"
 )
 
 // PlanOpCreateView implements the CREATE VIEW operator

sql3/planner/opdropview.go

@@ -6,8 +6,8 @@ import (
 	"context"
 	"fmt"
 
-	"github.com/molecula/featurebase/v3/sql3"
-	"github.com/molecula/featurebase/v3/sql3/planner/types"
+	"github.com/featurebasedb/featurebase/v3/sql3"
+	"github.com/featurebasedb/featurebase/v3/sql3/planner/types"
 )
 
 // PlanOpDropView plan operator to drop a view.

sql3/planner/opfeaturebasetables.go

@@ -135,11 +135,17 @@ func (i *showTablesRowIter) Next(ctx context.Context) (types.Row, error) {
 			}
 		default:
 			u := i.planner.systemAPI.DataDir()
-			u = fmt.Sprintf("%s/indexes/%s", u, indexName)
 
-			spaceUsed, err = pilosa.GetDiskUsage(u)
-			if err != nil {
-				return nil, err
+			// TODO(tlt): GetDiskUsage needs to be behind an interface because
+			// this doesn't work in serverless. For now I'm just going to skip
+			// it based on the emtpy DataDir, but let's do this the right way.
+			if u != "" {
+				u = fmt.Sprintf("%s/indexes/%s", u, indexName)
+
+				spaceUsed, err = pilosa.GetDiskUsage(u)
+				if err != nil {
+					return nil, err
+				}
 			}
 		}
 

sql3/planner/systemobjects.go

@@ -6,11 +6,11 @@ import (
 	"context"
 	"time"
 
-	pilosa "github.com/molecula/featurebase/v3"
-	"github.com/molecula/featurebase/v3/dax"
-	"github.com/molecula/featurebase/v3/sql3"
-	"github.com/molecula/featurebase/v3/sql3/parser"
-	"github.com/molecula/featurebase/v3/sql3/planner/types"
+	pilosa "github.com/featurebasedb/featurebase/v3"
+	"github.com/featurebasedb/featurebase/v3/dax"
+	"github.com/featurebasedb/featurebase/v3/sql3"
+	"github.com/featurebasedb/featurebase/v3/sql3/parser"
+	"github.com/featurebasedb/featurebase/v3/sql3/planner/types"
 )
 
 type viewSystemObject struct {