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