Skip to content

Commit dc8ed5e

Browse files
agneumagneum
committed
feat(engine): add CLI command to generate diff, use sockets to connect to clone databases
1 parent 6ceebec commit dc8ed5e

File tree

7 files changed

+127
-97
lines changed

7 files changed

+127
-97
lines changed

‎engine/cmd/cli/commands/clone/actions.go

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -253,7 +253,7 @@ func schemaDiff(cliCtx *cli.Context) error {
253253
return err
254254
}
255255

256-
_, err = fmt.Fprint(cliCtx.App.Writer, string(response))
256+
_, err = fmt.Fprintln(cliCtx.App.Writer, string(response))
257257

258258
return err
259259
}

‎engine/cmd/database-lab/main.go

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -133,7 +133,7 @@ func main() {
133133
shutdownDatabaseLabEngine(shutdownCtx, docker, engProps, pm.First())
134134
}
135135

136-
cloningSvc := cloning.NewBase(&cfg.Cloning, provisioner, tm, observingChan)
136+
cloningSvc := cloning.NewBase(&cfg.Cloning, &cfg.Global, provisioner, tm, observingChan)
137137
if err = cloningSvc.Run(ctx); err != nil {
138138
log.Err(err)
139139
emergencyShutdown()
@@ -154,7 +154,7 @@ func main() {
154154
})
155155

156156
embeddedUI := embeddedui.New(cfg.EmbeddedUI, engProps, runner, docker)
157-
schemaDiff := schema.NewDiff(docker)
157+
schemaDiff := schema.NewDiff(docker, cloningSvc, pm)
158158
server := srv.NewServer(&cfg.Server, &cfg.Global, engProps, docker, cloningSvc, provisioner, retrievalSvc, platformSvc,
159159
obs, est, schemaDiff, pm, tm)
160160
shutdownCh := setShutdownListener()

‎engine/internal/cloning/base.go

Lines changed: 8 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -24,6 +24,7 @@ import (
2424
"gitlab.com/postgres-ai/database-lab/v3/internal/provision/resources"
2525
"gitlab.com/postgres-ai/database-lab/v3/internal/telemetry"
2626
"gitlab.com/postgres-ai/database-lab/v3/pkg/client/dblabapi/types"
27+
"gitlab.com/postgres-ai/database-lab/v3/pkg/config/global"
2728
"gitlab.com/postgres-ai/database-lab/v3/pkg/log"
2829
"gitlab.com/postgres-ai/database-lab/v3/pkg/models"
2930
"gitlab.com/postgres-ai/database-lab/v3/pkg/util"
@@ -45,6 +46,7 @@ type Config struct {
4546
// Base provides cloning service.
4647
type Base struct {
4748
config *Config
49+
global *global.Config
4850
cloneMutex sync.RWMutex
4951
clones map[string]*CloneWrapper
5052
snapshotBox SnapshotBox
@@ -54,9 +56,10 @@ type Base struct {
5456
}
5557

5658
// NewBase instances a new Base service.
57-
func NewBase(cfg *Config, provision *provision.Provisioner, tm *telemetry.Agent, observingCh chan string) *Base {
59+
func NewBase(cfg *Config, global *global.Config, provision *provision.Provisioner, tm *telemetry.Agent, observingCh chan string) *Base {
5860
return &Base{
5961
config: cfg,
62+
global: global,
6063
clones: make(map[string]*CloneWrapper),
6164
provision: provision,
6265
tm: tm,
@@ -154,6 +157,10 @@ func (c *Base) CreateClone(cloneRequest *types.CloneCreateRequest) (*models.Clon
154157
}
155158
}
156159

160+
if cloneRequest.DB.DBName == "" {
161+
cloneRequest.DB.DBName = c.global.Database.Name()
162+
}
163+
157164
clone := &models.Clone{
158165
ID: cloneRequest.ID,
159166
Snapshot: snapshot,

‎engine/internal/cloning/storage_test.go

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -122,7 +122,7 @@ func TestSavingSessionState(t *testing.T) {
122122
prov, err := newProvisioner()
123123
assert.NoError(t, err)
124124

125-
s := NewBase(nil, prov, &telemetry.Agent{}, nil)
125+
s := NewBase(nil, nil, prov, &telemetry.Agent{}, nil)
126126
err = s.saveClonesState(f.Name())
127127
assert.NoError(t, err)
128128

@@ -166,7 +166,7 @@ func TestFilter(t *testing.T) {
166166
assert.NoError(t, err)
167167
defer func() { _ = os.Remove(filepath) }()
168168

169-
s := NewBase(nil, prov, &telemetry.Agent{}, nil)
169+
s := NewBase(nil, nil, prov, &telemetry.Agent{}, nil)
170170

171171
s.filterRunningClones(context.Background())
172172
assert.Equal(t, 0, len(s.clones))

‎engine/internal/schema/diff/diff.go

Lines changed: 0 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -15,9 +15,6 @@ func OptimizeQueries(queries string) (string, error) {
1515
return "", fmt.Errorf("failed to parse queries %w", err)
1616
}
1717

18-
log.Printf("Original query:\n%v\n\n", queries)
19-
log.Printf("Parse Tree:\n%#v\n\n", idxTree)
20-
2118
stmts := idxTree.GetStmts()
2219
nodes := processStmts(stmts)
2320
idxTree.Stmts = nodes

‎engine/internal/schema/schema.go

Lines changed: 113 additions & 53 deletions
Original file line numberDiff line numberDiff line change
@@ -12,12 +12,16 @@ import (
1212

1313
"github.com/docker/docker/api/types"
1414
"github.com/docker/docker/api/types/container"
15+
"github.com/docker/docker/api/types/mount"
1516
"github.com/docker/docker/api/types/network"
1617
"github.com/docker/docker/client"
1718
"github.com/docker/docker/pkg/stdcopy"
1819

20+
"gitlab.com/postgres-ai/database-lab/v3/internal/cloning"
21+
"gitlab.com/postgres-ai/database-lab/v3/internal/provision/pool"
1922
"gitlab.com/postgres-ai/database-lab/v3/internal/retrieval/engine/postgres/tools"
2023
"gitlab.com/postgres-ai/database-lab/v3/internal/retrieval/engine/postgres/tools/cont"
24+
dle_types "gitlab.com/postgres-ai/database-lab/v3/pkg/client/dblabapi/types"
2125
"gitlab.com/postgres-ai/database-lab/v3/pkg/log"
2226
"gitlab.com/postgres-ai/database-lab/v3/pkg/models"
2327
"gitlab.com/postgres-ai/database-lab/v3/pkg/util"
@@ -29,41 +33,84 @@ const schemaDiffImage = "supabase/pgadmin-schema-diff"
2933
// Diff defines a schema generator.
3034
type Diff struct {
3135
cli *client.Client
36+
cl *cloning.Base
37+
pm *pool.Manager
3238
}
3339

3440
// NewDiff creates a new Diff service.
35-
func NewDiff(cli *client.Client) *Diff {
36-
return &Diff{cli: cli}
37-
}
38-
39-
func connStr(clone *models.Clone) string {
40-
// TODO: fix params
41-
return fmt.Sprintf("postgres://%s:%s@%s:%s/%s",
42-
clone.DB.Username,
43-
"test", // clone.DB.Password,
44-
util.GetCloneNameStr(clone.DB.Port),
45-
clone.DB.Port,
46-
"test", // clone.DB.DBName,
47-
)
41+
func NewDiff(cli *client.Client, cl *cloning.Base, pm *pool.Manager) *Diff {
42+
return &Diff{cli: cli, cl: cl, pm: pm}
4843
}
4944

5045
// GenerateDiff generate difference between database schemas.
51-
func (d *Diff) GenerateDiff(actual, origin *models.Clone, instanceID string) (string, error) {
52-
log.Dbg("Origin clone:", origin)
53-
log.Dbg("Actual clone:", actual.DB.ConnStr+" password=test")
46+
func (d *Diff) GenerateDiff(ctx context.Context, actual *models.Clone, instanceID string) (string, error) {
47+
origin, err := d.createOriginClone(ctx, actual)
48+
if err != nil {
49+
return "", fmt.Errorf("cannot create a clone based on snapshot %s: %w", actual.Snapshot.ID, err)
50+
}
5451

55-
ctx := context.Background()
52+
defer func() {
53+
if err := d.cl.DestroyClone(origin.ID); err != nil {
54+
log.Err("Failed to destroy origin clone:", err)
55+
}
56+
}()
57+
58+
diffContID, err := d.startDiffContainer(ctx, actual, origin, instanceID)
59+
if err != nil {
60+
return "", fmt.Errorf("failed to start diff container: %w", err)
61+
}
5662

57-
if origin.Status.Code != models.StatusOK {
58-
if _, err := d.watchCloneStatus(ctx, origin, origin.Status.Code); err != nil {
59-
return "", fmt.Errorf("failed to watch the clone status: %w", err)
63+
defer func() {
64+
if err := d.cli.ContainerRemove(ctx, diffContID, types.ContainerRemoveOptions{
65+
RemoveVolumes: true,
66+
Force: true,
67+
}); err != nil {
68+
log.Err("failed to remove the diff container:", diffContID, err)
6069
}
70+
}()
71+
72+
statusCh, errCh := d.cli.ContainerWait(ctx, diffContID, container.WaitConditionNotRunning)
73+
select {
74+
case err := <-errCh:
75+
if err != nil {
76+
return "", fmt.Errorf("error on container waiting: %w", err)
77+
}
78+
case <-statusCh:
79+
}
80+
81+
out, err := d.cli.ContainerLogs(ctx, diffContID, types.ContainerLogsOptions{ShowStdout: true})
82+
if err != nil {
83+
return "", fmt.Errorf("failed to get container logs: %w", err)
6184
}
6285

86+
buf := bytes.NewBuffer([]byte{})
87+
88+
if _, err = stdcopy.StdCopy(buf, os.Stderr, out); err != nil {
89+
return "", fmt.Errorf("failed to copy container output: %w", err)
90+
}
91+
92+
filteredOutput, err := filterOutput(buf)
93+
if err != nil {
94+
return "", fmt.Errorf("failed to filter output: %w", err)
95+
}
96+
97+
return filteredOutput.String(), nil
98+
}
99+
100+
// startDiffContainer starts a new diff container.
101+
func (d *Diff) startDiffContainer(ctx context.Context, actual, origin *models.Clone, instanceID string) (string, error) {
63102
if err := tools.PullImage(ctx, d.cli, schemaDiffImage); err != nil {
64103
return "", fmt.Errorf("failed to pull image: %w", err)
65104
}
66105

106+
fsm, err := d.pm.GetFSManager(actual.Snapshot.Pool)
107+
if err != nil {
108+
return "", fmt.Errorf("failed to get pool filesystem manager %s: %w", actual.Snapshot.ID, err)
109+
}
110+
111+
originSocketDir := fsm.Pool().SocketCloneDir(util.GetCloneNameStr(origin.DB.Port))
112+
actualSocketDir := fsm.Pool().SocketCloneDir(util.GetCloneNameStr(actual.DB.Port))
113+
67114
diffCont, err := d.cli.ContainerCreate(ctx,
68115
&container.Config{
69116
Labels: map[string]string{
@@ -72,14 +119,16 @@ func (d *Diff) GenerateDiff(actual, origin *models.Clone, instanceID string) (st
72119
},
73120
Image: schemaDiffImage,
74121
Cmd: []string{
75-
connStr(actual),
76-
connStr(origin),
122+
connString(actual, actualSocketDir),
123+
connString(origin, originSocketDir),
77124
},
78125
},
79-
&container.HostConfig{},
126+
&container.HostConfig{
127+
Mounts: d.getDiffMounts(actualSocketDir, originSocketDir),
128+
},
80129
&network.NetworkingConfig{},
81130
nil,
82-
"clone-diff-"+actual.ID,
131+
d.cloneDiffName(actual),
83132
)
84133
if err != nil {
85134
return "", fmt.Errorf("failed to create diff container: %w", err)
@@ -89,47 +138,57 @@ func (d *Diff) GenerateDiff(actual, origin *models.Clone, instanceID string) (st
89138
return "", fmt.Errorf("failed to connect UI container to the internal Docker network: %w", err)
90139
}
91140

92-
err = d.cli.ContainerStart(ctx, diffCont.ID, types.ContainerStartOptions{})
93-
if err != nil {
141+
if err = d.cli.ContainerStart(ctx, diffCont.ID, types.ContainerStartOptions{}); err != nil {
94142
return "", fmt.Errorf("failed to create diff container: %w", err)
95143
}
96144

97-
defer func() {
98-
if err := d.cli.ContainerRemove(ctx, diffCont.ID, types.ContainerRemoveOptions{
99-
RemoveVolumes: true,
100-
Force: true,
101-
}); err != nil {
102-
log.Err("failed to remove the diff container:", diffCont.ID, err)
103-
}
104-
}()
145+
return diffCont.ID, nil
146+
}
105147

106-
statusCh, errCh := d.cli.ContainerWait(ctx, diffCont.ID, container.WaitConditionNotRunning)
107-
select {
108-
case err := <-errCh:
109-
if err != nil {
110-
return "", fmt.Errorf("error on container waiting: %w", err)
111-
}
112-
case <-statusCh:
148+
func (d *Diff) createOriginClone(ctx context.Context, clone *models.Clone) (*models.Clone, error) {
149+
originClone, err := d.cl.CreateClone(&dle_types.CloneCreateRequest{
150+
ID: d.cloneDiffName(clone),
151+
DB: &dle_types.DatabaseRequest{
152+
Username: clone.DB.Username,
153+
DBName: clone.DB.DBName,
154+
},
155+
Snapshot: &dle_types.SnapshotCloneFieldRequest{ID: clone.Snapshot.ID},
156+
})
157+
if err != nil {
158+
return nil, fmt.Errorf("cannot create a clone based on snapshot %s: %w", clone.Snapshot.ID, err)
113159
}
114160

115-
out, err := d.cli.ContainerLogs(ctx, diffCont.ID, types.ContainerLogsOptions{ShowStdout: true})
116-
if err != nil {
117-
return "", fmt.Errorf("failed to get container logs: %w", err)
161+
if originClone.Status.Code != models.StatusOK {
162+
if _, err := d.watchCloneStatus(ctx, originClone, originClone.Status.Code); err != nil {
163+
return nil, fmt.Errorf("failed to watch the clone status: %w", err)
164+
}
118165
}
119166

120-
buf := bytes.NewBuffer([]byte{})
167+
return originClone, nil
168+
}
121169

122-
_, err = stdcopy.StdCopy(buf, os.Stderr, out)
123-
if err != nil {
124-
return "", fmt.Errorf("failed to copy container output: %w", err)
170+
func (d *Diff) getDiffMounts(actualHost, originHost string) []mount.Mount {
171+
return []mount.Mount{
172+
{
173+
Type: mount.TypeBind,
174+
Source: actualHost,
175+
Target: actualHost,
176+
},
177+
{
178+
Type: mount.TypeBind,
179+
Source: originHost,
180+
Target: originHost,
181+
},
125182
}
183+
}
126184

127-
filteredOutput, err := filterOutput(buf)
128-
if err != nil {
129-
return "", fmt.Errorf("failed to filter output: %w", err)
130-
}
185+
func (d *Diff) cloneDiffName(actual *models.Clone) string {
186+
return "clone-diff-" + actual.ID
187+
}
131188

132-
return filteredOutput.String(), nil
189+
func connString(clone *models.Clone, socketDir string) string {
190+
return fmt.Sprintf("host=%s port=%s user=%s dbname=%s",
191+
socketDir, clone.DB.Port, clone.DB.Username, clone.DB.DBName)
133192
}
134193

135194
// watchCloneStatus checks the clone status for changing.
@@ -176,6 +235,7 @@ func filterOutput(b *bytes.Buffer) (*strings.Builder, error) {
176235
return nil, err
177236
}
178237

238+
// Filter empty lines, comments and warnings.
179239
if len(line) == 0 || bytes.HasPrefix(line, []byte("--")) || bytes.HasPrefix(line, []byte("NOTE:")) {
180240
continue
181241
}

‎engine/internal/srv/routes.go

Lines changed: 1 addition & 35 deletions
Original file line numberDiff line numberDiff line change
@@ -13,7 +13,6 @@ import (
1313
"github.com/gorilla/websocket"
1414
"github.com/jackc/pgtype/pgxtype"
1515
"github.com/pkg/errors"
16-
"github.com/sethvargo/go-password/password"
1716

1817
"gitlab.com/postgres-ai/database-lab/v3/internal/estimator"
1918
"gitlab.com/postgres-ai/database-lab/v3/internal/observer"
@@ -545,40 +544,7 @@ func (s *Server) schemaDiff(w http.ResponseWriter, r *http.Request) {
545544
return
546545
}
547546

548-
const (
549-
PasswordLength = 16
550-
PasswordMinDigits = 4
551-
PasswordMinSymbols = 0
552-
)
553-
554-
_, err = password.Generate(PasswordLength, PasswordMinDigits, PasswordMinSymbols, false, true)
555-
if err != nil {
556-
api.SendError(w, r, fmt.Errorf("failed to generate a password to a reference clone: %w", err))
557-
return
558-
}
559-
560-
originClone, err := s.Cloning.CreateClone(&types.CloneCreateRequest{
561-
ID: "diff-" + cloneID,
562-
DB: &types.DatabaseRequest{
563-
Username: clone.DB.Username,
564-
DBName: clone.DB.DBName,
565-
//Password: pwd,
566-
Password: "test", // TODO: update
567-
},
568-
Snapshot: &types.SnapshotCloneFieldRequest{ID: clone.Snapshot.ID},
569-
})
570-
if err != nil {
571-
api.SendError(w, r, fmt.Errorf("cannot create a clone based on snapshot %s: %w", clone.Snapshot.ID, err))
572-
return
573-
}
574-
575-
defer func() {
576-
if err := s.Cloning.DestroyClone(originClone.ID); err != nil {
577-
log.Err("Failed to destroy origin clone:", err)
578-
}
579-
}()
580-
581-
generateDiff, err := s.SchemaDiff.GenerateDiff(clone, originClone, s.engProps.InstanceID)
547+
generateDiff, err := s.SchemaDiff.GenerateDiff(r.Context(), clone, s.engProps.InstanceID)
582548
if err != nil {
583549
api.SendError(w, r, fmt.Errorf("cannot generate schema diff: %w", err))
584550
return

0 commit comments

Comments
 (0)