Move the upsert definition back into the code

This commit is contained in:
Devin Christensen 2016-01-22 09:47:02 -07:00
parent bfbdc72e03
commit 32b712ddb1
2 changed files with 37 additions and 46 deletions

View file

@ -8,7 +8,7 @@ import (
"time"
"github.com/armon/go-metrics"
_ "github.com/lib/pq"
"github.com/lib/pq"
)
// PostgreSQL Backend is a physical backend that stores data
@ -28,15 +28,11 @@ func newPostgreSQLBackend(conf map[string]string) (Backend, error) {
return nil, fmt.Errorf("missing connection_url")
}
table, ok := conf["table"]
unquoted_table, ok := conf["table"]
if !ok {
table = "vault"
}
upsert_function, ok := conf["upsert_function"]
if !ok {
upsert_function = "vault_upsert"
unquoted_table = "vault"
}
quoted_table := pq.QuoteIdentifier(unquoted_table)
// Create PostgreSQL handle for the database.
db, err := sql.Open("postgres", connURL)
@ -51,18 +47,45 @@ func newPostgreSQLBackend(conf map[string]string) (Backend, error) {
return nil, fmt.Errorf("failed to check for native upsert: %v", err)
}
// Setup our put strategy based on the presence or absence of a native
// upsert. The upsert function used is taken [from the PostgreSQL
// docs](http://www.postgresql.org/docs/current/static/plpgsql-control-structures.html#PLPGSQL-UPSERT-EXAMPLE)
// and chosen primarily for reasons [listed
// here](http://www.depesz.com/2012/06/10/why-is-upsert-so-complicated/)
var put_statement string
create_upsert_sql := `
CREATE OR REPLACE FUNCTION vault_upsert(_key TEXT, _value BYTEA) RETURNS VOID AS
$$
BEGIN
LOOP
UPDATE ` + quoted_table + ` SET vault_value = _value WHERE vault_key = _key;
IF found THEN
RETURN;
END IF;
BEGIN
INSERT INTO ` + quoted_table + ` (vault_key, vault_value) VALUES (_key, _value);
RETURN;
EXCEPTION WHEN unique_violation THEN
-- Do nothing, and loop to try the UPDATE again.
END;
END LOOP;
END;
$$
LANGUAGE plpgsql;`
if upsert_required {
put_statement = "SELECT " + upsert_function + "($1, $2)"
put_statement = "SELECT vault_upsert($1, $2)"
if _, err := db.Exec(create_upsert_sql); err != nil {
return nil, fmt.Errorf("failed to create upsert function: %v", err)
}
} else {
put_statement = "INSERT INTO " + table + " VALUES($1, $2)" +
put_statement = "INSERT INTO " + quoted_table + " VALUES($1, $2)" +
" ON CONFLICT (vault_key) DO " +
" UPDATE SET vault_value = $2"
}
// Setup the backend.
m := &PostgreSQLBackend{
table: table,
table: unquoted_table,
client: db,
statements: make(map[string]*sql.Stmt),
}
@ -70,9 +93,9 @@ func newPostgreSQLBackend(conf map[string]string) (Backend, error) {
// Prepare all the statements required
statements := map[string]string{
"put": put_statement,
"get": "SELECT vault_value FROM " + table + " WHERE vault_key = $1",
"delete": "DELETE FROM " + table + " WHERE vault_key = $1",
"list": "SELECT vault_key FROM " + table + " WHERE vault_key LIKE $1",
"get": "SELECT vault_value FROM " + quoted_table + " WHERE vault_key = $1",
"delete": "DELETE FROM " + quoted_table + " WHERE vault_key = $1",
"list": "SELECT vault_key FROM " + quoted_table + " WHERE vault_key LIKE $1",
}
for name, query := range statements {
if err := m.prepare(name, query); err != nil {

View file

@ -317,9 +317,6 @@ The PostgreSQL backend has the following options:
* `table` (optional) - The name of the table to write vault data to. Defaults
to "vault".
* `upsert_function` (optional) - The name of the upsert function. Defaults to
"vault_upsert". *This will only be used if you're running a version of PostgreSQL prior to 9.5*
Make sure the PostgreSQL database you choose (or create) for vault storage has
a table suitable for storing vault's data:
@ -330,35 +327,6 @@ CREATE TABLE vault (
);
```
If you're using a version of PostgreSQL prior to 9.5, create an upsert function
in the database you will be using for vault storage (taken from
[PostgreSQL documentation](http://www.postgresql.org/docs/9.4/static/plpgsql-control-structures.html#PLPGSQL-UPSERT-EXAMPLE):
```sql
CREATE FUNCTION vault_upsert(_key TEXT, _value BYTEA) RETURNS VOID AS
$$
BEGIN
LOOP
-- first try to update the key
UPDATE vault SET vault_value = _value WHERE vault_key = _key;
IF found THEN
RETURN;
END IF;
-- not there, so try to insert the key
-- if someone else inserts the same key concurrently,
-- we could get a unique-key failure
BEGIN
INSERT INTO vault (vault_key, vault_value) VALUES (_key, _value);
RETURN;
EXCEPTION WHEN unique_violation THEN
-- Do nothing, and loop to try the UPDATE again.
END;
END LOOP;
END;
$$
LANGUAGE plpgsql;
```
#### Backend Reference: Inmem
The in-memory backend has no configuration options.