Refactor the code a bit

* Storing SQL queries in a struct should ensure that `datastore.ModifySQLStatement` gets called on all of them. * A wrapper func around `db.Exec` reduces copypasta. * Actually call `InitRepositoryProvider` for API keys package
cloudfoundry · richard-cox · Aug 25, 2020 · Aug 18, 2020 · Aug 19, 2020 · Aug 20, 2020
commit 541b8f324ea5951be28167f440863b26fe8a177f
diff --git a/src/jetstream/main.go b/src/jetstream/main.go
@@ -37,6 +37,7 @@ import (
 
 	"github.com/cloudfoundry-incubator/stratos/src/jetstream/crypto"
 	"github.com/cloudfoundry-incubator/stratos/src/jetstream/datastore"
+	"github.com/cloudfoundry-incubator/stratos/src/jetstream/repository/apikeys"
 	"github.com/cloudfoundry-incubator/stratos/src/jetstream/repository/cnsis"
 	"github.com/cloudfoundry-incubator/stratos/src/jetstream/repository/console_config"
 	"github.com/cloudfoundry-incubator/stratos/src/jetstream/repository/interfaces"
@@ -178,6 +179,7 @@ func main() {
 	console_config.InitRepositoryProvider(dc.DatabaseProvider)
 	localusers.InitRepositoryProvider(dc.DatabaseProvider)
 	sessiondata.InitRepositoryProvider(dc.DatabaseProvider)
+	apikeys.InitRepositoryProvider(dc.DatabaseProvider)
 
 	// Establish a Postgresql connection pool
 	databaseConnectionPool, migratorConf, err := initConnPool(dc, envLookup)

diff --git a/src/jetstream/repository/apikeys/psql_apikeys.go b/src/jetstream/repository/apikeys/psql_apikeys.go
@@ -4,6 +4,8 @@ import (
 	"database/sql"
 	"encoding/base64"
 	"errors"
+	"fmt"
+	"reflect"
 	"time"
 
 	"github.com/cloudfoundry-incubator/stratos/src/jetstream/crypto"
@@ -13,13 +15,19 @@ import (
 	log "github.com/sirupsen/logrus"
 )
 
-var insertAPIKey = `INSERT INTO api_keys (guid, secret, user_guid, comment) VALUES ($1, $2, $3, $4)`
-var getAPIKeyBySecret = `SELECT guid, user_guid, comment, last_used FROM api_keys WHERE secret = $1`
-var listAPIKeys = `SELECT guid, user_guid, comment, last_used FROM api_keys WHERE user_guid = $1`
-var deleteAPIKey = `DELETE FROM api_keys WHERE user_guid = $1 AND guid = $2`
-var updateAPIKeyLastUsed = `UPDATE api_keys SET last_used = $1 WHERE guid = $2`
-
-// UpdateAPIKeyLastUsed
+var sqlQueries = struct {
+	InsertAPIKey         string
+	GetAPIKeyBySecret    string
+	ListAPIKeys          string
+	DeleteAPIKey         string
+	UpdateAPIKeyLastUsed string
+}{
+	InsertAPIKey:         `INSERT INTO api_keys (guid, secret, user_guid, comment) VALUES ($1, $2, $3, $4)`,
+	GetAPIKeyBySecret:    `SELECT guid, user_guid, comment, last_used FROM api_keys WHERE secret = $1`,
+	ListAPIKeys:          `SELECT guid, user_guid, comment, last_used FROM api_keys WHERE user_guid = $1`,
+	DeleteAPIKey:         `DELETE FROM api_keys WHERE user_guid = $1 AND guid = $2`,
+	UpdateAPIKeyLastUsed: `UPDATE api_keys SET last_used = $1 WHERE guid = $2`,
+}
 
 // PgsqlAPIKeysRepository - Postgresql-backed API keys repository
 type PgsqlAPIKeysRepository struct {
@@ -35,11 +43,19 @@ func NewPgsqlAPIKeysRepository(dcp *sql.DB) (Repository, error) {
 // InitRepositoryProvider - One time init for the given DB Provider
 func InitRepositoryProvider(databaseProvider string) {
 	// Modify the database statements if needed, for the given database type
-	insertAPIKey = datastore.ModifySQLStatement(insertAPIKey, databaseProvider)
-	getAPIKeyBySecret = datastore.ModifySQLStatement(getAPIKeyBySecret, databaseProvider)
-	deleteAPIKey = datastore.ModifySQLStatement(deleteAPIKey, databaseProvider)
-	listAPIKeys = datastore.ModifySQLStatement(listAPIKeys, databaseProvider)
-	updateAPIKeyLastUsed = datastore.ModifySQLStatement(updateAPIKeyLastUsed, databaseProvider)
+	// Iterating over the struct to ensure that all of the queries are updated
+	v := reflect.ValueOf(sqlQueries)
+	for i := 0; i < v.NumField(); i++ {
+		q := v.Field(i).Interface().(string)
+
+		reflect.
+			ValueOf(&sqlQueries).
+			Elem().
+			FieldByIndex([]int{i}).
+			SetString(
+				datastore.ModifySQLStatement(q, databaseProvider),
+			)
+	}
 }
 
 // AddAPIKey - Add a new API key to the datastore.
@@ -67,18 +83,9 @@ func (p *PgsqlAPIKeysRepository) AddAPIKey(userID string, comment string) (*inte
 	keyGUID := uuid.NewV4().String()
 	keySecret := base64.URLEncoding.EncodeToString(randomBytes)
 
-	var result sql.Result
-	if result, err = p.db.Exec(insertAPIKey, keyGUID, keySecret, userID, comment); err != nil {
-		log.Errorf("unable to INSERT API key: %v", err)
-		return nil, err
-	}
-
-	//Validate that 1 row has been updated
-	rowsUpdates, err := result.RowsAffected()
+	err = execQuery(p, sqlQueries.InsertAPIKey, keyGUID, keySecret, userID, comment)
 	if err != nil {
-		return nil, errors.New("unable to INSERT api key: could not determine number of rows that were updated")
-	} else if rowsUpdates < 1 {
-		return nil, errors.New("unable to INSERT api key: no rows were updated")
+		return nil, fmt.Errorf("AddAPIKey: %v", err)
 	}
 
 	apiKey := &interfaces.APIKey{
@@ -97,7 +104,7 @@ func (p *PgsqlAPIKeysRepository) GetAPIKeyBySecret(keySecret string) (*interface
 
 	var apiKey interfaces.APIKey
 
-	err := p.db.QueryRow(getAPIKeyBySecret, keySecret).Scan(
+	err := p.db.QueryRow(sqlQueries.GetAPIKeyBySecret, keySecret).Scan(
 		&apiKey.GUID,
 		&apiKey.UserGUID,
 		&apiKey.Comment,
@@ -115,7 +122,7 @@ func (p *PgsqlAPIKeysRepository) GetAPIKeyBySecret(keySecret string) (*interface
 func (p *PgsqlAPIKeysRepository) ListAPIKeys(userID string) ([]interfaces.APIKey, error) {
 	log.Debug("ListAPIKeys")
 
-	rows, err := p.db.Query(listAPIKeys, userID)
+	rows, err := p.db.Query(sqlQueries.ListAPIKeys, userID)
 	if err != nil {
 		log.Errorf("unable to list API keys: %v", err)
 		return nil, err
@@ -139,16 +146,9 @@ func (p *PgsqlAPIKeysRepository) ListAPIKeys(userID string) ([]interfaces.APIKey
 func (p *PgsqlAPIKeysRepository) DeleteAPIKey(userGUID string, keyGUID string) error {
 	log.Debug("DeleteAPIKey")
 
-	result, err := p.db.Exec(deleteAPIKey, userGUID, keyGUID)
-	if err != nil {
-		return err
-	}
-
-	rowsUpdates, err := result.RowsAffected()
+	err := execQuery(p, sqlQueries.DeleteAPIKey, userGUID, keyGUID)
 	if err != nil {
-		return errors.New("unable to DELETE api key: could not determine number of rows that were updated")
-	} else if rowsUpdates < 1 {
-		return errors.New("unable to DELETE api key: no rows were updated")
+		return fmt.Errorf("DeleteAPIKey: %v", err)
 	}
 
 	return nil
@@ -158,16 +158,26 @@ func (p *PgsqlAPIKeysRepository) DeleteAPIKey(userGUID string, keyGUID string) e
 func (p *PgsqlAPIKeysRepository) UpdateAPIKeyLastUsed(keyGUID string) error {
 	log.Debug("UpdateAPIKeyLastUsed")
 
-	result, err := p.db.Exec(updateAPIKeyLastUsed, time.Now(), keyGUID)
+	err := execQuery(p, sqlQueries.UpdateAPIKeyLastUsed, time.Now(), keyGUID)
+	if err != nil {
+		return fmt.Errorf("UpdateAPIKeyLastUsed: %v", err)
+	}
+
+	return nil
+}
+
+// A wrapper around db.Exec that validates that exactly 1 row has been inserted/deleted/updated
+func execQuery(p *PgsqlAPIKeysRepository, query string, args ...interface{}) error {
+	result, err := p.db.Exec(query, args...)
 	if err != nil {
 		return err
 	}
 
 	rowsUpdates, err := result.RowsAffected()
 	if err != nil {
-		return errors.New("unable to UPDATE api key: could not determine number of rows that were updated")
+		return errors.New("could not determine number of rows that were updated")
 	} else if rowsUpdates < 1 {
-		return errors.New("unable to UPDATE api key: no rows were updated")
+		return errors.New("no rows were updated")
 	}
 
 	return nil

diff --git a/src/jetstream/repository/apikeys/psql_apikeys_test.go b/src/jetstream/repository/apikeys/psql_apikeys_test.go
@@ -63,7 +63,7 @@ func TestAddAPIKey(t *testing.T) {
 			apiKey, err := repository.AddAPIKey(userID, comment)
 
 			Convey("an error should be returned", func() {
-				So(err, ShouldResemble, errors.New("unable to INSERT api key: no rows were updated"))
+				So(err, ShouldResemble, errors.New("AddAPIKey: no rows were updated"))
 			})
 
 			Convey("should return nil", func() {
@@ -270,7 +270,7 @@ func TestDeleteAPIKey(t *testing.T) {
 			err := repository.DeleteAPIKey(userID, keyID)
 
 			Convey("an error should be returned", func() {
-				So(err, ShouldResemble, errors.New("unable to DELETE api key: no rows were updated"))
+				So(err, ShouldResemble, errors.New("DeleteAPIKey: no rows were updated"))
 			})
 		})
 
@@ -316,7 +316,7 @@ func TestUpdateAPIKeyLastUsed(t *testing.T) {
 			err := repository.UpdateAPIKeyLastUsed(keyID)
 
 			Convey("an error should be returned", func() {
-				So(err, ShouldResemble, errors.New("unable to UPDATE api key: no rows were updated"))
+				So(err, ShouldResemble, errors.New("UpdateAPIKeyLastUsed: no rows were updated"))
 			})
 		})