github.com/hugorut/terraform@v1.1.3/src/backend/remote-state/pg/client.go (about)

     1  package pg
     2  
     3  import (
     4  	"crypto/md5"
     5  	"database/sql"
     6  	"fmt"
     7  
     8  	uuid "github.com/hashicorp/go-uuid"
     9  	"github.com/hugorut/terraform/src/states/remote"
    10  	"github.com/hugorut/terraform/src/states/statemgr"
    11  	_ "github.com/lib/pq"
    12  )
    13  
    14  // RemoteClient is a remote client that stores data in a Postgres database
    15  type RemoteClient struct {
    16  	Client     *sql.DB
    17  	Name       string
    18  	SchemaName string
    19  
    20  	info *statemgr.LockInfo
    21  }
    22  
    23  func (c *RemoteClient) Get() (*remote.Payload, error) {
    24  	query := `SELECT data FROM %s.%s WHERE name = $1`
    25  	row := c.Client.QueryRow(fmt.Sprintf(query, c.SchemaName, statesTableName), c.Name)
    26  	var data []byte
    27  	err := row.Scan(&data)
    28  	switch {
    29  	case err == sql.ErrNoRows:
    30  		// No existing state returns empty.
    31  		return nil, nil
    32  	case err != nil:
    33  		return nil, err
    34  	default:
    35  		md5 := md5.Sum(data)
    36  		return &remote.Payload{
    37  			Data: data,
    38  			MD5:  md5[:],
    39  		}, nil
    40  	}
    41  }
    42  
    43  func (c *RemoteClient) Put(data []byte) error {
    44  	query := `INSERT INTO %s.%s (name, data) VALUES ($1, $2)
    45  		ON CONFLICT (name) DO UPDATE
    46  		SET data = $2 WHERE %s.name = $1`
    47  	_, err := c.Client.Exec(fmt.Sprintf(query, c.SchemaName, statesTableName, statesTableName), c.Name, data)
    48  	if err != nil {
    49  		return err
    50  	}
    51  	return nil
    52  }
    53  
    54  func (c *RemoteClient) Delete() error {
    55  	query := `DELETE FROM %s.%s WHERE name = $1`
    56  	_, err := c.Client.Exec(fmt.Sprintf(query, c.SchemaName, statesTableName), c.Name)
    57  	if err != nil {
    58  		return err
    59  	}
    60  	return nil
    61  }
    62  
    63  func (c *RemoteClient) Lock(info *statemgr.LockInfo) (string, error) {
    64  	var err error
    65  	var lockID string
    66  
    67  	if info.ID == "" {
    68  		lockID, err = uuid.GenerateUUID()
    69  		if err != nil {
    70  			return "", err
    71  		}
    72  		info.ID = lockID
    73  	}
    74  
    75  	// Local helper function so we can call it multiple places
    76  	//
    77  	lockUnlock := func(pgLockId string) error {
    78  		query := `SELECT pg_advisory_unlock(%s)`
    79  		row := c.Client.QueryRow(fmt.Sprintf(query, pgLockId))
    80  		var didUnlock []byte
    81  		err := row.Scan(&didUnlock)
    82  		if err != nil {
    83  			return &statemgr.LockError{Info: info, Err: err}
    84  		}
    85  		return nil
    86  	}
    87  
    88  	// Try to acquire locks for the existing row `id` and the creation lock `-1`.
    89  	query := `SELECT %s.id, pg_try_advisory_lock(%s.id), pg_try_advisory_lock(-1) FROM %s.%s WHERE %s.name = $1`
    90  	row := c.Client.QueryRow(fmt.Sprintf(query, statesTableName, statesTableName, c.SchemaName, statesTableName, statesTableName), c.Name)
    91  	var pgLockId, didLock, didLockForCreate []byte
    92  	err = row.Scan(&pgLockId, &didLock, &didLockForCreate)
    93  	switch {
    94  	case err == sql.ErrNoRows:
    95  		// No rows means we're creating the workspace. Take the creation lock.
    96  		innerRow := c.Client.QueryRow(`SELECT pg_try_advisory_lock(-1)`)
    97  		var innerDidLock []byte
    98  		err := innerRow.Scan(&innerDidLock)
    99  		if err != nil {
   100  			return "", &statemgr.LockError{Info: info, Err: err}
   101  		}
   102  		if string(innerDidLock) == "false" {
   103  			return "", &statemgr.LockError{Info: info, Err: fmt.Errorf("Already locked for workspace creation: %s", c.Name)}
   104  		}
   105  		info.Path = "-1"
   106  	case err != nil:
   107  		return "", &statemgr.LockError{Info: info, Err: err}
   108  	case string(didLock) == "false":
   109  		// Existing workspace is already locked. Release the attempted creation lock.
   110  		lockUnlock("-1")
   111  		return "", &statemgr.LockError{Info: info, Err: fmt.Errorf("Workspace is already locked: %s", c.Name)}
   112  	case string(didLockForCreate) == "false":
   113  		// Someone has the creation lock already. Release the existing workspace because it might not be safe to touch.
   114  		lockUnlock(string(pgLockId))
   115  		return "", &statemgr.LockError{Info: info, Err: fmt.Errorf("Cannot lock workspace; already locked for workspace creation: %s", c.Name)}
   116  	default:
   117  		// Existing workspace is now locked. Release the attempted creation lock.
   118  		lockUnlock("-1")
   119  		info.Path = string(pgLockId)
   120  	}
   121  	c.info = info
   122  
   123  	return info.ID, nil
   124  }
   125  
   126  func (c *RemoteClient) getLockInfo() (*statemgr.LockInfo, error) {
   127  	return c.info, nil
   128  }
   129  
   130  func (c *RemoteClient) Unlock(id string) error {
   131  	if c.info != nil && c.info.Path != "" {
   132  		query := `SELECT pg_advisory_unlock(%s)`
   133  		row := c.Client.QueryRow(fmt.Sprintf(query, c.info.Path))
   134  		var didUnlock []byte
   135  		err := row.Scan(&didUnlock)
   136  		if err != nil {
   137  			return &statemgr.LockError{Info: c.info, Err: err}
   138  		}
   139  		c.info = nil
   140  	}
   141  	return nil
   142  }