Thanks to visit codestin.com
Credit goes to github.com

Skip to content

Commit 06034af

Browse files
spikecurtisspikecurtis
committed
chore: integrate broker into dbtestutil
1 parent ef3db0e commit 06034af

File tree

3 files changed

+106
-58
lines changed

3 files changed

+106
-58
lines changed

‎cli/templatepull_test.go

Lines changed: 6 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -263,8 +263,13 @@ func TestTemplatePull_ToDir(t *testing.T) {
263263
// nolint: paralleltest // These tests change the current working dir, and is therefore unsuitable for parallelisation.
264264
for _, tc := range tests {
265265
t.Run(tc.name, func(t *testing.T) {
266-
dir := t.TempDir()
266+
// create coderd first, because our postgres cloning code needs to be run from somewhere in the package
267+
// hierarchy, before we change directories.
268+
client := coderdtest.New(t, &coderdtest.Options{
269+
IncludeProvisionerDaemon: true,
270+
})
267271

272+
dir := t.TempDir()
268273
cwd, err := os.Getwd()
269274
require.NoError(t, err)
270275
t.Cleanup(func() {
@@ -282,9 +287,6 @@ func TestTemplatePull_ToDir(t *testing.T) {
282287
actualDest = filepath.Join(dir, "actual")
283288
}
284289

285-
client := coderdtest.New(t, &coderdtest.Options{
286-
IncludeProvisionerDaemon: true,
287-
})
288290
owner := coderdtest.CreateFirstUser(t, client)
289291
templateAdmin, _ := coderdtest.CreateAnotherUser(t, client, owner.OrganizationID, rbac.RoleTemplateAdmin())
290292

‎coderd/database/dbtestutil/postgres.go

Lines changed: 75 additions & 54 deletions
Original file line numberDiff line numberDiff line change
@@ -21,8 +21,8 @@ import (
2121
"github.com/ory/dockertest/v3/docker"
2222
"golang.org/x/xerrors"
2323

24+
"github.com/coder/coder/v2/coderd/database/dbtestutil/broker"
2425
"github.com/coder/coder/v2/coderd/database/migrations"
25-
"github.com/coder/coder/v2/cryptorand"
2626
"github.com/coder/retry"
2727
)
2828

@@ -52,6 +52,7 @@ var (
5252
"connection refused", // nothing is listening on the port
5353
"No connection could be made", // Windows variant of the above
5454
}
55+
defaultBrokerClientSingleton = &broker.Singleton{}
5556
)
5657

5758
// initDefaultConnection initializes the default postgres connection parameters.
@@ -173,6 +174,7 @@ type TBSubset interface {
173174
// Otherwise, it will start a new postgres container.
174175
func Open(t TBSubset, opts ...OpenOption) (string, error) {
175176
t.Helper()
177+
t.Logf("%s dbtestutil: Open new database", time.Now().Format(time.StampMilli))
176178

177179
connectionParamsInitOnce.Do(func() {
178180
errDefaultConnectionParamsInit = initDefaultConnection(t)
@@ -193,40 +195,29 @@ func Open(t TBSubset, opts ...OpenOption) (string, error) {
193195
port = defaultConnectionParams.Port
194196
)
195197

196-
// Use a time-based prefix to make it easier to find the database
197-
// when debugging.
198-
now := time.Now().Format("test_2006_01_02_15_04_05")
199-
dbSuffix, err := cryptorand.StringCharset(cryptorand.Lower, 10)
200-
if err != nil {
201-
return "", xerrors.Errorf("generate db suffix: %w", err)
202-
}
203-
dbName := now + "_" + dbSuffix
204-
205198
// if empty createDatabaseFromTemplate will create a new template db
206199
templateDBName := os.Getenv("DB_FROM")
207200
if openOptions.DBFrom != nil {
208201
templateDBName = *openOptions.DBFrom
209202
}
210-
if err = createDatabaseFromTemplate(t, defaultConnectionParams, dbName, templateDBName); err != nil {
203+
dbName, err := createDatabaseFromTemplate(t, defaultConnectionParams, templateDBName)
204+
if err != nil {
211205
return "", xerrors.Errorf("create database: %w", err)
212206
}
213207

208+
// initialize outside the function so that the singleton keeps the client around
209+
brokerClient, err := defaultBrokerClientSingleton.Get(t)
210+
if err != nil {
211+
return "", xerrors.Errorf("get broker client: %w", err)
212+
}
214213
t.Cleanup(func() {
215-
cleanupDbURL := defaultConnectionParams.DSN()
216-
cleanupConn, err := sql.Open("postgres", cleanupDbURL)
217-
if err != nil {
218-
t.Logf("cleanup database %q: failed to connect to postgres: %s\n", dbName, err.Error())
219-
return
220-
}
221-
defer func() {
222-
if err := cleanupConn.Close(); err != nil {
223-
t.Logf("cleanup database %q: failed to close connection: %s\n", dbName, err.Error())
224-
}
225-
}()
226-
_, err = cleanupConn.Exec("DROP DATABASE " + dbName + ";")
214+
resp, err := brokerClient.Discard(context.Background(), &broker.DiscardRequest{
215+
DbName: dbName,
216+
})
227217
if err != nil {
228-
t.Logf("failed to clean up database %q: %s\n", dbName, err.Error())
229-
return
218+
t.Logf("discard database: %s", err)
219+
} else if resp.Status.Code != broker.Status_OK {
220+
t.Logf("discard database: (%s) %s", resp.Status.Code.String(), resp.Status.Message)
230221
}
231222
})
232223

@@ -249,9 +240,67 @@ func Open(t TBSubset, opts ...OpenOption) (string, error) {
249240
// If templateDBName is empty, it will create a new template database based on
250241
// the current migrations, and name it "tpl_<migrations_hash>". Or if it's
251242
// already been created, it will use that.
252-
func createDatabaseFromTemplate(t TBSubset, connParams ConnectionParams, newDBName string, templateDBName string) error {
243+
func createDatabaseFromTemplate(t TBSubset, connParams ConnectionParams, templateDBName string) (string, error) {
253244
t.Helper()
245+
// If this takes longer than 20s, we are in big trouble since many tests time out after 10s. So bail out here to
246+
// make it more clear what the root cause is.
247+
ctx, cancel := context.WithTimeout(context.Background(), 20*time.Second)
248+
defer cancel()
249+
250+
emptyTemplateDBName := templateDBName == ""
251+
if emptyTemplateDBName {
252+
templateDBName = fmt.Sprintf("tpl_%s", migrations.GetMigrationsHash()[:32])
253+
}
254+
tplDbDoesNotExistOccurred := false
255+
256+
brokerClient, err := defaultBrokerClientSingleton.Get(t)
257+
if err != nil {
258+
return "", xerrors.Errorf("get broker client: %w", err)
259+
}
260+
t.Logf("%s dbtestutil: got brokerClient", time.Now().Format(time.StampMilli))
261+
queryResp, err := brokerClient.Query(ctx, &broker.QueryRequest{
262+
DbName: templateDBName,
263+
})
264+
t.Logf("%s dbtestutil: queried template database", time.Now().Format(time.StampMilli))
265+
if err != nil {
266+
return "", xerrors.Errorf("query template database: %w", err)
267+
}
268+
if queryResp.Status.Code == broker.Status_ERR_DB_NOT_FOUND {
269+
tplDbDoesNotExistOccurred = true
270+
} else if queryResp.Status.Code != broker.Status_OK {
271+
return "", xerrors.Errorf(
272+
"query template database: (%s) %s", queryResp.Status.Code.String(), queryResp.Status.Message)
273+
}
274+
if tplDbDoesNotExistOccurred && !emptyTemplateDBName {
275+
return "", xerrors.Errorf(`template database '%s' does not exist`, templateDBName)
276+
}
277+
if tplDbDoesNotExistOccurred {
278+
// We need to create the templateDatabase because one was not explicitly passed to us
279+
err = createTemplateDatabase(t, connParams, templateDBName)
280+
if err != nil {
281+
return "", err
282+
}
283+
}
284+
285+
t.Logf("%s dbtestutil: cloning database from template", time.Now().Format(time.StampMilli))
286+
cloneResp, err := brokerClient.Clone(ctx, &broker.CloneRequest{
287+
TemplateDbName: templateDBName,
288+
})
289+
if err != nil {
290+
return "", xerrors.Errorf("clone template database: %w", err)
291+
}
292+
if cloneResp.Status.Code != broker.Status_OK {
293+
return "", xerrors.Errorf("clone template database: (%s) %s",
294+
cloneResp.Status.Code.String(), cloneResp.Status.Message)
295+
}
296+
t.Logf("%s dbtestutil: clone complete", time.Now().Format(time.StampMilli))
297+
return cloneResp.DbName, nil
298+
}
254299

300+
func createTemplateDatabase(t TBSubset, connParams ConnectionParams, templateDBName string) error {
301+
// We need to do this directly, and not thru the broker because we want to be 100% sure the migrations we apply
302+
// match *this test binary*, which might get compiled at a different time than the broker.
303+
t.Logf("%s dbtestutil: creating template database", time.Now().Format(time.StampMilli))
255304
dbURL := connParams.DSN()
256305
db, err := sql.Open("postgres", dbURL)
257306
if err != nil {
@@ -263,26 +312,6 @@ func createDatabaseFromTemplate(t TBSubset, connParams ConnectionParams, newDBNa
263312
}
264313
}()
265314

266-
emptyTemplateDBName := templateDBName == ""
267-
if emptyTemplateDBName {
268-
templateDBName = fmt.Sprintf("tpl_%s", migrations.GetMigrationsHash()[:32])
269-
}
270-
_, err = db.Exec("CREATE DATABASE " + newDBName + " WITH TEMPLATE " + templateDBName)
271-
if err == nil {
272-
// Template database already exists and we successfully created the new database.
273-
return nil
274-
}
275-
tplDbDoesNotExistOccurred := strings.Contains(err.Error(), "template database") && strings.Contains(err.Error(), "does not exist")
276-
if (tplDbDoesNotExistOccurred && !emptyTemplateDBName) || !tplDbDoesNotExistOccurred {
277-
// First and case: user passed a templateDBName that doesn't exist.
278-
// Second and case: some other error.
279-
return xerrors.Errorf("create db with template: %w", err)
280-
}
281-
if !emptyTemplateDBName {
282-
// sanity check
283-
panic("templateDBName is not empty. there's a bug in the code above")
284-
}
285-
// The templateDBName is empty, so we need to create the template database.
286315
// We will use a tx to obtain a lock, so another test or process doesn't race with us.
287316
tx, err := db.BeginTx(context.Background(), nil)
288317
if err != nil {
@@ -351,14 +380,6 @@ func createDatabaseFromTemplate(t TBSubset, connParams ConnectionParams, newDBNa
351380
return xerrors.Errorf("rename tmp template db: %w", err)
352381
}
353382
}
354-
355-
// Try to create the database again now that a template exists.
356-
if _, err = db.Exec("CREATE DATABASE " + newDBName + " WITH TEMPLATE " + templateDBName); err != nil {
357-
return xerrors.Errorf("create db with template after migrations: %w", err)
358-
}
359-
if err = tx.Commit(); err != nil {
360-
return xerrors.Errorf("commit tx: %w", err)
361-
}
362383
return nil
363384
}
364385

‎coderd/database/dbtestutil/postgres_test.go

Lines changed: 25 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -3,6 +3,7 @@ package dbtestutil_test
33
import (
44
"database/sql"
55
"testing"
6+
"time"
67

78
_ "github.com/lib/pq"
89
"github.com/stretchr/testify/require"
@@ -110,3 +111,27 @@ func TestOpen_ValidDBFrom(t *testing.T) {
110111
require.True(t, rows.Next())
111112
require.NoError(t, rows.Close())
112113
}
114+
115+
func TestOpen_Panic(t *testing.T) {
116+
t.Skip("unskip this to manually test that we don't leak a database into postgres")
117+
t.Parallel()
118+
if !dbtestutil.WillUsePostgres() {
119+
t.Skip("this test requires postgres")
120+
}
121+
122+
_, err := dbtestutil.Open(t)
123+
require.NoError(t, err)
124+
panic("now check SELECT datname FROM pg_database;")
125+
}
126+
127+
func TestOpen_Timeout(t *testing.T) {
128+
t.Skip("unskip this and set a short timeout to manually test that we don't leak a database into postgres")
129+
t.Parallel()
130+
if !dbtestutil.WillUsePostgres() {
131+
t.Skip("this test requires postgres")
132+
}
133+
134+
_, err := dbtestutil.Open(t)
135+
require.NoError(t, err)
136+
time.Sleep(11 * time.Minute)
137+
}

0 commit comments

Comments
 (0)