Add a health check function for the key DB store that verifies we can access the required table

Signed-off-by: Ying Li <ying.li@docker.com>

cmd/notary-signer/main.go

@@ -13,11 +13,12 @@ import (
 	"os"
 	"path/filepath"
 	"strings"
+    "time"
 
 	"google.golang.org/grpc"
 	"google.golang.org/grpc/credentials"
 
-	_ "github.com/docker/distribution/health"
+	"github.com/docker/distribution/health"
 	"github.com/docker/notary/cryptoservice"
 	"github.com/docker/notary/signer"
 	"github.com/docker/notary/signer/api"
@@ -147,6 +148,10 @@ func main() {
 	if err != nil {
 		log.Fatalf("failed to create a new keydbstore: %v", err)
 	}
+
+    health.RegisterPeriodicFunc(
+        "DB connectable and valid", keyStore.HealthCheck, time.Second * 60)
+
 	cryptoService := cryptoservice.NewCryptoService("", keyStore)
 
 	cryptoServices[data.ED25519Key] = cryptoService

signer/keydbstore.go

@@ -195,3 +195,17 @@ func (s *KeyDBStore) RotateKeyPassphrase(name, newPassphraseAlias string) error
 
 	return nil
 }
+
+// DB should exist and be query-able
+func (s *KeyDBStore) HealthCheck() error {
+	dbPrivateKey := GormPrivateKey{}
+	tableOk := s.db.HasTable(&dbPrivateKey)
+	switch {
+	case s.db.Error != nil:
+		return s.db.Error
+	case !tableOk:
+		return fmt.Errorf(
+			"Cannot access table: %s", dbPrivateKey.TableName())
+	}
+	return nil
+}

signer/keydbstore_test.go

@@ -27,6 +27,7 @@ var anotherRetriever = func(keyName, alias string, createNew bool, attempts int)
 	return "", false, errors.New("password alias no found")
 }
 
+
 func TestCreateRead(t *testing.T) {
 	tempBaseDir, err := ioutil.TempDir("", "notary-test-")
 	defer os.RemoveAll(tempBaseDir)
@@ -161,3 +162,36 @@ func TestKeyRotation(t *testing.T) {
 	err = dbStore.RotateKeyPassphrase(testKey.ID(), "alias_3")
 	assert.Error(t, err, "password alias no found")
 }
+
+
+func TestDBHealthCheck(t *testing.T) {
+    tempBaseDir, err := ioutil.TempDir("", "notary-test-")
+    defer os.RemoveAll(tempBaseDir)
+
+    // We are using SQLite for the tests
+    db, err := sql.Open("sqlite3", tempBaseDir+"test_db")
+    assert.NoError(t, err)
+
+    // Create a new KeyDB store
+    dbStore, err := NewKeyDBStore(retriever, "", "sqlite3", db)
+    assert.NoError(t, err)
+
+    // No key table, health check fails
+    err = dbStore.HealthCheck()
+    assert.Error(t, err, "Cannot access table:")
+
+    // Ensure that the private_key table exists
+    dbStore.db.CreateTable(&GormPrivateKey{})
+
+    // Heath check success because the table exists
+    err = dbStore.HealthCheck()
+    assert.NoError(t, err)
+
+    // Close the connection
+    err = dbStore.db.Close()
+    assert.NoError(t, err)
+
+    // Heath check fail because the connection is closed
+    err = dbStore.HealthCheck()
+    assert.Error(t, err, "Cannot access table:")
+}