Repository: apigee/registry (go)

Issue:
Registry server: reuse db connections
DB connections are being opened and closed for each database action. The DB connection should be maintained for the lifetime of the server to allow connection pooling and other optimizations such as reuse of prepared statements.

Reference [sql.Open](https://pkg.go.dev/database/sql#Open):

> The returned DB is safe for concurrent use by multiple goroutines and maintains its own pool of idle connections. Thus, the Open function should be called just once. It is rarely necessary to close a DB.



Pull request:
Registry server: reuse db connections
Closes #622

Patch:
diff --git a/cmd/registry-server/main.go b/cmd/registry-server/main.go
index 8671872..6fbfc00 100644
--- a/cmd/registry-server/main.go
+++ b/cmd/registry-server/main.go
@@ -141,6 +141,7 @@ func main() {
 	signal.Notify(done, os.Interrupt, syscall.SIGTERM)
 	<-done
 	server.GracefulStop()
+	registryServer.Close()
 }
 
 func validateConfig() error {
diff --git a/server/registry/actions_apis.go b/server/registry/actions_apis.go
index 46829e7..596ed3d 100644
--- a/server/registry/actions_apis.go
+++ b/server/registry/actions_apis.go
@@ -45,7 +45,6 @@ func (s *RegistryServer) createApi(ctx context.Context, name names.Api, body *rp
 	if err != nil {
 		return nil, status.Error(codes.Unavailable, err.Error())
 	}
-	defer db.Close()
 
 	if err := name.Validate(); err != nil {
 		return nil, status.Error(codes.InvalidArgument, err.Error())
@@ -80,7 +79,6 @@ func (s *RegistryServer) DeleteApi(ctx context.Context, req *rpc.DeleteApiReques
 	if err != nil {
 		return nil, status.Error(codes.Unavailable, err.Error())
 	}
-	defer db.Close()
 
 	name, err := names.ParseApi(req.GetName())
 	if err != nil {
@@ -106,7 +104,6 @@ func (s *RegistryServer) GetApi(ctx context.Context, req *rpc.GetApiRequest) (*r
 	if err != nil {
 		return nil, status.Error(codes.Unavailable, err.Error())
 	}
-	defer db.Close()
 
 	name, err := names.ParseApi(req.GetName())
 	if err != nil {
@@ -132,7 +129,6 @@ func (s *RegistryServer) ListApis(ctx context.Context, req *rpc.ListApisRequest)
 	if err != nil {
 		return nil, status.Error(codes.Unavailable, err.Error())
 	}
-	defer db.Close()
 
 	if req.GetPageSize() < 0 {
 		return nil, status.Errorf(codes.InvalidArgument, "invalid page_size %d: must not be negative", req.GetPageSize())
@@ -198,7 +194,6 @@ func (s *RegistryServer) UpdateApi(ctx context.Context, req *rpc.UpdateApiReques
 	if err != nil {
 		return nil, status.Error(codes.Unavailable, err.Error())
 	}
-	defer db.Close()
 
 	// TODO: the Get/Update/Save calls below should be a transaction.
 	api, err := db.GetApi(ctx, name)
diff --git a/server/registry/actions_artifacts.go b/server/registry/actions_artifacts.go
index 1ab7ef1..a31a6a0 100644
--- a/server/registry/actions_artifacts.go
+++ b/server/registry/actions_artifacts.go
@@ -55,7 +55,6 @@ func (s *RegistryServer) CreateArtifact(ctx context.Context, req *rpc.CreateArti
 	if err != nil {
 		return nil, status.Error(codes.Unavailable, err.Error())
 	}
-	defer db.Close()
 
 	if req.GetArtifact() == nil {
 		return nil, status.Errorf(codes.InvalidArgument, "invalid artifact %+v: body must be provided", req.GetArtifact())
@@ -123,7 +122,6 @@ func (s *RegistryServer) DeleteArtifact(ctx context.Context, req *rpc.DeleteArti
 	if err != nil {
 		return nil, status.Error(codes.Unavailable, err.Error())
 	}
-	defer db.Close()
 
 	name, err := names.ParseArtifact(req.GetName())
 	if err != nil {
@@ -149,7 +147,6 @@ func (s *RegistryServer) GetArtifact(ctx context.Context, req *rpc.GetArtifactRe
 	if err != nil {
 		return nil, status.Error(codes.Unavailable, err.Error())
 	}
-	defer db.Close()
 
 	name, err := names.ParseArtifact(req.GetName())
 	if err != nil {
@@ -170,7 +167,6 @@ func (s *RegistryServer) GetArtifactContents(ctx context.Context, req *rpc.GetAr
 	if err != nil {
 		return nil, status.Error(codes.Unavailable, err.Error())
 	}
-	defer db.Close()
 
 	name, err := names.ParseArtifact(req.GetName())
 	if err != nil {
@@ -207,7 +203,6 @@ func (s *RegistryServer) ListArtifacts(ctx context.Context, req *rpc.ListArtifac
 	if err != nil {
 		return nil, status.Error(codes.Unavailable, err.Error())
 	}
-	defer db.Close()
 
 	if req.GetPageSize() < 0 {
 		return nil, status.Errorf(codes.InvalidArgument, "invalid page_size %d: must not be negative", req.GetPageSize())
@@ -277,7 +272,6 @@ func (s *RegistryServer) ReplaceArtifact(ctx context.Context, req *rpc.ReplaceAr
 	if err != nil {
 		return nil, status.Error(codes.Unavailable, err.Error())
 	}
-	defer db.Close()
 
 	name, err := names.ParseArtifact(req.Artifact.GetName())
 	if err != nil {
diff --git a/server/registry/actions_deployment_revisions.go b/server/registry/actions_deployment_revisions.go
index 27ecdca..d20e115 100644
--- a/server/registry/actions_deployment_revisions.go
+++ b/server/registry/actions_deployment_revisions.go
@@ -31,7 +31,6 @@ func (s *RegistryServer) ListApiDeploymentRevisions(ctx context.Context, req *rp
 	if err != nil {
 		return nil, status.Error(codes.Unavailable, err.Error())
 	}
-	defer db.Close()
 
 	if req.GetPageSize() < 0 {
 		return nil, status.Errorf(codes.InvalidArgument, "invalid page_size %d: must not be negative", req.GetPageSize())
@@ -75,7 +74,6 @@ func (s *RegistryServer) DeleteApiDeploymentRevision(ctx context.Context, req *r
 	if err != nil {
 		return nil, status.Error(codes.Unavailable, err.Error())
 	}
-	defer db.Close()
 
 	name, err := names.ParseDeploymentRevision(req.GetName())
 	if err != nil {
@@ -116,7 +114,6 @@ func (s *RegistryServer) TagApiDeploymentRevision(ctx context.Context, req *rpc.
 	if err != nil {
 		return nil, status.Error(codes.Unavailable, err.Error())
 	}
-	defer db.Close()
 
 	if req.GetTag() == "" {
 		return nil, status.Errorf(codes.InvalidArgument, "invalid tag %q, must not be empty", req.GetTag())
@@ -162,7 +159,6 @@ func (s *RegistryServer) RollbackApiDeployment(ctx context.Context, req *rpc.Rol
 	if err != nil {
 		return nil, status.Error(codes.Unavailable, err.Error())
 	}
-	defer db.Close()
 
 	if req.GetRevisionId() == "" {
 		return nil, status.Errorf(codes.InvalidArgument, "invalid revision ID %q, must not be empty", req.GetRevisionId())
diff --git a/server/registry/actions_deployments.go b/server/registry/actions_deployments.go
index f98e8b9..d1d05c6 100644
--- a/server/registry/actions_deployments.go
+++ b/server/registry/actions_deployments.go
@@ -48,7 +48,6 @@ func (s *RegistryServer) createDeployment(ctx context.Context, name names.Deploy
 	if err != nil {
 		return nil, status.Error(codes.Unavailable, err.Error())
 	}
-	defer db.Close()
 
 	if _, err := db.GetDeployment(ctx, name); err == nil {
 		return nil, status.Errorf(codes.AlreadyExists, "API deployment %q already exists", name)
@@ -89,7 +88,6 @@ func (s *RegistryServer) DeleteApiDeployment(ctx context.Context, req *rpc.Delet
 	if err != nil {
 		return nil, status.Error(codes.Unavailable, err.Error())
 	}
-	defer db.Close()
 
 	name, err := names.ParseDeployment(req.GetName())
 	if err != nil {
@@ -125,7 +123,6 @@ func (s *RegistryServer) getApiDeployment(ctx context.Context, name names.Deploy
 	if err != nil {
 		return nil, status.Error(codes.Unavailable, err.Error())
 	}
-	defer db.Close()
 
 	deployment, err := db.GetDeployment(ctx, name)
 	if err != nil {
@@ -145,7 +142,6 @@ func (s *RegistryServer) getApiDeploymentRevision(ctx context.Context, name name
 	if err != nil {
 		return nil, status.Error(codes.Unavailable, err.Error())
 	}
-	defer db.Close()
 
 	revision, err := db.GetDeploymentRevision(ctx, name)
 	if err != nil {
@@ -166,7 +162,6 @@ func (s *RegistryServer) ListApiDeployments(ctx context.Context, req *rpc.ListAp
 	if err != nil {
 		return nil, status.Error(codes.Unavailable, err.Error())
 	}
-	defer db.Close()
 
 	if req.GetPageSize() < 0 {
 		return nil, status.Errorf(codes.InvalidArgument, "invalid page_size %d: must not be negative", req.GetPageSize())
@@ -213,7 +208,6 @@ func (s *RegistryServer) UpdateApiDeployment(ctx context.Context, req *rpc.Updat
 	if err != nil {
 		return nil, status.Error(codes.Unavailable, err.Error())
 	}
-	defer db.Close()
 
 	if req.GetApiDeployment() == nil {
 		return nil, status.Errorf(codes.InvalidArgument, "invalid api_deployment %+v: body must be provided", req.GetApiDeployment())
diff --git a/server/registry/actions_migrate_database.go b/server/registry/actions_migrate_database.go
index a1d7e39..095578b 100644
--- a/server/registry/actions_migrate_database.go
+++ b/server/registry/actions_migrate_database.go
@@ -33,7 +33,6 @@ func (s *RegistryServer) MigrateDatabase(ctx context.Context, req *rpc.MigrateDa
 	if err != nil {
 		return nil, status.Error(codes.Unavailable, err.Error())
 	}
-	defer db.Close()
 
 	err = db.Migrate(ctx)
 	if err != nil {
diff --git a/server/registry/actions_projects.go b/server/registry/actions_projects.go
index 7fd7e00..b4472a6 100644
--- a/server/registry/actions_projects.go
+++ b/server/registry/actions_projects.go
@@ -40,7 +40,6 @@ func (s *RegistryServer) createProject(ctx context.Context, name names.Project,
 	if err != nil {
 		return nil, status.Error(codes.Unavailable, err.Error())
 	}
-	defer db.Close()
 
 	if err := name.Validate(); err != nil {
 		return nil, status.Error(codes.InvalidArgument, err.Error())
@@ -61,7 +60,6 @@ func (s *RegistryServer) DeleteProject(ctx context.Context, req *rpc.DeleteProje
 	if err != nil {
 		return nil, status.Error(codes.Unavailable, err.Error())
 	}
-	defer db.Close()
 
 	name, err := names.ParseProject(req.GetName())
 	if err != nil {
@@ -87,7 +85,6 @@ func (s *RegistryServer) GetProject(ctx context.Context, req *rpc.GetProjectRequ
 	if err != nil {
 		return nil, status.Error(codes.Unavailable, err.Error())
 	}
-	defer db.Close()
 
 	name, err := names.ParseProject(req.GetName())
 	if err != nil {
@@ -108,7 +105,6 @@ func (s *RegistryServer) ListProjects(ctx context.Context, req *rpc.ListProjects
 	if err != nil {
 		return nil, status.Error(codes.Unavailable, err.Error())
 	}
-	defer db.Close()
 
 	if req.GetPageSize() < 0 {
 		return nil, status.Errorf(codes.InvalidArgument, "invalid page_size %d: must not be negative", req.GetPageSize())
@@ -166,7 +162,6 @@ func (s *RegistryServer) UpdateProject(ctx context.Context, req *rpc.UpdateProje
 	if err != nil {
 		return nil, status.Error(codes.Unavailable, err.Error())
 	}
-	defer db.Close()
 
 	// TODO: the Get/Update/Save calls below should be a transaction.
 	project, err := db.GetProject(ctx, name)
diff --git a/server/registry/actions_spec_revisions.go b/server/registry/actions_spec_revisions.go
index f8630f0..a5bd436 100644
--- a/server/registry/actions_spec_revisions.go
+++ b/server/registry/actions_spec_revisions.go
@@ -31,7 +31,6 @@ func (s *RegistryServer) ListApiSpecRevisions(ctx context.Context, req *rpc.List
 	if err != nil {
 		return nil, status.Error(codes.Unavailable, err.Error())
 	}
-	defer db.Close()
 
 	if req.GetPageSize() < 0 {
 		return nil, status.Errorf(codes.InvalidArgument, "invalid page_size %d: must not be negative", req.GetPageSize())
@@ -75,7 +74,6 @@ func (s *RegistryServer) DeleteApiSpecRevision(ctx context.Context, req *rpc.Del
 	if err != nil {
 		return nil, status.Error(codes.Unavailable, err.Error())
 	}
-	defer db.Close()
 
 	name, err := names.ParseSpecRevision(req.GetName())
 	if err != nil {
@@ -116,7 +114,6 @@ func (s *RegistryServer) TagApiSpecRevision(ctx context.Context, req *rpc.TagApi
 	if err != nil {
 		return nil, status.Error(codes.Unavailable, err.Error())
 	}
-	defer db.Close()
 
 	if req.GetTag() == "" {
 		return nil, status.Errorf(codes.InvalidArgument, "invalid tag %q, must not be empty", req.GetTag())
@@ -162,7 +159,6 @@ func (s *RegistryServer) RollbackApiSpec(ctx context.Context, req *rpc.RollbackA
 	if err != nil {
 		return nil, status.Error(codes.Unavailable, err.Error())
 	}
-	defer db.Close()
 
 	if req.GetRevisionId() == "" {
 		return nil, status.Errorf(codes.InvalidArgument, "invalid revision ID %q, must not be empty", req.GetRevisionId())
diff --git a/server/registry/actions_specs.go b/server/registry/actions_specs.go
index 58bc57c..5916b7a 100644
--- a/server/registry/actions_specs.go
+++ b/server/registry/actions_specs.go
@@ -52,7 +52,6 @@ func (s *RegistryServer) createSpec(ctx context.Context, name names.Spec, body *
 	if err != nil {
 		return nil, status.Error(codes.Unavailable, err.Error())
 	}
-	defer db.Close()
 
 	if _, err := db.GetSpec(ctx, name); err == nil {
 		return nil, status.Errorf(codes.AlreadyExists, "API spec %q already exists", name)
@@ -97,7 +96,6 @@ func (s *RegistryServer) DeleteApiSpec(ctx context.Context, req *rpc.DeleteApiSp
 	if err != nil {
 		return nil, status.Error(codes.Unavailable, err.Error())
 	}
-	defer db.Close()
 
 	name, err := names.ParseSpec(req.GetName())
 	if err != nil {
@@ -133,7 +131,6 @@ func (s *RegistryServer) getApiSpec(ctx context.Context, name names.Spec) (*rpc.
 	if err != nil {
 		return nil, status.Error(codes.Unavailable, err.Error())
 	}
-	defer db.Close()
 
 	spec, err := db.GetSpec(ctx, name)
 	if err != nil {
@@ -153,7 +150,6 @@ func (s *RegistryServer) getApiSpecRevision(ctx context.Context, name names.Spec
 	if err != nil {
 		return nil, status.Error(codes.Unavailable, err.Error())
 	}
-	defer db.Close()
 
 	revision, err := db.GetSpecRevision(ctx, name)
 	if err != nil {
@@ -174,7 +170,6 @@ func (s *RegistryServer) GetApiSpecContents(ctx context.Context, req *rpc.GetApi
 	if err != nil {
 		return nil, status.Error(codes.Unavailable, err.Error())
 	}
-	defer db.Close()
 
 	var specName = req.GetName()
 	var spec *models.Spec
@@ -219,7 +214,6 @@ func (s *RegistryServer) ListApiSpecs(ctx context.Context, req *rpc.ListApiSpecs
 	if err != nil {
 		return nil, status.Error(codes.Unavailable, err.Error())
 	}
-	defer db.Close()
 
 	if req.GetPageSize() < 0 {
 		return nil, status.Errorf(codes.InvalidArgument, "invalid page_size %d: must not be negative", req.GetPageSize())
@@ -266,7 +260,6 @@ func (s *RegistryServer) UpdateApiSpec(ctx context.Context, req *rpc.UpdateApiSp
 	if err != nil {
 		return nil, status.Error(codes.Unavailable, err.Error())
 	}
-	defer db.Close()
 
 	if req.GetApiSpec() == nil {
 		return nil, status.Errorf(codes.InvalidArgument, "invalid api_spec %+v: body must be provided", req.GetApiSpec())
diff --git a/server/registry/actions_storage.go b/server/registry/actions_storage.go
index ebe6c31..89585eb 100644
--- a/server/registry/actions_storage.go
+++ b/server/registry/actions_storage.go
@@ -29,7 +29,6 @@ func (s *RegistryServer) GetStorage(ctx context.Context, req *emptypb.Empty) (*r
 	if err != nil {
 		return nil, status.Error(codes.Unavailable, err.Error())
 	}
-	defer db.Close()
 	tableNames, err := db.TableNames(ctx)
 	if err != nil {
 		return nil, err
diff --git a/server/registry/actions_versions.go b/server/registry/actions_versions.go
index c27fef7..ef2bb12 100644
--- a/server/registry/actions_versions.go
+++ b/server/registry/actions_versions.go
@@ -45,7 +45,6 @@ func (s *RegistryServer) createApiVersion(ctx context.Context, name names.Versio
 	if err != nil {
 		return nil, status.Error(codes.Unavailable, err.Error())
 	}
-	defer db.Close()
 
 	if err := name.Validate(); err != nil {
 		return nil, status.Error(codes.InvalidArgument, err.Error())
@@ -80,7 +79,6 @@ func (s *RegistryServer) DeleteApiVersion(ctx context.Context, req *rpc.DeleteAp
 	if err != nil {
 		return nil, status.Error(codes.Unavailable, err.Error())
 	}
-	defer db.Close()
 
 	name, err := names.ParseVersion(req.GetName())
 	if err != nil {
@@ -106,7 +104,6 @@ func (s *RegistryServer) GetApiVersion(ctx context.Context, req *rpc.GetApiVersi
 	if err != nil {
 		return nil, status.Error(codes.Unavailable, err.Error())
 	}
-	defer db.Close()
 
 	name, err := names.ParseVersion(req.GetName())
 	if err != nil {
@@ -132,7 +129,6 @@ func (s *RegistryServer) ListApiVersions(ctx context.Context, req *rpc.ListApiVe
 	if err != nil {
 		return nil, status.Error(codes.Unavailable, err.Error())
 	}
-	defer db.Close()
 
 	if req.GetPageSize() < 0 {
 		return nil, status.Errorf(codes.InvalidArgument, "invalid page_size %d: must not be negative", req.GetPageSize())
@@ -198,7 +194,6 @@ func (s *RegistryServer) UpdateApiVersion(ctx context.Context, req *rpc.UpdateAp
 	if err != nil {
 		return nil, status.Error(codes.Unavailable, err.Error())
 	}
-	defer db.Close()
 
 	// TODO: the Get/Update/Save calls below should be a transaction.
 	version, err := db.GetVersion(ctx, name)
diff --git a/server/registry/internal/storage/client.go b/server/registry/internal/storage/client.go
index cb08ba8..d00e7f8 100644
--- a/server/registry/internal/storage/client.go
+++ b/server/registry/internal/storage/client.go
@@ -55,7 +55,8 @@ func NewClient(ctx context.Context, driver, dsn string) (*Client, error) {
 	switch driver {
 	case "sqlite3":
 		db, err := gorm.Open(sqlite.Open(dsn), &gorm.Config{
-			Logger: NewGormLogger(ctx),
+			Logger:      NewGormLogger(ctx),
+			PrepareStmt: true,
 		})
 		if err != nil {
 			c := &Client{db: db}
diff --git a/server/registry/server.go b/server/registry/server.go
index 99e17ae..de1e679 100644
--- a/server/registry/server.go
+++ b/server/registry/server.go
@@ -43,6 +43,7 @@ type RegistryServer struct {
 	dbConfig      string
 	notifyEnabled bool
 	projectID     string
+	client        *storage.Client
 
 	rpc.UnimplementedRegistryServer
 	rpc.UnimplementedAdminServer
@@ -61,19 +62,23 @@ func New(config Config) (*RegistryServer, error) {
 		s.dbConfig = "/tmp/registry.db"
 	}
 
-	db, err := storage.NewClient(context.Background(), s.database, s.dbConfig)
+	var err error
+	s.client, err = storage.NewClient(context.Background(), s.database, s.dbConfig)
 	if err != nil {
 		return nil, err
 	}
-	defer db.Close()
-	if err := db.EnsureTables(); err != nil {
+	if err := s.client.EnsureTables(); err != nil {
 		return nil, err
 	}
 	return s, nil
 }
 
 func (s *RegistryServer) getStorageClient(ctx context.Context) (*storage.Client, error) {
-	return storage.NewClient(ctx, s.database, s.dbConfig)
+	return s.client, nil
+}
+
+func (s *RegistryServer) Close() {
+	s.client.Close()
 }
 
 func isNotFound(err error) bool {


Test patch:
diff --git a/connection/grpctest/server.go b/connection/grpctest/server.go
index 54501af..5326d10 100644
--- a/connection/grpctest/server.go
+++ b/connection/grpctest/server.go
@@ -83,7 +83,7 @@ func TestMain(m *testing.M, rc registry.Config) {
 //   ... run test here ...
 // }
 func NewServer(rc registry.Config) (*Server, error) {
-	gl := &Server{}
+	s := &Server{}
 	var err error
 	if rc.Database == "" {
 		rc.Database = "sqlite3"
@@ -94,15 +94,15 @@ func NewServer(rc registry.Config) (*Server, error) {
 			return nil, err
 		}
 		rc.DBConfig = f.Name()
-		gl.TmpDir = f.Name()
+		s.TmpDir = f.Name()
 	}
 
-	rs, err := registry.New(rc)
+	s.Registry, err = registry.New(rc)
 	if err != nil {
 		return nil, err
 	}
 
-	gl.Listener, gl.Server, err = rs.ServeGRPC(&net.TCPAddr{
+	s.Listener, s.Server, err = s.Registry.ServeGRPC(&net.TCPAddr{
 		IP:   net.IPv4(127, 0, 0, 1),
 		Port: 0}) // random port
 	if err != nil {
@@ -110,44 +110,48 @@ func NewServer(rc registry.Config) (*Server, error) {
 	}
 
 	// set for internal client
-	addr := fmt.Sprintf("localhost:%d", gl.Port())
+	addr := fmt.Sprintf("localhost:%d", s.Port())
 	os.Setenv("APG_REGISTRY_ADDRESS", addr)
 	os.Setenv("APG_REGISTRY_AUDIENCES", fmt.Sprintf("http://%s", addr))
 	os.Setenv("APG_REGISTRY_INSECURE", "1")
 
-	return gl, nil
+	return s, nil
 }
 
 // Server wraps a gRPC Server for a RegistryServer
 type Server struct {
 	Server   *grpc.Server
 	Listener net.Listener
+	Registry *registry.RegistryServer
 	TmpDir   string
 }
 
 // Close will gracefully stop the server and remove tmpDir
-func (g *Server) Close() {
-	if g.Server != nil {
-		g.Server.GracefulStop()
+func (s *Server) Close() {
+	if s.Server != nil {
+		s.Server.GracefulStop()
 	}
-	if g.TmpDir != "" {
-		os.RemoveAll(g.TmpDir)
-		g.TmpDir = ""
+	if s.Registry != nil {
+		s.Registry.Close()
+	}
+	if s.TmpDir != "" {
+		os.RemoveAll(s.TmpDir)
+		s.TmpDir = ""
 	}
 }
 
 // Address returns the address of the listener
-func (g *Server) Address() string {
-	if g.Port() == 0 {
+func (s *Server) Address() string {
+	if s.Port() == 0 {
 		return ""
 	}
-	return fmt.Sprintf("localhost:%d", g.Port())
+	return fmt.Sprintf("localhost:%d", s.Port())
 }
 
 // Port returns the port of the listener
-func (g *Server) Port() int {
-	if g.Listener == nil {
+func (s *Server) Port() int {
+	if s.Listener == nil {
 		return 0
 	}
-	return g.Listener.Addr().(*net.TCPAddr).Port
+	return s.Listener.Addr().(*net.TCPAddr).Port
 }
diff --git a/server/registry/server_test.go b/server/registry/server_test.go
index d2cc5dd..d70fc16 100644
--- a/server/registry/server_test.go
+++ b/server/registry/server_test.go
@@ -39,38 +39,41 @@ func init() {
 	flag.BoolVar(&usePostgres, "postgresql", false, "perform server tests using postgresql")
 }
 
+// defaultTestServer will call server.Close() when test completes
 func defaultTestServer(t *testing.T) *RegistryServer {
 	t.Helper()
-
-	if !usePostgres {
-		if server, err := serverWithSQLite(t); err != nil {
-			t.Fatalf("Setup: failed to get server with SQLite: %s", err)
-		} else {
-			return server
+	var err error
+	var server *RegistryServer
+
+	if usePostgres {
+		server, err = serverWithPostgres(t)
+		if err != nil {
+			t.Errorf("Setup: failed to get server with postgres: %s", err)
+			t.Log("Falling back to server with SQLite storage")
 		}
 	}
-
-	server, err := serverWithPostgres(t)
-	if err != nil {
-		t.Errorf("Setup: failed to get server with postgres: %s", err)
-		t.Log("Falling back to server with SQLite storage")
-		if server, err := serverWithSQLite(t); err != nil {
+	if server == nil {
+		if server, err = serverWithSQLite(t); err != nil {
 			t.Fatalf("Setup: failed to get server with SQLite: %s", err)
-		} else {
-			return server
 		}
 	}
 
 	return server
 }
 
+// serverWithSQLite will call server.Close() when test completes
 func serverWithSQLite(t *testing.T) (*RegistryServer, error) {
-	return New(Config{
+	server, err := New(Config{
 		Database: "sqlite3",
 		DBConfig: fmt.Sprintf("%s/registry.db", t.TempDir()),
 	})
+	if server != nil {
+		t.Cleanup(server.Close)
+	}
+	return server, err
 }
 
+// serverWithPostgres will call server.Close() when test completes
 func serverWithPostgres(t *testing.T) (*RegistryServer, error) {
 	sharedStorage.Lock()
 	t.Cleanup(sharedStorage.Unlock)
@@ -79,10 +82,14 @@ func serverWithPostgres(t *testing.T) (*RegistryServer, error) {
 		return nil, fmt.Errorf("failed to reset database: %s", err)
 	}
 
-	return New(Config{
+	server, err := New(Config{
 		Database: postgresDriver,
 		DBConfig: postgresDBConfig,
 	})
+	if server != nil {
+		t.Cleanup(server.Close)
+	}
+	return server, err
 }
 
 func resetPostgres() error {
diff --git a/server/registry/storage_test.go b/server/registry/storage_test.go
index 9199b30..1f2e55a 100644
--- a/server/registry/storage_test.go
+++ b/server/registry/storage_test.go
@@ -32,7 +32,6 @@ func TestDuplicateCreation(t *testing.T) {
 	if err != nil {
 		t.Fatalf("Setup/Seeding: Failed to seed registry: %s", err)
 	}
-	defer db.Close()
 
 	// Create a project.
 	project := models.NewProject(names.Project{ProjectID: "duplicate"}, &rpc.Project{})