tfstated/pkg/database/states.go

package database

import (
	"database/sql"
	"errors"
	"fmt"
	"slices"
)

// returns true in case of successful deletion
func (db *DB) DeleteState(name string) (bool, error) {
	result, err := db.Exec(`DELETE FROM states WHERE name = ?;`, name)
	if err != nil {
		return false, err
	}
	n, err := result.RowsAffected()
	if err != nil {
		return false, err
	}
	return n == 1, nil
}

func (db *DB) GetState(name string) ([]byte, error) {
	var encryptedData []byte
	err := db.QueryRow(`SELECT versions.data FROM versions JOIN states ON states.id = versions.state_id  WHERE states.name = ? ORDER BY versions.id DESC LIMIT 1;`, name).Scan(&encryptedData)
	if err != nil {
		if errors.Is(err, sql.ErrNoRows) {
			return []byte{}, nil
		}
		return nil, err
	}
	if encryptedData == nil {
		return []byte{}, nil
	}
	return db.dataEncryptionKey.DecryptAES256(encryptedData)
}

// returns true in case of id mismatch
func (db *DB) SetState(name string, data []byte, lockID string) (bool, error) {
	encryptedData, err := db.dataEncryptionKey.EncryptAES256(data)
	if err != nil {
		return false, err
	}
	tx, err := db.Begin()
	if err != nil {
		return false, err
	}
	defer func() {
		if err != nil {
			_ = tx.Rollback()
		}
	}()
	var (
		stateID  int64
		lockData []byte
	)
	if err = tx.QueryRowContext(db.ctx, `SELECT id, lock->>'ID' FROM states WHERE name = ?;`, name).Scan(&stateID, &lockData); err != nil {
		if errors.Is(err, sql.ErrNoRows) {
			var result sql.Result
			result, err = tx.ExecContext(db.ctx, `INSERT INTO states(name) VALUES (?)`, name)
			if err != nil {
				return false, err
			}
			stateID, err = result.LastInsertId()
			if err != nil {
				return false, err
			}
		} else {
			return false, err
		}
	}

	if lockID != "" && slices.Compare([]byte(lockID), lockData) != 0 {
		err = fmt.Errorf("failed to update state, lock ID does not match")
		return true, err
	}
	_, err = tx.ExecContext(db.ctx, `INSERT INTO versions(state_id, data) VALUES (?, ?);`, stateID, encryptedData)
	if err != nil {
		return false, err
	}
	// TODO delete old states
	err = tx.Commit()
	return false, err
}