Repository: google/exposure-notifications-server
Pull request #410 (2020-05-21 22:18:44)

Separate db/model federation

google-prow-robot: [APPROVALNOTIFIER] This PR is **APPROVED**

This pull-request has been approved by: *<a href="https://github.com/google/exposure-notifications-server/pull/410#" title="Author self-approved">crwilcox</a>*

The full list of commands accepted by this bot can be found [here](https://go.k8s.io/bot-commands?repo=google%2Fexposure-notifications-server).

The pull request process is described [here](https://git.k8s.io/community/contributors/guide/owners.md#the-code-review-process)

<details >
Needs approval from an approver in each of these files:

- ~~[OWNERS](https://github.com/google/exposure-notifications-server/blob/master/OWNERS)~~ [crwilcox]

Approvers can indicate their approval by writing `/approve` in a comment
Approvers can cancel approval by writing `/approve cancel` in a comment
</details>
<!-- META={"approvers":[]} -->
google-prow-robot: @crwilcox: The following test **failed**, say `/retest` to rerun all failed tests:

Test name | Commit | Details | Rerun command
--- | --- | --- | ---
pull-release-unit | 6fedbbc2efcf4a2203d7b366813e618b7325b5d3 | [link](https://prow-internal.gcpnode.com/view/gcs/gke-prow/pr-logs/pull/google_exposure-notifications-server/410/pull-release-unit/1263595133559902208) | `/test pull-release-unit`

[Full PR test history](https://prow-internal.gcpnode.com/pr-history?org=google&repo=exposure-notifications-server&pr=410). [Your PR dashboard](https://gubernator-internal.googleplex.com/pr/crwilcox).

<details>

Instructions for interacting with me using PR comments are available [here](https://git.k8s.io/community/contributors/guide/pull-requests.md).  If you have questions or suggestions related to my behavior, please file an issue against the [kubernetes/test-infra](https://github.com/kubernetes/test-infra/issues/new?title=Prow%20issue:) repository. I understand the commands that are listed [here](https://go.k8s.io/bot-commands).
</details>
<!-- test report -->
crwilcox: I came across an issue where tests were contending over create/drop of DB:
```
{"severity":"info","ts":1590100649.22637,"logger":"default","caller":"database/connection.go:58","message":"Creating connection pool."}
2020/05/21 15:37:34 creating test DB: ERROR: database "exposure-server-test" is being accessed by other users (SQLSTATE 55006)
FAIL    github.com/google/exposure-notifications-server/internal/publish/database       5.610s
```

This change now also includes a fix for that.

Patch:
diff --git a/internal/database/federationin.go b/internal/federationin/database/federationin.go
similarity index 76%
rename from internal/database/federationin.go
rename to internal/federationin/database/federationin.go
index 22ee0d8..0d6c7ef 100644
--- a/internal/database/federationin.go
+++ b/internal/federationin/database/federationin.go
@@ -19,17 +19,34 @@ import (
 	"fmt"
 	"time"
 
+	"github.com/google/exposure-notifications-server/internal/database"
+	"github.com/google/exposure-notifications-server/internal/federationin/model"
 	pgx "github.com/jackc/pgx/v4"
 )
 
+type FederationInDB struct {
+	db *database.DB
+}
+
+func New(db *database.DB) *FederationInDB {
+	return &FederationInDB{
+		db: db,
+	}
+}
+
 // FinalizeSyncFn is used to finalize a historical sync record.
 type FinalizeSyncFn func(maxTimestamp time.Time, totalInserted int) error
 
 type queryRowFn func(ctx context.Context, query string, args ...interface{}) pgx.Row
 
+// Lock acquires lock with given name that times out after ttl. Returns an UnlockFn that can be used to unlock the lock. ErrAlreadyLocked will be returned if there is already a lock in use.
+func (db *FederationInDB) Lock(ctx context.Context, lockID string, ttl time.Duration) (database.UnlockFn, error) {
+	return db.db.Lock(ctx, lockID, ttl)
+}
+
 // GetFederationInQuery returns a query for given queryID. If not found, ErrNotFound will be returned.
-func (db *DB) GetFederationInQuery(ctx context.Context, queryID string) (*FederationInQuery, error) {
-	conn, err := db.Pool.Acquire(ctx)
+func (db *FederationInDB) GetFederationInQuery(ctx context.Context, queryID string) (*model.FederationInQuery, error) {
+	conn, err := db.db.Pool.Acquire(ctx)
 	if err != nil {
 		return nil, fmt.Errorf("acquiring connection: %w", err)
 	}
@@ -38,7 +55,7 @@ func (db *DB) GetFederationInQuery(ctx context.Context, queryID string) (*Federa
 	return getFederationInQuery(ctx, queryID, conn.QueryRow)
 }
 
-func getFederationInQuery(ctx context.Context, queryID string, queryRow queryRowFn) (*FederationInQuery, error) {
+func getFederationInQuery(ctx context.Context, queryID string, queryRow queryRowFn) (*model.FederationInQuery, error) {
 	row := queryRow(ctx, `
 		SELECT
 			query_id, server_addr, oidc_audience, include_regions, exclude_regions, last_timestamp
@@ -49,10 +66,10 @@ func getFederationInQuery(ctx context.Context, queryID string, queryRow queryRow
 		`, queryID)
 
 	// See https://www.opsdash.com/blog/postgres-arrays-golang.html for working with Postgres arrays in Go.
-	q := FederationInQuery{}
+	q := model.FederationInQuery{}
 	if err := row.Scan(&q.QueryID, &q.ServerAddr, &q.Audience, &q.IncludeRegions, &q.ExcludeRegions, &q.LastTimestamp); err != nil {
 		if err == pgx.ErrNoRows {
-			return nil, ErrNotFound
+			return nil, database.ErrNotFound
 		}
 		return nil, fmt.Errorf("scanning results: %w", err)
 	}
@@ -60,8 +77,8 @@ func getFederationInQuery(ctx context.Context, queryID string, queryRow queryRow
 }
 
 // AddFederationInQuery adds a FederationInQuery entity. It will overwrite a query with matching q.queryID if it exists.
-func (db *DB) AddFederationInQuery(ctx context.Context, q *FederationInQuery) error {
-	return db.InTx(ctx, pgx.Serializable, func(tx pgx.Tx) error {
+func (db *FederationInDB) AddFederationInQuery(ctx context.Context, q *model.FederationInQuery) error {
+	return db.db.InTx(ctx, pgx.Serializable, func(tx pgx.Tx) error {
 		query := `
 			INSERT INTO
 				FederationInQuery
@@ -82,8 +99,8 @@ func (db *DB) AddFederationInQuery(ctx context.Context, q *FederationInQuery) er
 }
 
 // GetFederationInSync returns a federation sync record for given syncID. If not found, ErrNotFound will be returned.
-func (db *DB) GetFederationInSync(ctx context.Context, syncID int64) (*FederationInSync, error) {
-	conn, err := db.Pool.Acquire(ctx)
+func (db *FederationInDB) GetFederationInSync(ctx context.Context, syncID int64) (*model.FederationInSync, error) {
+	conn, err := db.db.Pool.Acquire(ctx)
 	if err != nil {
 		return nil, fmt.Errorf("acquiring connection: %w", err)
 	}
@@ -92,7 +109,7 @@ func (db *DB) GetFederationInSync(ctx context.Context, syncID int64) (*Federatio
 	return getFederationInSync(ctx, syncID, conn.QueryRow)
 }
 
-func getFederationInSync(ctx context.Context, syncID int64, queryRowContext queryRowFn) (*FederationInSync, error) {
+func getFederationInSync(ctx context.Context, syncID int64, queryRowContext queryRowFn) (*model.FederationInSync, error) {
 	row := queryRowContext(ctx, `
 		SELECT
 			sync_id, query_id, started, completed, insertions, max_timestamp
@@ -102,14 +119,14 @@ func getFederationInSync(ctx context.Context, syncID int64, queryRowContext quer
 			sync_id=$1
 		`, syncID)
 
-	s := FederationInSync{}
+	s := model.FederationInSync{}
 	var (
 		completed, max *time.Time
 		insertions     *int
 	)
 	if err := row.Scan(&s.SyncID, &s.QueryID, &s.Started, &completed, &insertions, &max); err != nil {
 		if err == pgx.ErrNoRows {
-			return nil, ErrNotFound
+			return nil, database.ErrNotFound
 		}
 		return nil, fmt.Errorf("scanning results: %w", err)
 	}
@@ -126,8 +143,8 @@ func getFederationInSync(ctx context.Context, syncID int64, queryRowContext quer
 }
 
 // StartFederationInSync stores a historical record of a query sync starting. It returns a FederationInSync key, and a FinalizeSyncFn that must be invoked to finalize the historical record.
-func (db *DB) StartFederationInSync(ctx context.Context, q *FederationInQuery, started time.Time) (int64, FinalizeSyncFn, error) {
-	conn, err := db.Pool.Acquire(ctx)
+func (db *FederationInDB) StartFederationInSync(ctx context.Context, q *model.FederationInQuery, started time.Time) (int64, FinalizeSyncFn, error) {
+	conn, err := db.db.Pool.Acquire(ctx)
 	if err != nil {
 		return 0, nil, fmt.Errorf("acquiring connection: %w", err)
 	}
@@ -153,7 +170,7 @@ func (db *DB) StartFederationInSync(ctx context.Context, q *FederationInQuery, s
 	finalize := func(maxTimestamp time.Time, totalInserted int) error {
 		completed := started.Add(time.Now().Sub(startedTimer))
 
-		return db.InTx(ctx, pgx.Serializable, func(tx pgx.Tx) error {
+		return db.db.InTx(ctx, pgx.Serializable, func(tx pgx.Tx) error {
 			// Special case: when no keys are pulled, the maxTimestamp will be 0, so we don't update the
 			// FederationQuery in this case to prevent it from going back and fetching old keys from the past.
 			if totalInserted > 0 {
diff --git a/internal/federationin/federationin.go b/internal/federationin/federationin.go
index 1eb3163..bb263d1 100644
--- a/internal/federationin/federationin.go
+++ b/internal/federationin/federationin.go
@@ -28,11 +28,13 @@ import (
 	"time"
 
 	coredb "github.com/google/exposure-notifications-server/internal/database"
+	"github.com/google/exposure-notifications-server/internal/federationin/database"
+	"github.com/google/exposure-notifications-server/internal/federationin/model"
 	"github.com/google/exposure-notifications-server/internal/logging"
 	"github.com/google/exposure-notifications-server/internal/metrics"
 	"github.com/google/exposure-notifications-server/internal/pb"
-	"github.com/google/exposure-notifications-server/internal/publish/database"
-	"github.com/google/exposure-notifications-server/internal/publish/model"
+	publishdb "github.com/google/exposure-notifications-server/internal/publish/database"
+	publishmodel "github.com/google/exposure-notifications-server/internal/publish/model"
 	"github.com/google/exposure-notifications-server/internal/serverenv"
 
 	"google.golang.org/api/idtoken"
@@ -46,13 +48,13 @@ const (
 )
 
 var (
-	fetchBatchSize = database.InsertExposuresBatchSize
+	fetchBatchSize = publishdb.InsertExposuresBatchSize
 )
 
 type (
 	fetchFn               func(context.Context, *pb.FederationFetchRequest, ...grpc.CallOption) (*pb.FederationFetchResponse, error)
-	insertExposuresFn     func(context.Context, []*model.Exposure) error
-	startFederationSyncFn func(context.Context, *coredb.FederationInQuery, time.Time) (int64, coredb.FinalizeSyncFn, error)
+	insertExposuresFn     func(context.Context, []*publishmodel.Exposure) error
+	startFederationSyncFn func(context.Context, *model.FederationInQuery, time.Time) (int64, database.FinalizeSyncFn, error)
 )
 
 type pullDependencies struct {
@@ -66,16 +68,16 @@ type pullDependencies struct {
 func NewHandler(env *serverenv.ServerEnv, config *Config) http.Handler {
 	return &handler{
 		env:       env,
-		db:        env.Database(),
-		publishdb: database.New(env.Database()),
+		db:        database.New(env.Database()),
+		publishdb: publishdb.New(env.Database()),
 		config:    config,
 	}
 }
 
 type handler struct {
 	env       *serverenv.ServerEnv
-	db        *coredb.DB
-	publishdb *database.PublishDB
+	db        *database.FederationInDB
+	publishdb *publishdb.PublishDB
 	config    *Config
 }
 
@@ -188,7 +190,7 @@ func (h *handler) ServeHTTP(w http.ResponseWriter, r *http.Request) {
 	}
 }
 
-func pull(ctx context.Context, metrics metrics.Exporter, deps pullDependencies, q *coredb.FederationInQuery, batchStart time.Time, truncateWindow time.Duration) error {
+func pull(ctx context.Context, metrics metrics.Exporter, deps pullDependencies, q *model.FederationInQuery, batchStart time.Time, truncateWindow time.Duration) error {
 	logger := logging.FromContext(ctx)
 	logger.Infof("Processing query %q", q.QueryID)
 
@@ -209,7 +211,7 @@ func pull(ctx context.Context, metrics metrics.Exporter, deps pullDependencies,
 		logger.Infof("Inserted %d keys", total)
 	}()
 
-	createdAt := model.TruncateWindow(batchStart, truncateWindow)
+	createdAt := publishmodel.TruncateWindow(batchStart, truncateWindow)
 	partial := true
 	for partial {
 
@@ -225,8 +227,8 @@ func pull(ctx context.Context, metrics metrics.Exporter, deps pullDependencies,
 			maxTimestamp = responseTimestamp
 		}
 
-		// Loop through the result set, storing in database.
-		var exposures []*model.Exposure
+		// Loop through the result set, storing in publishdb.
+		var exposures []*publishmodel.Exposure
 		for _, ctr := range response.Response {
 
 			var upperRegions []string
@@ -238,12 +240,12 @@ func pull(ctx context.Context, metrics metrics.Exporter, deps pullDependencies,
 			for _, cti := range ctr.ContactTracingInfo {
 				for _, key := range cti.ExposureKeys {
 
-					if cti.TransmissionRisk < model.MinTransmissionRisk || cti.TransmissionRisk > model.MaxTransmissionRisk {
+					if cti.TransmissionRisk < publishmodel.MinTransmissionRisk || cti.TransmissionRisk > publishmodel.MaxTransmissionRisk {
 						logger.Errorf("invalid transmission risk %v - dropping record.", cti.TransmissionRisk)
 						continue
 					}
 
-					exposures = append(exposures, &model.Exposure{
+					exposures = append(exposures, &publishmodel.Exposure{
 						TransmissionRisk: int(cti.TransmissionRisk),
 						ExposureKey:      key.ExposureKey,
 						Regions:          upperRegions,
diff --git a/internal/database/federation_model.go b/internal/federationin/model/federation_model.go
similarity index 99%
rename from internal/database/federation_model.go
rename to internal/federationin/model/federation_model.go
index fee9372..f34d88f 100644
--- a/internal/database/federation_model.go
+++ b/internal/federationin/model/federation_model.go
@@ -12,7 +12,7 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 
-package database
+package model
 
 import (
 	"time"
diff --git a/internal/database/federationout.go b/internal/federationout/database/federationout.go
similarity index 73%
rename from internal/database/federationout.go
rename to internal/federationout/database/federationout.go
index 6887a54..985f8c3 100644
--- a/internal/database/federationout.go
+++ b/internal/federationout/database/federationout.go
@@ -18,12 +18,24 @@ import (
 	"context"
 	"fmt"
 
+	"github.com/google/exposure-notifications-server/internal/database"
+	"github.com/google/exposure-notifications-server/internal/federationin/model"
 	pgx "github.com/jackc/pgx/v4"
 )
 
+type FederationOutDB struct {
+	db *database.DB
+}
+
+func New(db *database.DB) *FederationOutDB {
+	return &FederationOutDB{
+		db: db,
+	}
+}
+
 // AddFederationOutAuthorization adds or updates a FederationOutAuthorization record.
-func (db *DB) AddFederationOutAuthorization(ctx context.Context, auth *FederationOutAuthorization) error {
-	return db.InTx(ctx, pgx.Serializable, func(tx pgx.Tx) error {
+func (db *FederationOutDB) AddFederationOutAuthorization(ctx context.Context, auth *model.FederationOutAuthorization) error {
+	return db.db.InTx(ctx, pgx.Serializable, func(tx pgx.Tx) error {
 		q := `
 			INSERT INTO
 				FederationOutAuthorization
@@ -44,8 +56,8 @@ func (db *DB) AddFederationOutAuthorization(ctx context.Context, auth *Federatio
 }
 
 // GetFederationOutAuthorization returns a FederationOutAuthorization record, or ErrNotFound if not found.
-func (db *DB) GetFederationOutAuthorization(ctx context.Context, issuer, subject string) (*FederationOutAuthorization, error) {
-	conn, err := db.Pool.Acquire(ctx)
+func (db *FederationOutDB) GetFederationOutAuthorization(ctx context.Context, issuer, subject string) (*model.FederationOutAuthorization, error) {
+	conn, err := db.db.Pool.Acquire(ctx)
 	if err != nil {
 		return nil, fmt.Errorf("acquiring connection: %w", err)
 	}
@@ -62,10 +74,10 @@ func (db *DB) GetFederationOutAuthorization(ctx context.Context, issuer, subject
 			oidc_subject = $2
 		LIMIT 1
 		`, issuer, subject)
-	auth := FederationOutAuthorization{}
+	auth := model.FederationOutAuthorization{}
 	if err := row.Scan(&auth.Issuer, &auth.Subject, &auth.Audience, &auth.Note, &auth.IncludeRegions, &auth.ExcludeRegions); err != nil {
 		if err == pgx.ErrNoRows {
-			return nil, ErrNotFound
+			return nil, database.ErrNotFound
 		}
 		return nil, fmt.Errorf("scanning results: %w", err)
 	}
diff --git a/internal/federationout/federationout.go b/internal/federationout/federationout.go
index b4c8817..ee3d4ac 100644
--- a/internal/federationout/federationout.go
+++ b/internal/federationout/federationout.go
@@ -23,12 +23,15 @@ import (
 	"strings"
 	"time"
 
+	"github.com/google/exposure-notifications-server/internal/federationin/model"
 	"github.com/google/exposure-notifications-server/internal/logging"
 	"github.com/google/exposure-notifications-server/internal/pb"
 
 	coredb "github.com/google/exposure-notifications-server/internal/database"
-	"github.com/google/exposure-notifications-server/internal/publish/database"
-	"github.com/google/exposure-notifications-server/internal/publish/model"
+	"github.com/google/exposure-notifications-server/internal/federationout/database"
+	publishdb "github.com/google/exposure-notifications-server/internal/publish/database"
+	publishmodel "github.com/google/exposure-notifications-server/internal/publish/model"
+
 	"github.com/google/exposure-notifications-server/internal/serverenv"
 	"google.golang.org/api/idtoken"
 	"google.golang.org/grpc"
@@ -45,22 +48,22 @@ const (
 // Compile time assert that this server implements the required grpc interface.
 var _ pb.FederationServer = (*Server)(nil)
 
-type iterateExposuresFunc func(context.Context, database.IterateExposuresCriteria, func(*model.Exposure) error) (string, error)
+type iterateExposuresFunc func(context.Context, publishdb.IterateExposuresCriteria, func(*publishmodel.Exposure) error) (string, error)
 
 // NewServer builds a new FederationServer.
 func NewServer(env *serverenv.ServerEnv, config *Config) pb.FederationServer {
 	return &Server{
 		env:       env,
-		db:        env.Database(),
-		publishdb: database.New(env.Database()),
+		db:        database.New(env.Database()),
+		publishdb: publishdb.New(env.Database()),
 		config:    config,
 	}
 }
 
 type Server struct {
 	env       *serverenv.ServerEnv
-	db        *coredb.DB
-	publishdb *database.PublishDB
+	db        *database.FederationOutDB
+	publishdb *publishdb.PublishDB
 	config    *Config
 }
 
@@ -71,7 +74,7 @@ func (s Server) Fetch(ctx context.Context, req *pb.FederationFetchRequest) (*pb.
 	ctx, cancel := context.WithTimeout(ctx, s.config.Timeout)
 	defer cancel()
 	logger := logging.FromContext(ctx)
-	response, err := s.fetch(ctx, req, s.publishdb.IterateExposures, model.TruncateWindow(time.Now(), s.config.TruncateWindow)) // Don't fetch the current window, which isn't complete yet. TODO(squee1945): should I double this for safety?
+	response, err := s.fetch(ctx, req, s.publishdb.IterateExposures, publishmodel.TruncateWindow(time.Now(), s.config.TruncateWindow)) // Don't fetch the current window, which isn't complete yet. TODO(squee1945): should I double this for safety?
 	if err != nil {
 		s.env.MetricsExporter(ctx).WriteInt("federation-fetch-failed", true, 1)
 		logger.Errorf("Fetch error: %v", err)
@@ -96,14 +99,14 @@ func (s Server) fetch(ctx context.Context, req *pb.FederationFetchRequest, itFun
 	logger.Infof("Processing client request %#v", req)
 
 	// If there is a FederationAuthorization on the context, set the query to operate within its limits.
-	if auth, ok := ctx.Value(authKey{}).(*coredb.FederationOutAuthorization); ok {
+	if auth, ok := ctx.Value(authKey{}).(*model.FederationOutAuthorization); ok {
 		// For included regions, we INTERSECT the requested included regions with the configured included regions.
 		req.RegionIdentifiers = intersect(req.RegionIdentifiers, auth.IncludeRegions)
 		// For excluded regions, we UNION the the requested excluded regions with the configured excluded regions.
 		req.ExcludeRegionIdentifiers = union(req.ExcludeRegionIdentifiers, auth.ExcludeRegions)
 	}
 
-	criteria := database.IterateExposuresCriteria{
+	criteria := publishdb.IterateExposuresCriteria{
 		IncludeRegions:      req.RegionIdentifiers,
 		ExcludeRegions:      req.ExcludeRegionIdentifiers,
 		SinceTimestamp:      time.Unix(req.LastFetchResponseKeyTimestamp, 0),
@@ -130,7 +133,7 @@ func (s Server) fetch(ctx context.Context, req *pb.FederationFetchRequest, itFun
 	ctiMap := map[string]*pb.ContactTracingInfo{}     // local index into the response being assembled; keys on unique set of (ctrMap key, transmissionRisk, verificationAuthorityName)
 	response := &pb.FederationFetchResponse{}
 	count := 0
-	cursor, err := itFunc(ctx, criteria, func(inf *model.Exposure) error {
+	cursor, err := itFunc(ctx, criteria, func(inf *publishmodel.Exposure) error {
 		// If the diagnosis key is empty, it's malformed, so skip it.
 		if len(inf.ExposureKey) == 0 {
 			logger.Debugf("Exposure %s missing ExposureKey, skipping.", inf.ExposureKey)
diff --git a/terraform/database.tf b/terraform/database.tf
index f8cdf10..4f647ad 100644
--- a/terraform/database.tf
+++ b/terraform/database.tf
@@ -50,6 +50,12 @@ resource "google_sql_database_instance" "db-inst" {
   lifecycle {
     # This prevents accidential deletion of the database.
     prevent_destroy = true
+
+    # Earlier versions of the database had a different name, and its not
+    # possible to rename Cloud SQL instances.
+    ignore_changes = [
+      name,
+    ]
   }
 
   depends_on = [
diff --git a/tools/federationin-query/main.go b/tools/federationin-query/main.go
index 6bae60e..e0978d6 100644
--- a/tools/federationin-query/main.go
+++ b/tools/federationin-query/main.go
@@ -22,8 +22,10 @@ import (
 	"regexp"
 	"time"
 
-	"github.com/google/exposure-notifications-server/internal/database"
+	coredb "github.com/google/exposure-notifications-server/internal/database"
 	"github.com/google/exposure-notifications-server/internal/federationin"
+	"github.com/google/exposure-notifications-server/internal/federationin/database"
+	"github.com/google/exposure-notifications-server/internal/federationin/model"
 	cflag "github.com/google/exposure-notifications-server/internal/flag"
 	"github.com/kelseyhightower/envconfig"
 )
@@ -75,19 +77,20 @@ func main() {
 	}
 
 	ctx := context.Background()
-	var config database.Config
+	var config coredb.Config
 	err := envconfig.Process("database", &config)
 	if err != nil {
 		log.Fatalf("error loading environment variables: %v", err)
 	}
 
-	db, err := database.NewFromEnv(ctx, &config)
+	coredb, err := coredb.NewFromEnv(ctx, &config)
 	if err != nil {
 		log.Fatalf("unable to connect to database: %v", err)
 	}
-	defer db.Close(ctx)
+	defer coredb.Close(ctx)
+	db := database.New(coredb)
 
-	query := &database.FederationInQuery{
+	query := &model.FederationInQuery{
 		QueryID:        *queryID,
 		ServerAddr:     *serverAddr,
 		Audience:       *audience,
diff --git a/tools/federationout-authorization/main.go b/tools/federationout-authorization/main.go
index 9158c1a..d9605b9 100644
--- a/tools/federationout-authorization/main.go
+++ b/tools/federationout-authorization/main.go
@@ -21,8 +21,10 @@ import (
 	"flag"
 	"log"
 
-	"github.com/google/exposure-notifications-server/internal/database"
+	coredb "github.com/google/exposure-notifications-server/internal/database"
 	"github.com/google/exposure-notifications-server/internal/federationin"
+	"github.com/google/exposure-notifications-server/internal/federationin/model"
+	"github.com/google/exposure-notifications-server/internal/federationout/database"
 	cflag "github.com/google/exposure-notifications-server/internal/flag"
 	"github.com/kelseyhightower/envconfig"
 )
@@ -68,19 +70,20 @@ func main() {
 	}
 
 	ctx := context.Background()
-	var config database.Config
+	var config coredb.Config
 	err := envconfig.Process("database", &config)
 	if err != nil {
 		log.Fatalf("error loading environment variables: %v", err)
 	}
 
-	db, err := database.NewFromEnv(ctx, &config)
+	coredb, err := coredb.NewFromEnv(ctx, &config)
 	if err != nil {
 		log.Fatalf("unable to connect to database: %v", err)
 	}
-	defer db.Close(ctx)
+	defer coredb.Close(ctx)
+	db := database.New(coredb)
 
-	auth := &database.FederationOutAuthorization{
+	auth := &model.FederationOutAuthorization{
 		Issuer:         defaultIssuer, // Authorization interceptor currently only supports defaultIssuer.
 		Subject:        *subject,
 		Audience:       *audience,


Test patch:
diff --git a/internal/authorizedapp/database/authorized_app_test.go b/internal/authorizedapp/database/authorized_app_test.go
index 4781166..2e55520 100644
--- a/internal/authorizedapp/database/authorized_app_test.go
+++ b/internal/authorizedapp/database/authorized_app_test.go
@@ -52,7 +52,7 @@ func TestMain(m *testing.M) {
 
 	if os.Getenv("DB_USER") != "" {
 		var err error
-		testDB, err = coredb.CreateTestDB(ctx)
+		testDB, err = coredb.CreateTestDB(ctx, "authorizedapp")
 		if err != nil {
 			log.Fatalf("creating test DB: %v", err)
 		}
diff --git a/internal/database/database_test_util.go b/internal/database/database_test_util.go
index 14fab99..1be69a3 100644
--- a/internal/database/database_test_util.go
+++ b/internal/database/database_test_util.go
@@ -39,7 +39,7 @@ func TestMain(m *testing.M) {
 
 	if os.Getenv("DB_USER") != "" {
 		var err error
-		testDB, err = CreateTestDB(ctx)
+		testDB, err = CreateTestDB(ctx, "core")
 		if err != nil {
 			log.Fatalf("creating test DB: %v", err)
 		}
@@ -50,8 +50,8 @@ func TestMain(m *testing.M) {
 // openTestDB connects to the Postgres server specified by the DB_XXX environment
 // variables, creates an empty test database on it, and returns a *DB connected
 // to that database.
-func CreateTestDB(ctx context.Context) (*DB, error) {
-	const testDBName = "exposure-server-test"
+func CreateTestDB(ctx context.Context, suffix string) (*DB, error) {
+	var testDBName = fmt.Sprintf("exposure-server-test-%s", suffix)
 
 	// Connect to the default database to create the test database.
 	var config Config
diff --git a/internal/export/database/export_test.go b/internal/export/database/export_test.go
index c1f1d92..28517b8 100644
--- a/internal/export/database/export_test.go
+++ b/internal/export/database/export_test.go
@@ -38,7 +38,7 @@ func TestMain(m *testing.M) {
 
 	if os.Getenv("DB_USER") != "" {
 		var err error
-		testDB, err = database.CreateTestDB(ctx)
+		testDB, err = database.CreateTestDB(ctx, "export")
 		if err != nil {
 			log.Fatalf("creating test DB: %v", err)
 		}
diff --git a/internal/database/federationin_test.go b/internal/federationin/database/federationin_test.go
similarity index 66%
rename from internal/database/federationin_test.go
rename to internal/federationin/database/federationin_test.go
index d62de9f..05f8fde 100644
--- a/internal/database/federationin_test.go
+++ b/internal/federationin/database/federationin_test.go
@@ -17,23 +17,43 @@ package database
 import (
 	"context"
 	"errors"
+	"log"
+	"os"
 	"testing"
 	"time"
 
+	coredb "github.com/google/exposure-notifications-server/internal/database"
+	"github.com/google/exposure-notifications-server/internal/federationin/model"
+
 	"github.com/google/go-cmp/cmp"
 	"github.com/google/go-cmp/cmp/cmpopts"
 )
 
+var testDB *coredb.DB
+
+func TestMain(m *testing.M) {
+	ctx := context.Background()
+
+	if os.Getenv("DB_USER") != "" {
+		var err error
+		testDB, err = coredb.CreateTestDB(ctx, "federationin")
+		if err != nil {
+			log.Fatalf("creating test DB: %v", err)
+		}
+	}
+	os.Exit(m.Run())
+}
+
 // TestFederationIn tests functions operating over FederationInQuery, FederationInSync.
 func TestFederationIn(t *testing.T) {
 	if testDB == nil {
 		t.Skip("no test DB")
 	}
-	defer ResetTestDB(t, testDB)
+	defer coredb.ResetTestDB(t, testDB)
 	ctx := context.Background()
 
 	ts := time.Date(2020, 5, 6, 0, 0, 0, 0, time.UTC)
-	want := &FederationInQuery{
+	want := &model.FederationInQuery{
 		QueryID:        "qid",
 		ServerAddr:     "addr",
 		IncludeRegions: []string{"MX"},
@@ -41,15 +61,15 @@ func TestFederationIn(t *testing.T) {
 		LastTimestamp:  ts,
 	}
 	// GetFederationQuery should fail if not found.
-	if _, err := testDB.GetFederationInQuery(ctx, want.QueryID); !errors.Is(err, ErrNotFound) {
+	if _, err := New(testDB).GetFederationInQuery(ctx, want.QueryID); !errors.Is(err, coredb.ErrNotFound) {
 		t.Errorf("got %v, want ErrNotFound", err)
 	}
 
 	// Add a query, then get it.
-	if err := testDB.AddFederationInQuery(ctx, want); err != nil {
+	if err := New(testDB).AddFederationInQuery(ctx, want); err != nil {
 		t.Fatal(err)
 	}
-	got, err := testDB.GetFederationInQuery(ctx, want.QueryID)
+	got, err := New(testDB).GetFederationInQuery(ctx, want.QueryID)
 	if err != nil {
 		t.Fatal(err)
 	}
@@ -59,10 +79,10 @@ func TestFederationIn(t *testing.T) {
 
 	// AddFederationQuery should overwrite.
 	want.ServerAddr = "addr2"
-	if err := testDB.AddFederationInQuery(ctx, want); err != nil {
+	if err := New(testDB).AddFederationInQuery(ctx, want); err != nil {
 		t.Fatal(err)
 	}
-	got, err = testDB.GetFederationInQuery(ctx, want.QueryID)
+	got, err = New(testDB).GetFederationInQuery(ctx, want.QueryID)
 	if err != nil {
 		t.Fatal(err)
 	}
@@ -71,21 +91,21 @@ func TestFederationIn(t *testing.T) {
 	}
 
 	// GetFederationSync should fail if not found.
-	if _, err := testDB.GetFederationInSync(ctx, 1); !errors.Is(err, ErrNotFound) {
+	if _, err := New(testDB).GetFederationInSync(ctx, 1); !errors.Is(err, coredb.ErrNotFound) {
 		t.Errorf("got %v, want ErrNotFound", err)
 	}
 
 	// Start a sync.
 	now := time.Now().Truncate(time.Microsecond)
-	syncID, finalize, err := testDB.StartFederationInSync(ctx, want, now)
+	syncID, finalize, err := New(testDB).StartFederationInSync(ctx, want, now)
 	if err != nil {
 		t.Fatal(err)
 	}
-	gotSync, err := testDB.GetFederationInSync(ctx, syncID)
+	gotSync, err := New(testDB).GetFederationInSync(ctx, syncID)
 	if err != nil {
 		t.Fatal(err)
 	}
-	wantSync := &FederationInSync{
+	wantSync := &model.FederationInSync{
 		SyncID:  syncID,
 		QueryID: want.QueryID,
 		Started: now,
@@ -103,7 +123,7 @@ func TestFederationIn(t *testing.T) {
 	if err := finalize(wantSync.MaxTimestamp, wantSync.Insertions); err != nil {
 		t.Fatal(err)
 	}
-	gotSync, err = testDB.GetFederationInSync(ctx, syncID)
+	gotSync, err = New(testDB).GetFederationInSync(ctx, syncID)
 	if err != nil {
 		t.Fatal(err)
 	}
diff --git a/internal/federationin/federationin_test.go b/internal/federationin/federationin_test.go
index 8daef98..1121c62 100644
--- a/internal/federationin/federationin_test.go
+++ b/internal/federationin/federationin_test.go
@@ -19,8 +19,9 @@ import (
 	"testing"
 	"time"
 
-	"github.com/google/exposure-notifications-server/internal/database"
-	"github.com/google/exposure-notifications-server/internal/publish/model"
+	"github.com/google/exposure-notifications-server/internal/federationin/database"
+	"github.com/google/exposure-notifications-server/internal/federationin/model"
+	publishmodel "github.com/google/exposure-notifications-server/internal/publish/model"
 
 	"github.com/google/exposure-notifications-server/internal/metrics"
 	"github.com/google/exposure-notifications-server/internal/pb"
@@ -39,8 +40,8 @@ var (
 	ddd = &pb.ExposureKey{ExposureKey: []byte("ddd"), IntervalNumber: 4}
 )
 
-// makeRemoteExposure returns a mock model.Exposure with LocalProvenance=false.
-func makeRemoteExposure(diagKey *pb.ExposureKey, diagStatus int, regions ...string) *model.Exposure {
+// makeRemoteExposure returns a mock publishmodel.Exposure with LocalProvenance=false.
+func makeRemoteExposure(diagKey *pb.ExposureKey, diagStatus int, regions ...string) *publishmodel.Exposure {
 	inf := makeExposure(diagKey, diagStatus, regions...)
 	inf.LocalProvenance = false
 	inf.FederationSyncID = syncID
@@ -66,10 +67,10 @@ func (r *remoteFetchServer) fetch(ctx context.Context, req *pb.FederationFetchRe
 
 // publishDB mocks the database, recording exposure insertions.
 type publishDB struct {
-	exposures []*model.Exposure
+	exposures []*publishmodel.Exposure
 }
 
-func (idb *publishDB) insertExposures(ctx context.Context, exposures []*model.Exposure) error {
+func (idb *publishDB) insertExposures(ctx context.Context, exposures []*publishmodel.Exposure) error {
 	idb.exposures = append(idb.exposures, exposures...)
 	return nil
 }
@@ -83,7 +84,7 @@ type syncDB struct {
 	totalInserted int
 }
 
-func (sdb *syncDB) startFederationSync(ctx context.Context, query *database.FederationInQuery, start time.Time) (int64, database.FinalizeSyncFn, error) {
+func (sdb *syncDB) startFederationSync(ctx context.Context, query *model.FederationInQuery, start time.Time) (int64, database.FinalizeSyncFn, error) {
 	sdb.syncStarted = true
 	timerStart := time.Now()
 	return syncID, func(maxTimestamp time.Time, totalInserted int) error {
@@ -101,7 +102,7 @@ func TestFederationPull(t *testing.T) {
 		name             string
 		batchSize        int
 		fetchResponses   []*pb.FederationFetchResponse
-		wantExposures    []*model.Exposure
+		wantExposures    []*publishmodel.Exposure
 		wantTokens       []string
 		wantMaxTimestamp time.Time
 	}{
@@ -137,7 +138,7 @@ func TestFederationPull(t *testing.T) {
 					FetchResponseKeyTimestamp: 400,
 				},
 			},
-			wantExposures: []*model.Exposure{
+			wantExposures: []*publishmodel.Exposure{
 				makeRemoteExposure(aaa, 1, "US"),
 				makeRemoteExposure(bbb, 1, "US"),
 				makeRemoteExposure(ccc, 2, "CA", "US"),
@@ -173,7 +174,7 @@ func TestFederationPull(t *testing.T) {
 					FetchResponseKeyTimestamp: 400,
 				},
 			},
-			wantExposures: []*model.Exposure{
+			wantExposures: []*publishmodel.Exposure{
 				makeRemoteExposure(aaa, 1, "US"),
 				makeRemoteExposure(bbb, 1, "US"),
 			},
@@ -214,7 +215,7 @@ func TestFederationPull(t *testing.T) {
 					FetchResponseKeyTimestamp: 400,
 				},
 			},
-			wantExposures: []*model.Exposure{
+			wantExposures: []*publishmodel.Exposure{
 				makeRemoteExposure(aaa, 8, "US"),
 				makeRemoteExposure(bbb, 8, "US"),
 				makeRemoteExposure(ccc, 7, "US"),
@@ -251,7 +252,7 @@ func TestFederationPull(t *testing.T) {
 					FetchResponseKeyTimestamp: 400,
 				},
 			},
-			wantExposures: []*model.Exposure{
+			wantExposures: []*publishmodel.Exposure{
 				makeRemoteExposure(aaa, 3, "US"),
 				makeRemoteExposure(bbb, 3, "US"),
 				makeRemoteExposure(ccc, 2, "CA", "US"),
@@ -265,7 +266,7 @@ func TestFederationPull(t *testing.T) {
 	for _, tc := range testCases {
 		t.Run(tc.name, func(t *testing.T) {
 			ctx := context.Background()
-			query := &database.FederationInQuery{}
+			query := &model.FederationInQuery{}
 			remote := remoteFetchServer{responses: tc.fetchResponses}
 			idb := publishDB{}
 			sdb := syncDB{}
@@ -286,7 +287,7 @@ func TestFederationPull(t *testing.T) {
 				t.Fatalf("pull returned err=%v, want err=nil", err)
 			}
 
-			if diff := cmp.Diff(tc.wantExposures, idb.exposures, cmpopts.IgnoreFields(model.Exposure{}, "CreatedAt")); diff != "" {
+			if diff := cmp.Diff(tc.wantExposures, idb.exposures, cmpopts.IgnoreFields(publishmodel.Exposure{}, "CreatedAt")); diff != "" {
 				t.Errorf("exposures mismatch (-want +got):\n%s", diff)
 			}
 			if diff := cmp.Diff(tc.wantTokens, remote.gotTokens); diff != "" {
@@ -311,8 +312,8 @@ func TestFederationPull(t *testing.T) {
 	}
 }
 
-func makeExposure(diagKey *pb.ExposureKey, diagStatus int, regions ...string) *model.Exposure {
-	return &model.Exposure{
+func makeExposure(diagKey *pb.ExposureKey, diagStatus int, regions ...string) *publishmodel.Exposure {
+	return &publishmodel.Exposure{
 		Regions:          regions,
 		TransmissionRisk: diagStatus,
 		ExposureKey:      diagKey.ExposureKey,
diff --git a/internal/database/federationout_test.go b/internal/federationout/database/federationout_test.go
similarity index 62%
rename from internal/database/federationout_test.go
rename to internal/federationout/database/federationout_test.go
index 3f0676b..a9533d5 100644
--- a/internal/database/federationout_test.go
+++ b/internal/federationout/database/federationout_test.go
@@ -17,20 +17,40 @@ package database
 import (
 	"context"
 	"errors"
+	"log"
+	"os"
 	"testing"
 
+	coredb "github.com/google/exposure-notifications-server/internal/database"
+	"github.com/google/exposure-notifications-server/internal/federationin/model"
+
 	"github.com/google/go-cmp/cmp"
 )
 
+var testDB *coredb.DB
+
+func TestMain(m *testing.M) {
+	ctx := context.Background()
+
+	if os.Getenv("DB_USER") != "" {
+		var err error
+		testDB, err = coredb.CreateTestDB(ctx, "federationout")
+		if err != nil {
+			log.Fatalf("creating test DB: %v", err)
+		}
+	}
+	os.Exit(m.Run())
+}
+
 // TestFederationOutAuthorization tests the functions accessing the FederationOutAuthorization table.
 func TestFederationOutAuthorization(t *testing.T) {
 	if testDB == nil {
 		t.Skip("no test DB")
 	}
-	defer ResetTestDB(t, testDB)
+	defer coredb.ResetTestDB(t, testDB)
 	ctx := context.Background()
 
-	want := &FederationOutAuthorization{
+	want := &model.FederationOutAuthorization{
 		Issuer:         "iss",
 		Subject:        "sub",
 		Audience:       "aud",
@@ -40,14 +60,14 @@ func TestFederationOutAuthorization(t *testing.T) {
 	}
 
 	// GetFederationOutAuthorization should fail if not found.
-	if _, err := testDB.GetFederationOutAuthorization(ctx, want.Issuer, want.Subject); !errors.Is(err, ErrNotFound) {
+	if _, err := New(testDB).GetFederationOutAuthorization(ctx, want.Issuer, want.Subject); !errors.Is(err, coredb.ErrNotFound) {
 		t.Errorf("got %v, want ErrNotFound", err)
 	}
 	// Add a query, then get it.
-	if err := testDB.AddFederationOutAuthorization(ctx, want); err != nil {
+	if err := New(testDB).AddFederationOutAuthorization(ctx, want); err != nil {
 		t.Fatal(err)
 	}
-	got, err := testDB.GetFederationOutAuthorization(ctx, want.Issuer, want.Subject)
+	got, err := New(testDB).GetFederationOutAuthorization(ctx, want.Issuer, want.Subject)
 	if err != nil {
 		t.Fatal(err)
 	}
@@ -57,10 +77,10 @@ func TestFederationOutAuthorization(t *testing.T) {
 
 	// AddFederationOutAuthorization should overwrite.
 	want.Note = "a different note"
-	if err := testDB.AddFederationOutAuthorization(ctx, want); err != nil {
+	if err := New(testDB).AddFederationOutAuthorization(ctx, want); err != nil {
 		t.Fatal(err)
 	}
-	got, err = testDB.GetFederationOutAuthorization(ctx, want.Issuer, want.Subject)
+	got, err = New(testDB).GetFederationOutAuthorization(ctx, want.Issuer, want.Subject)
 	if err != nil {
 		t.Fatal(err)
 	}
diff --git a/internal/publish/database/exposure_test.go b/internal/publish/database/exposure_test.go
index 329bb99..0b039bf 100644
--- a/internal/publish/database/exposure_test.go
+++ b/internal/publish/database/exposure_test.go
@@ -36,7 +36,7 @@ func TestMain(m *testing.M) {
 
 	if os.Getenv("DB_USER") != "" {
 		var err error
-		testDB, err = coredb.CreateTestDB(ctx)
+		testDB, err = coredb.CreateTestDB(ctx, "publish")
 		if err != nil {
 			log.Fatalf("creating test DB: %v", err)
 		}