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

Skip to content

Commit ee8c6e5

Browse files
hugodutkahugodutka
committed
dbpool
1 parent d93a9cf commit ee8c6e5

File tree

4 files changed

+482
-21
lines changed

4 files changed

+482
-21
lines changed

‎coderd/database/dbtestutil/dbpool/main.go‎

Lines changed: 31 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,31 @@
1+
package dbpool
2+
3+
import "net/rpc"
4+
5+
type Client struct {
6+
rpcClient *rpc.Client
7+
}
8+
9+
func NewClient(addr string) (*Client, error) {
10+
rpcClient, err := rpc.DialHTTP("tcp", addr)
11+
if err != nil {
12+
return nil, err
13+
}
14+
return &Client{rpcClient: rpcClient}, nil
15+
}
16+
17+
func (c *Client) GetDB() (string, error) {
18+
var arg int
19+
var reply string
20+
err := c.rpcClient.Call("DBPool.GetDB", &arg, &reply)
21+
return reply, err
22+
}
23+
24+
func (c *Client) DisposeDB(dbURL string) error {
25+
var reply int
26+
return c.rpcClient.Call("DBPool.DisposeDB", &dbURL, &reply)
27+
}
28+
29+
func (c *Client) Close() error {
30+
return c.rpcClient.Close()
31+
}

‎coderd/database/dbtestutil/postgres.go‎

Lines changed: 92 additions & 21 deletions
Original file line numberDiff line numberDiff line change
@@ -8,6 +8,7 @@ import (
88
"errors"
99
"fmt"
1010
"net"
11+
"net/url"
1112
"os"
1213
"path/filepath"
1314
"strconv"
@@ -21,6 +22,7 @@ import (
2122
"github.com/ory/dockertest/v3/docker"
2223
"golang.org/x/xerrors"
2324

25+
"github.com/coder/coder/v2/coderd/database/dbtestutil/dbpool"
2426
"github.com/coder/coder/v2/coderd/database/migrations"
2527
"github.com/coder/coder/v2/cryptorand"
2628
"github.com/coder/retry"
@@ -38,6 +40,39 @@ func (p ConnectionParams) DSN() string {
3840
return fmt.Sprintf("postgres://%s:%s@%s:%s/%s?sslmode=disable", p.Username, p.Password, p.Host, p.Port, p.DBName)
3941
}
4042

43+
func ParseDSN(dsn string) (ConnectionParams, error) {
44+
u, err := url.Parse(dsn)
45+
if err != nil {
46+
return ConnectionParams{}, xerrors.Errorf("parse dsn: %w", err)
47+
}
48+
49+
if u.Scheme != "postgres" {
50+
return ConnectionParams{}, xerrors.Errorf("invalid dsn scheme: %s", u.Scheme)
51+
}
52+
53+
var params ConnectionParams
54+
if u.User != nil {
55+
params.Username = u.User.Username()
56+
params.Password, _ = u.User.Password()
57+
}
58+
59+
params.Host = u.Hostname()
60+
params.Port = u.Port()
61+
if params.Port == "" {
62+
// Default PostgreSQL port
63+
params.Port = "5432"
64+
}
65+
66+
// The path includes a leading slash, remove it.
67+
if len(u.Path) > 1 {
68+
params.DBName = u.Path[1:]
69+
} else {
70+
return ConnectionParams{}, xerrors.New("database name missing in dsn")
71+
}
72+
73+
return params, nil
74+
}
75+
4176
// These variables are global because all tests share them.
4277
var (
4378
connectionParamsInitOnce sync.Once
@@ -138,24 +173,75 @@ type TBSubset interface {
138173
Logf(format string, args ...any)
139174
}
140175

176+
func RemoveDB(t TBSubset, dbName string) error {
177+
cleanupDbURL := defaultConnectionParams.DSN()
178+
cleanupConn, err := sql.Open("postgres", cleanupDbURL)
179+
if err != nil {
180+
return xerrors.Errorf("cleanup database %q: failed to connect to postgres: %w", dbName, err)
181+
}
182+
defer func() {
183+
if err := cleanupConn.Close(); err != nil {
184+
t.Logf("cleanup database %q: failed to close connection: %s\n", dbName, err.Error())
185+
}
186+
}()
187+
_, err = cleanupConn.Exec("DROP DATABASE " + dbName + ";")
188+
if err != nil {
189+
return xerrors.Errorf("cleanup database %q: failed to drop database: %w", dbName, err)
190+
}
191+
return nil
192+
}
193+
194+
func getDBPoolClient() (*dbpool.Client, error) {
195+
dbpoolURL := os.Getenv("DBPOOL")
196+
if dbpoolURL == "" {
197+
return nil, nil //nolint:nilnil
198+
}
199+
client, err := dbpool.NewClient(dbpoolURL)
200+
if err != nil {
201+
return nil, xerrors.Errorf("create db pool client: %w", err)
202+
}
203+
return client, nil
204+
}
205+
141206
// Open creates a new PostgreSQL database instance.
142207
// If there's a database running at localhost:5432, it will use that.
143208
// Otherwise, it will start a new postgres container.
144209
func Open(t TBSubset, opts ...OpenOption) (string, error) {
145210
t.Helper()
146211

212+
openOptions := OpenOptions{}
213+
for _, opt := range opts {
214+
opt(&openOptions)
215+
}
216+
217+
if openOptions.DBFrom == nil {
218+
dbPoolClient, err := getDBPoolClient()
219+
if err != nil {
220+
return "", xerrors.Errorf("get db pool client: %w", err)
221+
}
222+
if dbPoolClient != nil {
223+
dbURL, err := dbPoolClient.GetDB()
224+
if err != nil {
225+
return "", xerrors.Errorf("get db from pool: %w", err)
226+
}
227+
t.Cleanup(func() {
228+
defer dbPoolClient.Close()
229+
err := dbPoolClient.DisposeDB(dbURL)
230+
if err != nil {
231+
t.Logf("cleanup database %s: failed to dispose db: %+v\n", dbURL, err)
232+
}
233+
})
234+
return dbURL, nil
235+
}
236+
}
237+
147238
connectionParamsInitOnce.Do(func() {
148239
errDefaultConnectionParamsInit = initDefaultConnection(t)
149240
})
150241
if errDefaultConnectionParamsInit != nil {
151242
return "", xerrors.Errorf("init default connection params: %w", errDefaultConnectionParamsInit)
152243
}
153244

154-
openOptions := OpenOptions{}
155-
for _, opt := range opts {
156-
opt(&openOptions)
157-
}
158-
159245
var (
160246
username = defaultConnectionParams.Username
161247
password = defaultConnectionParams.Password
@@ -182,22 +268,7 @@ func Open(t TBSubset, opts ...OpenOption) (string, error) {
182268
}
183269

184270
t.Cleanup(func() {
185-
cleanupDbURL := defaultConnectionParams.DSN()
186-
cleanupConn, err := sql.Open("postgres", cleanupDbURL)
187-
if err != nil {
188-
t.Logf("cleanup database %q: failed to connect to postgres: %s\n", dbName, err.Error())
189-
return
190-
}
191-
defer func() {
192-
if err := cleanupConn.Close(); err != nil {
193-
t.Logf("cleanup database %q: failed to close connection: %s\n", dbName, err.Error())
194-
}
195-
}()
196-
_, err = cleanupConn.Exec("DROP DATABASE " + dbName + ";")
197-
if err != nil {
198-
t.Logf("failed to clean up database %q: %s\n", dbName, err.Error())
199-
return
200-
}
271+
RemoveDB(t, dbName)
201272
})
202273

203274
dsn := ConnectionParams{

0 commit comments

Comments
 (0)