Skip to content

Commit d9cf9dd

Browse files
agneumagneum
committed
independent branch datasets (#564)
1 parent 9a5b3f5 commit d9cf9dd

File tree

14 files changed

+419
-43
lines changed

14 files changed

+419
-43
lines changed

‎engine/internal/provision/databases/postgres/postgres.go

Lines changed: 27 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -186,6 +186,33 @@ func getPgConnStr(host, dbname, username string, port uint) string {
186186
return sb.String()
187187
}
188188

189+
// runExistsSQL executes simple SQL commands which returns one bool value.
190+
func runExistsSQL(command, connStr string) (bool, error) {
191+
db, err := sql.Open("postgres", connStr)
192+
193+
if err != nil {
194+
return false, fmt.Errorf("cannot connect to database: %w", err)
195+
}
196+
197+
var result bool
198+
199+
row := db.QueryRow(command)
200+
err = row.Scan(&result)
201+
202+
defer func() {
203+
err := db.Close()
204+
if err != nil {
205+
log.Err("Cannot close database connection.")
206+
}
207+
}()
208+
209+
if err != nil && err == sql.ErrNoRows {
210+
return false, nil
211+
}
212+
213+
return result, err
214+
}
215+
189216
// runSimpleSQL executes simple SQL commands which returns one string value.
190217
func runSimpleSQL(command, connStr string) (string, error) {
191218
db, err := sql.Open("postgres", connStr)

‎engine/internal/provision/databases/postgres/postgres_mgmt.go

Lines changed: 52 additions & 13 deletions
Original file line numberDiff line numberDiff line change
@@ -82,19 +82,37 @@ func CreateUser(c *resources.AppConfig, user resources.EphemeralUser) error {
8282
dbName = user.AvailableDB
8383
}
8484

85+
// check user
86+
pgConnStr := getPgConnStr(c.Host, dbName, c.DB.Username, c.Port)
87+
88+
userExists, err := runExistsSQL(userExistsQuery(user.Name), pgConnStr)
89+
if err != nil {
90+
return fmt.Errorf("failed to check if user exists: %w", err)
91+
}
92+
8593
if user.Restricted {
86-
// create restricted user
87-
query = restrictedUserQuery(user.Name, user.Password)
88-
out, err := runSimpleSQL(query, getPgConnStr(c.Host, dbName, c.DB.Username, c.Port))
94+
// Create or alter restricted user.
95+
query = restrictedUserQuery(user.Name, user.Password, userExists)
96+
out, err := runSimpleSQL(query, pgConnStr)
8997

9098
if err != nil {
9199
return fmt.Errorf("failed to create restricted user: %w", err)
92100
}
93101

94102
log.Dbg("Restricted user has been created: ", out)
95103

96-
// set restricted user as owner for database objects
97-
databaseList, err := runSQLSelectQuery(selectAllDatabases, getPgConnStr(c.Host, dbName, c.DB.Username, c.Port))
104+
// Change user ownership.
105+
query = restrictedUserOwnershipQuery(user.Name, user.Password)
106+
out, err = runSimpleSQL(query, pgConnStr)
107+
108+
if err != nil {
109+
return fmt.Errorf("failed to create restricted user: %w", err)
110+
}
111+
112+
log.Dbg("Database ownership has been changed: ", out)
113+
114+
// Set restricted user as owner for database objects.
115+
databaseList, err := runSQLSelectQuery(selectAllDatabases, pgConnStr)
98116

99117
if err != nil {
100118
return fmt.Errorf("failed list all databases: %w", err)
@@ -111,26 +129,47 @@ func CreateUser(c *resources.AppConfig, user resources.EphemeralUser) error {
111129
log.Dbg("Objects restriction applied", database, out)
112130
}
113131
} else {
114-
query = superuserQuery(user.Name, user.Password)
132+
query = superuserQuery(user.Name, user.Password, userExists)
115133

116-
out, err := runSimpleSQL(query, getPgConnStr(c.Host, dbName, c.DB.Username, c.Port))
134+
out, err := runSimpleSQL(query, pgConnStr)
117135
if err != nil {
118136
return fmt.Errorf("failed to create superuser: %w", err)
119137
}
120138

121-
log.Dbg("Super user has been created: ", out)
139+
log.Dbg("Superuser has been created: ", out)
140+
141+
return nil
122142
}
123143

124144
return nil
125145
}
126146

127-
func superuserQuery(username, password string) string {
128-
return fmt.Sprintf(`create user %s with password %s login superuser;`, pq.QuoteIdentifier(username), pq.QuoteLiteral(password))
147+
func superuserQuery(username, password string, exists bool) string {
148+
if exists {
149+
return fmt.Sprintf(`alter role %s with password %s login superuser;`,
150+
pq.QuoteIdentifier(username), pq.QuoteLiteral(password))
151+
}
152+
153+
return fmt.Sprintf(`create user %s with password %s login superuser;`,
154+
pq.QuoteIdentifier(username), pq.QuoteLiteral(password))
155+
}
156+
157+
func restrictedUserQuery(username, password string, exists bool) string {
158+
if exists {
159+
return fmt.Sprintf(`alter role %s with password %s login;`,
160+
pq.QuoteIdentifier(username), pq.QuoteLiteral(password))
161+
}
162+
163+
return fmt.Sprintf(`create user %s with password %s login;`,
164+
pq.QuoteIdentifier(username), pq.QuoteLiteral(password))
165+
}
166+
167+
func userExistsQuery(username string) string {
168+
return fmt.Sprintf(`select exists (select from pg_roles where rolname = %s)`, pq.QuoteLiteral(username))
129169
}
130170

131171
const restrictionUserCreationTemplate = `
132-
-- create a new user
133-
create user @username with password @password login;
172+
-- change owner
134173
do $$
135174
declare
136175
new_owner text;
@@ -307,7 +346,7 @@ end
307346
$$;
308347
`
309348

310-
func restrictedUserQuery(username, password string) string {
349+
func restrictedUserOwnershipQuery(username, password string) string {
311350
repl := strings.NewReplacer(
312351
"@usernameStr", pq.QuoteLiteral(username),
313352
"@username", pq.QuoteIdentifier(username),

‎engine/internal/provision/databases/postgres/postgres_mgmt_test.go

Lines changed: 54 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -11,45 +11,89 @@ import (
1111
)
1212

1313
func TestSuperuserQuery(t *testing.T) {
14+
const (
15+
user = "user1"
16+
userTest = "user.test\""
17+
pwd = "pwd"
18+
pwdQuote = "pwd\\'--"
19+
)
20+
21+
t.Run("username and password must be quoted", func(t *testing.T) {
22+
assert.Equal(t, `create user "user1" with password 'pwd' login superuser;`, superuserQuery(user, pwd, false))
23+
})
24+
25+
t.Run("username and password must be quoted", func(t *testing.T) {
26+
assert.Equal(t, `alter role "user1" with password 'pwd' login superuser;`, superuserQuery(user, pwd, true))
27+
})
28+
29+
t.Run("special chars must be quoted", func(t *testing.T) {
30+
31+
assert.Equal(t, `create user "user.test""" with password E'pwd\\''--' login superuser;`,
32+
superuserQuery(userTest, pwdQuote, false))
33+
})
34+
35+
t.Run("special chars must be quoted", func(t *testing.T) {
36+
assert.Equal(t, `alter role "user.test""" with password E'pwd\\''--' login superuser;`,
37+
superuserQuery(userTest, pwdQuote, true))
38+
})
39+
}
40+
41+
func TestRestrictedUserQuery(t *testing.T) {
1442
t.Run("username and password must be quoted", func(t *testing.T) {
1543
user := "user1"
1644
pwd := "pwd"
17-
assert.Equal(t, `create user "user1" with password 'pwd' login superuser;`, superuserQuery(user, pwd))
45+
query := restrictedUserQuery(user, pwd, false)
46+
47+
assert.Contains(t, query, `create user "user1" with password 'pwd' login;`)
48+
})
49+
50+
t.Run("username and password must be quoted", func(t *testing.T) {
51+
user := "user1"
52+
pwd := "pwd"
53+
query := restrictedUserQuery(user, pwd, true)
54+
55+
assert.Contains(t, query, `alter role "user1" with password 'pwd' login;`)
56+
})
57+
58+
t.Run("special chars must be quoted", func(t *testing.T) {
59+
user := "user.test\""
60+
pwd := "pwd\\'--"
61+
query := restrictedUserQuery(user, pwd, false)
62+
63+
assert.Contains(t, query, `create user "user.test""" with password E'pwd\\''--' login;`)
1864
})
1965

2066
t.Run("special chars must be quoted", func(t *testing.T) {
2167
user := "user.test\""
2268
pwd := "pwd\\'--"
23-
assert.Equal(t, `create user "user.test""" with password E'pwd\\''--' login superuser;`, superuserQuery(user, pwd))
69+
query := restrictedUserQuery(user, pwd, true)
70+
71+
assert.Contains(t, query, `alter role "user.test""" with password E'pwd\\''--' login;`)
2472
})
2573
}
2674

27-
func TestRestrictedUserQuery(t *testing.T) {
75+
func TestRestrictedUserOwnershipQuery(t *testing.T) {
2876
t.Run("username and password must be quoted", func(t *testing.T) {
2977
user := "user1"
3078
pwd := "pwd"
31-
query := restrictedUserQuery(user, pwd)
79+
query := restrictedUserOwnershipQuery(user, pwd)
3280

33-
assert.Contains(t, query, `create user "user1" with password 'pwd' login;`)
3481
assert.Contains(t, query, `new_owner := 'user1'`)
35-
3682
})
3783

3884
t.Run("special chars must be quoted", func(t *testing.T) {
3985
user := "user.test\""
4086
pwd := "pwd\\'--"
41-
query := restrictedUserQuery(user, pwd)
87+
query := restrictedUserOwnershipQuery(user, pwd)
4288

43-
assert.Contains(t, query, `create user "user.test""" with password E'pwd\\''--' login;`)
4489
assert.Contains(t, query, `new_owner := 'user.test"'`)
4590
})
4691

4792
t.Run("change owner of all databases", func(t *testing.T) {
4893
user := "user.test"
4994
pwd := "pwd"
50-
query := restrictedUserQuery(user, pwd)
95+
query := restrictedUserOwnershipQuery(user, pwd)
5196

5297
assert.Contains(t, query, `select datname from pg_catalog.pg_database where not datistemplat`)
5398
})
54-
5599
}

‎engine/internal/provision/docker/docker.go

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -221,7 +221,7 @@ func RemoveContainer(r runners.Runner, cloneName string) (string, error) {
221221

222222
// ListContainers lists container names.
223223
func ListContainers(r runners.Runner, clonePool string) ([]string, error) {
224-
dockerListCmd := fmt.Sprintf(`docker container ls --filter "label=%s" --filter "label=%s" --all --format '{{.Names}}'`,
224+
dockerListCmd := fmt.Sprintf(`docker container ls --filter "label=%s=%s" --all --format '{{.Names}}'`,
225225
LabelClone, clonePool)
226226

227227
out, err := r.Run(dockerListCmd, false)

‎engine/internal/provision/mode_local_test.go

Lines changed: 8 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -118,6 +118,10 @@ func (m mockFSManager) VerifyBranchMetadata() error {
118118
return nil
119119
}
120120

121+
func (m mockFSManager) CreateDataset(_ string) error {
122+
return nil
123+
}
124+
121125
func (m mockFSManager) CreateBranch(_, _ string) error {
122126
return nil
123127
}
@@ -174,6 +178,10 @@ func (m mockFSManager) SetMountpoint(_, _ string) error {
174178
return nil
175179
}
176180

181+
func (m mockFSManager) Move(_, _, _ string) error {
182+
return nil
183+
}
184+
177185
func (m mockFSManager) Rename(_, _ string) error {
178186
return nil
179187
}

‎engine/internal/provision/pool/manager.go

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -55,13 +55,15 @@ type Snapshotter interface {
5555
type Branching interface {
5656
InitBranching() error
5757
VerifyBranchMetadata() error
58+
CreateDataset(datasetName string) error
5859
CreateBranch(branchName, snapshotID string) error
5960
ListBranches() (map[string]string, error)
6061
ListAllBranches() ([]models.BranchEntity, error)
6162
GetRepo() (*models.Repo, error)
6263
GetAllRepo() (*models.Repo, error)
6364
SetRelation(parent, snapshotName string) error
6465
Snapshot(snapshotName string) error
66+
Move(baseSnap, currentSnap, target string) error
6567
SetMountpoint(path, branch string) error
6668
Rename(oldName, branch string) error
6769
AddBranchProp(branch, snapshotName string) error

‎engine/internal/provision/resources/pool.go

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -95,6 +95,11 @@ func (p *Pool) BranchPath(branchName string) string {
9595
return path.Join(p.BranchDir(), branchName)
9696
}
9797

98+
// BranchName returns a full branch name in the data pool.
99+
func (p *Pool) BranchName(poolName, branchName string) string {
100+
return path.Join(poolName, branchDir, branchName)
101+
}
102+
98103
// Status gets the pool status.
99104
func (p *Pool) Status() PoolStatus {
100105
p.mu.RLock()

‎engine/internal/provision/thinclones/lvm/lvmanager.go

Lines changed: 14 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -156,6 +156,13 @@ func (m *LVManager) VerifyBranchMetadata() error {
156156
return nil
157157
}
158158

159+
// CreateDataset creates a new dataset.
160+
func (m *LVManager) CreateDataset(_ string) error {
161+
log.Msg("CreateDataset is not supported for LVM. Skip the operation")
162+
163+
return nil
164+
}
165+
159166
// CreateBranch clones data as a new branch.
160167
func (m *LVManager) CreateBranch(_, _ string) error {
161168
log.Msg("CreateBranch is not supported for LVM. Skip the operation")
@@ -275,6 +282,13 @@ func (m *LVManager) Rename(_, _ string) error {
275282
return nil
276283
}
277284

285+
// Move moves snapshot diff.
286+
func (m *LVManager) Move(_, _, _ string) error {
287+
log.Msg("Move is not supported for LVM. Skip the operation")
288+
289+
return nil
290+
}
291+
278292
// HasDependentEntity checks if snapshot has dependent entities.
279293
func (m *LVManager) HasDependentEntity(_ string) error {
280294
log.Msg("HasDependentEntity is not supported for LVM. Skip the operation")

‎engine/internal/provision/thinclones/zfs/branching.go

Lines changed: 23 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -96,6 +96,13 @@ func (m *Manager) InitBranching() error {
9696
leader = follower
9797
}
9898

99+
// If not exists pool/branch/main, init main branch dataset.
100+
brName := m.Pool().BranchName(m.Pool().Name, branching.DefaultBranch)
101+
102+
if err := m.CreateDataset(brName); err != nil {
103+
return fmt.Errorf("failed to init main branch dataset: %w", err)
104+
}
105+
99106
log.Msg("data branching has been successfully initialized")
100107

101108
return nil
@@ -151,11 +158,9 @@ func (m *Manager) VerifyBranchMetadata() error {
151158

152159
// CreateBranch clones data as a new branch.
153160
func (m *Manager) CreateBranch(branchName, snapshotID string) error {
154-
branchPath := m.config.Pool.BranchPath(branchName)
155-
156161
// zfs clone -p pool@snapshot_20221019094237 pool/branch/001-branch
157162
cmd := []string{
158-
"zfs clone -p", snapshotID, branchPath,
163+
"zfs clone -p", snapshotID, branchName,
159164
}
160165

161166
out, err := m.runner.Run(strings.Join(cmd, " "))
@@ -169,7 +174,7 @@ func (m *Manager) CreateBranch(branchName, snapshotID string) error {
169174
// Snapshot takes a snapshot of the current data state.
170175
func (m *Manager) Snapshot(snapshotName string) error {
171176
cmd := []string{
172-
"zfs snapshot -r", snapshotName,
177+
"zfs snapshot ", snapshotName,
173178
}
174179

175180
out, err := m.runner.Run(strings.Join(cmd, " "))
@@ -180,6 +185,20 @@ func (m *Manager) Snapshot(snapshotName string) error {
180185
return nil
181186
}
182187

188+
// Move sends and receives snapshot diff.
189+
func (m *Manager) Move(baseSnap, currentSnap, target string) error {
190+
cmd := fmt.Sprintf(
191+
"zfs send -I %s %s | zfs receive -F %s", baseSnap, currentSnap, target,
192+
)
193+
194+
out, err := m.runner.Run(cmd)
195+
if err != nil {
196+
return fmt.Errorf("zfs moving snapshot error: %w. Out: %v", err, out)
197+
}
198+
199+
return nil
200+
}
201+
183202
// Rename renames clone.
184203
func (m *Manager) Rename(oldName, newName string) error {
185204
cmd := []string{

0 commit comments

Comments
 (0)