feat(tfstated): implement states locking

This commit is contained in:
Julien Dessaux 2024-10-03 00:13:09 +02:00
parent 8949cee8b3
commit 3d74311931
Signed by: adyxax
GPG key ID: F92E51B86E07177E
7 changed files with 174 additions and 13 deletions

49
cmd/tfstated/lock.go Normal file
View file

@ -0,0 +1,49 @@
package main
import (
"database/sql"
"errors"
"fmt"
"net/http"
"time"
"git.adyxax.org/adyxax/tfstated/pkg/database"
)
type lockRequest struct {
Created time.Time `json:"Created"`
ID string `json:"ID"`
Info string `json:"Info"`
Operation string `json:"Operation"`
Path string `json:"Path"`
Version string `json:"Version"`
Who string `json:"Who"`
}
func handleLock(db *database.DB) http.Handler {
return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
if r.URL.Path == "/" {
_ = encode(w, http.StatusBadRequest,
fmt.Errorf("no state path provided, cannot LOCK /"))
return
}
var lock lockRequest
if err := decode(r, &lock); err != nil {
_ = encode(w, http.StatusBadRequest, err)
return
}
if success, err := db.SetLockOrGetExistingLock(r.URL.Path, &lock); err != nil {
if errors.Is(err, sql.ErrNoRows) {
_ = encode(w, http.StatusNotFound,
fmt.Errorf("state path not found: %s", r.URL.Path))
} else {
_ = errorResponse(w, http.StatusInternalServerError, err)
}
} else if success {
w.WriteHeader(http.StatusOK)
} else {
_ = encode(w, http.StatusConflict, lock)
}
})
}

View file

@ -16,13 +16,20 @@ func handlePost(db *database.DB) http.Handler {
)
return
}
id := r.URL.Query().Get("ID")
data, err := io.ReadAll(r.Body)
if err != nil {
_ = errorResponse(w, http.StatusBadRequest, err)
return
}
if err := db.SetState(r.URL.Path, data); err != nil {
_ = errorResponse(w, http.StatusInternalServerError, err)
if idMismatch, err := db.SetState(r.URL.Path, data, id); err != nil {
if idMismatch {
_ = errorResponse(w, http.StatusConflict, err)
} else {
_ = errorResponse(w, http.StatusInternalServerError, err)
}
} else {
w.WriteHeader(http.StatusOK)
}

View file

@ -14,5 +14,7 @@ func addRoutes(
mux.Handle("DELETE /", handleDelete(db))
mux.Handle("GET /", handleGet(db))
mux.Handle("LOCK /", handleLock(db))
mux.Handle("POST /", handlePost(db))
mux.Handle("UNLOCK /", handleUnlock(db))
}

38
cmd/tfstated/unlock.go Normal file
View file

@ -0,0 +1,38 @@
package main
import (
"database/sql"
"errors"
"fmt"
"net/http"
"git.adyxax.org/adyxax/tfstated/pkg/database"
)
func handleUnlock(db *database.DB) http.Handler {
return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
if r.URL.Path == "/" {
_ = encode(w, http.StatusBadRequest,
fmt.Errorf("no state path provided, cannot LOCK /"))
return
}
var lock lockRequest
if err := decode(r, &lock); err != nil {
_ = encode(w, http.StatusBadRequest, err)
return
}
if success, err := db.Unlock(r.URL.Path, &lock); err != nil {
if errors.Is(err, sql.ErrNoRows) {
_ = encode(w, http.StatusNotFound,
fmt.Errorf("state path not found: %s", r.URL.Path))
} else {
_ = errorResponse(w, http.StatusInternalServerError, err)
}
} else if success {
w.WriteHeader(http.StatusOK)
} else {
_ = encode(w, http.StatusConflict, lock)
}
})
}

50
pkg/database/locks.go Normal file
View file

@ -0,0 +1,50 @@
package database
import (
"encoding/json"
)
// Atomically check if there is an existing lock in place on the state. Returns
// true if it can be set, otherwise returns false and lock is set to the value
// of the existing lock
func (db *DB) SetLockOrGetExistingLock(name string, lock any) (bool, error) {
tx, err := db.Begin()
if err != nil {
return false, err
}
defer func() {
if err != nil {
_ = tx.Rollback()
}
}()
var data []byte
if err = tx.QueryRowContext(db.ctx, `SELECT lock FROM states WHERE name = ?;`, name).Scan(&data); err != nil {
return false, err
}
if data != nil {
err = json.Unmarshal(data, lock)
return false, err
}
if data, err = json.Marshal(lock); err != nil {
return false, err
}
_, err = tx.Exec(`UPDATE states SET lock = json(?) WHERE name = ?;`, data, name)
if err != nil {
return false, err
}
err = tx.Commit()
return true, err
}
func (db *DB) Unlock(name, lock any) (bool, error) {
data, err := json.Marshal(lock)
if err != nil {
return false, err
}
result, err := db.Exec(`UPDATE states SET lock = NULL WHERE name = ? and lock = json(?);`, name, data)
if err != nil {
return false, err
}
n, err := result.RowsAffected()
return n == 1, err
}

View file

@ -5,6 +5,7 @@ CREATE TABLE schema_version (
CREATE TABLE states (
id INTEGER PRIMARY KEY,
name TEXT NOT NULL,
data BLOB NOT NULL
data BLOB NOT NULL,
lock TEXT
) STRICT;
CREATE UNIQUE INDEX states_name on states(name);

View file

@ -2,6 +2,7 @@ package database
import (
"database/sql"
"fmt"
)
func (db *DB) DeleteState(name string) error {
@ -18,18 +19,31 @@ func (db *DB) GetState(name string) ([]byte, error) {
return db.dataEncryptionKey.DecryptAES256(encryptedData)
}
func (db *DB) SetState(name string, data []byte) error {
// returns true in case of id mismatch
func (db *DB) SetState(name string, data []byte, id string) (bool, error) {
encryptedData, err := db.dataEncryptionKey.EncryptAES256(data)
if err != nil {
return err
return false, err
}
_, err = db.Exec(
`INSERT INTO states(name, data) VALUES (:name, :data) ON CONFLICT DO UPDATE SET data = :data WHERE name = :name;`,
sql.Named("data", encryptedData),
sql.Named("name", name),
)
if err != nil {
return err
if id == "" {
_, err = db.Exec(
`INSERT INTO states(name, data) VALUES (:name, :data) ON CONFLICT DO UPDATE SET data = :data WHERE name = :name;`,
sql.Named("data", encryptedData),
sql.Named("name", name),
)
return false, err
} else {
result, err := db.Exec(`UPDATE states SET data = ? WHERE name = ? and lock->>'ID' = ?;`, encryptedData, name, id)
if err != nil {
return false, err
}
n, err := result.RowsAffected()
if err != nil {
return false, err
}
if n != 1 {
return true, fmt.Errorf("failed to update state, lock ID does not match")
}
return false, nil
}
return nil
}