Release database resources on each iteration of a loop (#4305)

hashicorp · Apr 17, 2018 · f23b14a · f23b14a
1 parent 7390abd
commit f23b14a
Show file tree

Hide file tree

Showing 10 changed files with 112 additions and 123 deletions.
diff --git a/builtin/logical/mssql/path_creds_create.go b/builtin/logical/mssql/path_creds_create.go
@@ -6,6 +6,7 @@ import (
 	"strings"
 
 	"github.com/hashicorp/go-uuid"
+	"github.com/hashicorp/vault/helper/dbtxn"
 	"github.com/hashicorp/vault/helper/strutil"
 	"github.com/hashicorp/vault/logical"
 	"github.com/hashicorp/vault/logical/framework"
@@ -90,15 +91,11 @@ func (b *backend) pathCredsCreateRead(ctx context.Context, req *logical.Request,
 			continue
 		}
 
-		stmt, err := tx.Prepare(Query(query, map[string]string{
+		m := map[string]string{
 			"name":     username,
 			"password": password,
-		}))
-		if err != nil {
-			return nil, err
 		}
-		defer stmt.Close()
-		if _, err := stmt.Exec(); err != nil {
+		if err := dbtxn.ExecuteTxQuery(ctx, tx, m, query); err != nil {
 			return nil, err
 		}
 	}

diff --git a/builtin/logical/mssql/secret_creds.go b/builtin/logical/mssql/secret_creds.go
@@ -6,6 +6,7 @@ import (
 	"fmt"
 
 	"github.com/hashicorp/errwrap"
+	"github.com/hashicorp/vault/helper/dbtxn"
 	"github.com/hashicorp/vault/logical"
 	"github.com/hashicorp/vault/logical/framework"
 )
@@ -130,16 +131,11 @@ func (b *backend) secretCredsRevoke(ctx context.Context, req *logical.Request, d
 	// many permissions as possible right now
 	var lastStmtError error
 	for _, query := range revokeStmts {
-		stmt, err := db.Prepare(query)
-		if err != nil {
+
+		if err := dbtxn.ExecuteDBQuery(ctx, db, nil, query); err != nil {
 			lastStmtError = err
 			continue
 		}
-		defer stmt.Close()
-		_, err = stmt.Exec()
-		if err != nil {
-			lastStmtError = err
-		}
 	}
 
 	// can't drop if not all database users are dropped

diff --git a/builtin/logical/mysql/path_role_create.go b/builtin/logical/mysql/path_role_create.go
@@ -6,6 +6,7 @@ import (
 	"strings"
 
 	"github.com/hashicorp/go-uuid"
+	"github.com/hashicorp/vault/helper/dbtxn"
 	"github.com/hashicorp/vault/helper/strutil"
 	"github.com/hashicorp/vault/logical"
 	"github.com/hashicorp/vault/logical/framework"
@@ -103,15 +104,11 @@ func (b *backend) pathRoleCreateRead(ctx context.Context, req *logical.Request,
 			continue
 		}
 
-		stmt, err := tx.Prepare(Query(query, map[string]string{
+		m := map[string]string{
 			"name":     username,
 			"password": password,
-		}))
-		if err != nil {
-			return nil, err
 		}
-		defer stmt.Close()
-		if _, err := stmt.Exec(); err != nil {
+		if err := dbtxn.ExecuteTxQuery(ctx, tx, m, query); err != nil {
 			return nil, err
 		}
 	}

diff --git a/builtin/logical/postgresql/path_role_create.go b/builtin/logical/postgresql/path_role_create.go
@@ -7,6 +7,7 @@ import (
 	"time"
 
 	"github.com/hashicorp/go-uuid"
+	"github.com/hashicorp/vault/helper/dbtxn"
 	"github.com/hashicorp/vault/helper/strutil"
 	"github.com/hashicorp/vault/logical"
 	"github.com/hashicorp/vault/logical/framework"
@@ -106,16 +107,13 @@ func (b *backend) pathRoleCreateRead(ctx context.Context, req *logical.Request,
 			continue
 		}
 
-		stmt, err := tx.Prepare(Query(query, map[string]string{
+		m := map[string]string{
 			"name":       username,
 			"password":   password,
 			"expiration": expiration,
-		}))
-		if err != nil {
-			return nil, err
 		}
-		defer stmt.Close()
-		if _, err := stmt.Exec(); err != nil {
+
+		if err := dbtxn.ExecuteTxQuery(ctx, tx, m, query); err != nil {
 			return nil, err
 		}
 	}

diff --git a/builtin/logical/postgresql/secret_creds.go b/builtin/logical/postgresql/secret_creds.go
@@ -8,6 +8,7 @@ import (
 	"time"
 
 	"github.com/hashicorp/errwrap"
+	"github.com/hashicorp/vault/helper/dbtxn"
 	"github.com/hashicorp/vault/helper/strutil"
 	"github.com/hashicorp/vault/logical"
 	"github.com/hashicorp/vault/logical/framework"
@@ -211,14 +212,7 @@ func (b *backend) secretCredsRevoke(ctx context.Context, req *logical.Request, d
 		// many permissions as possible right now
 		var lastStmtError error
 		for _, query := range revocationStmts {
-			stmt, err := db.Prepare(query)
-			if err != nil {
-				lastStmtError = err
-				continue
-			}
-			defer stmt.Close()
-			_, err = stmt.Exec()
-			if err != nil {
+			if err := dbtxn.ExecuteDBQuery(ctx, db, nil, query); err != nil {
 				lastStmtError = err
 			}
 		}
@@ -258,15 +252,10 @@ func (b *backend) secretCredsRevoke(ctx context.Context, req *logical.Request, d
 				continue
 			}
 
-			stmt, err := tx.Prepare(Query(query, map[string]string{
+			m := map[string]string{
 				"name": username,
-			}))
-			if err != nil {
-				return nil, err
 			}
-			defer stmt.Close()
-
-			if _, err := stmt.Exec(); err != nil {
+			if err := dbtxn.ExecuteTxQuery(ctx, tx, m, query); err != nil {
 				return nil, err
 			}
 		}

diff --git a/helper/dbtxn/dbtxn.go b/helper/dbtxn/dbtxn.go
@@ -0,0 +1,63 @@
+package dbtxn
+
+import (
+	"context"
+	"database/sql"
+	"fmt"
+	"strings"
+)
+
+// ExecuteDBQuery handles executing one single statement, while properly releasing its resources.
+// - ctx: 	Required
+// - db: 	Required
+// - config: 	Optional, may be nil
+// - query: 	Required
+func ExecuteDBQuery(ctx context.Context, db *sql.DB, params map[string]string, query string) error {
+
+	parsedQuery := parseQuery(params, query)
+
+	stmt, err := db.PrepareContext(ctx, parsedQuery)
+	if err != nil {
+		return err
+	}
+	defer stmt.Close()
+
+	return execute(ctx, stmt)
+}
+
+// ExecuteTxQuery handles executing one single statement, while properly releasing its resources.
+// - ctx: 	Required
+// - tx: 	Required
+// - config: 	Optional, may be nil
+// - query: 	Required
+func ExecuteTxQuery(ctx context.Context, tx *sql.Tx, params map[string]string, query string) error {
+
+	parsedQuery := parseQuery(params, query)
+
+	stmt, err := tx.PrepareContext(ctx, parsedQuery)
+	if err != nil {
+		return err
+	}
+	defer stmt.Close()
+
+	return execute(ctx, stmt)
+}
+
+func execute(ctx context.Context, stmt *sql.Stmt) error {
+	if _, err := stmt.ExecContext(ctx); err != nil {
+		return err
+	}
+	return nil
+}
+
+func parseQuery(m map[string]string, tpl string) string {
+
+	if m == nil || len(m) <= 0 {
+		return tpl
+	}
+
+	for k, v := range m {
+		tpl = strings.Replace(tpl, fmt.Sprintf("{{%s}}", k), v, -1)
+	}
+	return tpl
+}
diff --git a/plugins/database/hana/hana.go b/plugins/database/hana/hana.go
@@ -11,6 +11,7 @@ import (
 	_ "github.com/SAP/go-hdb/driver"
 	"github.com/hashicorp/vault/api"
 	"github.com/hashicorp/vault/builtin/logical/database/dbplugin"
+	"github.com/hashicorp/vault/helper/dbtxn"
 	"github.com/hashicorp/vault/helper/strutil"
 	"github.com/hashicorp/vault/plugins"
 	"github.com/hashicorp/vault/plugins/helper/database/connutil"
@@ -143,16 +144,12 @@ func (h *HANA) CreateUser(ctx context.Context, statements dbplugin.Statements, u
 				continue
 			}
 
-			stmt, err := tx.PrepareContext(ctx, dbutil.QueryHelper(query, map[string]string{
+			m := map[string]string{
 				"name":       username,
 				"password":   password,
 				"expiration": expirationStr,
-			}))
-			if err != nil {
-				return "", "", err
 			}
-			defer stmt.Close()
-			if _, err := stmt.ExecContext(ctx); err != nil {
+			if err := dbtxn.ExecuteTxQuery(ctx, tx, m, query); err != nil {
 				return "", "", err
 			}
 		}
@@ -238,14 +235,10 @@ func (h *HANA) RevokeUser(ctx context.Context, statements dbplugin.Statements, u
 				continue
 			}
 
-			stmt, err := tx.PrepareContext(ctx, dbutil.QueryHelper(query, map[string]string{
+			m := map[string]string{
 				"name": username,
-			}))
-			if err != nil {
-				return err
 			}
-			defer stmt.Close()
-			if _, err := stmt.ExecContext(ctx); err != nil {
+			if err := dbtxn.ExecuteTxQuery(ctx, tx, m, query); err != nil {
 				return err
 			}
 		}

diff --git a/plugins/database/mssql/mssql.go b/plugins/database/mssql/mssql.go
@@ -12,6 +12,7 @@ import (
 	"github.com/hashicorp/errwrap"
 	"github.com/hashicorp/vault/api"
 	"github.com/hashicorp/vault/builtin/logical/database/dbplugin"
+	"github.com/hashicorp/vault/helper/dbtxn"
 	"github.com/hashicorp/vault/helper/strutil"
 	"github.com/hashicorp/vault/plugins"
 	"github.com/hashicorp/vault/plugins/helper/database/connutil"
@@ -129,16 +130,13 @@ func (m *MSSQL) CreateUser(ctx context.Context, statements dbplugin.Statements,
 				continue
 			}
 
-			stmt, err := tx.PrepareContext(ctx, dbutil.QueryHelper(query, map[string]string{
+			m := map[string]string{
 				"name":       username,
 				"password":   password,
 				"expiration": expirationStr,
-			}))
-			if err != nil {
-				return "", "", err
 			}
-			defer stmt.Close()
-			if _, err := stmt.ExecContext(ctx); err != nil {
+
+			if err := dbtxn.ExecuteTxQuery(ctx, tx, m, query); err != nil {
 				return "", "", err
 			}
 		}
@@ -189,14 +187,10 @@ func (m *MSSQL) RevokeUser(ctx context.Context, statements dbplugin.Statements,
 				continue
 			}
 
-			stmt, err := tx.PrepareContext(ctx, dbutil.QueryHelper(query, map[string]string{
+			m := map[string]string{
 				"name": username,
-			}))
-			if err != nil {
-				return err
 			}
-			defer stmt.Close()
-			if _, err := stmt.ExecContext(ctx); err != nil {
+			if err := dbtxn.ExecuteTxQuery(ctx, tx, m, query); err != nil {
 				return err
 			}
 		}
@@ -285,14 +279,7 @@ func (m *MSSQL) revokeUserDefault(ctx context.Context, username string) error {
 	// many permissions as possible right now
 	var lastStmtError error
 	for _, query := range revokeStmts {
-		stmt, err := db.PrepareContext(ctx, query)
-		if err != nil {
-			lastStmtError = err
-			continue
-		}
-		defer stmt.Close()
-		_, err = stmt.ExecContext(ctx)
-		if err != nil {
+		if err := dbtxn.ExecuteDBQuery(ctx, db, nil, query); err != nil {
 			lastStmtError = err
 		}
 	}
@@ -355,16 +342,12 @@ func (m *MSSQL) RotateRootCredentials(ctx context.Context, statements []string)
 			if len(query) == 0 {
 				continue
 			}
-			stmt, err := tx.PrepareContext(ctx, dbutil.QueryHelper(query, map[string]string{
+
+			m := map[string]string{
 				"username": m.Username,
 				"password": password,
-			}))
-			if err != nil {
-				return nil, err
 			}
-
-			defer stmt.Close()
-			if _, err := stmt.ExecContext(ctx); err != nil {
+			if err := dbtxn.ExecuteTxQuery(ctx, tx, m, query); err != nil {
 				return nil, err
 			}
 		}

diff --git a/plugins/database/mysql/mysql.go b/plugins/database/mysql/mysql.go
@@ -10,6 +10,7 @@ import (
 	stdmysql "github.com/go-sql-driver/mysql"
 	"github.com/hashicorp/vault/api"
 	"github.com/hashicorp/vault/builtin/logical/database/dbplugin"
+	"github.com/hashicorp/vault/helper/dbtxn"
 	"github.com/hashicorp/vault/helper/strutil"
 	"github.com/hashicorp/vault/plugins"
 	"github.com/hashicorp/vault/plugins/helper/database/connutil"
@@ -182,10 +183,11 @@ func (m *MySQL) CreateUser(ctx context.Context, statements dbplugin.Statements,
 
 				return "", "", err
 			}
-			defer stmt.Close()
 			if _, err := stmt.ExecContext(ctx); err != nil {
+				stmt.Close()
 				return "", "", err
 			}
+			stmt.Close()
 		}
 	}
 
@@ -291,16 +293,12 @@ func (m *MySQL) RotateRootCredentials(ctx context.Context, statements []string)
 			if len(query) == 0 {
 				continue
 			}
-			stmt, err := tx.PrepareContext(ctx, dbutil.QueryHelper(query, map[string]string{
+
+			m := map[string]string{
 				"username": m.Username,
 				"password": password,
-			}))
-			if err != nil {
-				return nil, err
 			}
-
-			defer stmt.Close()
-			if _, err := stmt.ExecContext(ctx); err != nil {
+			if err := dbtxn.ExecuteTxQuery(ctx, tx, m, query); err != nil {
 				return nil, err
 			}
 		}