diff --git a/db/connection.go b/db/connection.go index b2384b1..6cf1b60 100644 --- a/db/connection.go +++ b/db/connection.go @@ -2,6 +2,8 @@ package db import ( "database/sql" + "net/url" + "strings" "sync" ) @@ -9,8 +11,11 @@ var ( dbPool = map[string]*sql.DB{} cmux sync.Mutex - //DriverName defaults to postgres + // DriverName defaults to postgres driverName = "postgres" + + // Connection params to be appended to the connection string + connParams string ) // ConnFn is the database connection builder function that @@ -39,11 +44,22 @@ func ConnectionFn(url string, opts ...connectionOption) ConnFn { v() } - conn, err := sql.Open(driverName, url) + // Modify the URL to include connection params + // if any. + modURL := url + if strings.Contains(modURL, "?") { + modURL = modURL + "&" + connParams + } else if connParams != "" { + modURL = modURL + "?" + connParams + } + + conn, err := sql.Open(driverName, modURL) if err != nil { return nil, err } + // This uses url instead of modURL while the params apply + // to all connections. dbPool[url] = conn return conn, nil @@ -57,3 +73,24 @@ func WithDriver(name string) connectionOption { driverName = name } } + +// Params allows to specify additional connection parameters +// that will be encoded as URL params next to the connection string. +// params should be in key,value,key,value,... format. +// e.g Params("sslmode", "disable", "timezone", "UTC") +func Params(params ...string) connectionOption { + vals := url.Values{} + for i := 0; i < len(params); i += 2 { + key := params[i] + if i+1 >= len(params) { + vals.Add(key, "") + break + } + + vals.Add(key, params[i+1]) + } + + return func() { + connParams = vals.Encode() + } +} diff --git a/db/connection_test.go b/db/connection_test.go index 6519127..fa718d7 100644 --- a/db/connection_test.go +++ b/db/connection_test.go @@ -11,17 +11,17 @@ func TestConnection(t *testing.T) { urls := []string{"test_1.db", "test_2.db", "test_3.db"} createDatabases := func() { - for _, dbUrl := range urls { - if err := db.Create(dbUrl); err != nil { + for _, url := range urls { + if err := db.Create(url); err != nil { t.Errorf("error creating db: %v", err) } } - } + t.Run("Correct - creating multiple connections", func(t *testing.T) { t.Cleanup(func() { - for _, dbUrl := range urls { - if err := db.Drop(dbUrl); err != nil { + for _, url := range urls { + if err := db.Drop(url); err != nil { t.Errorf("error dropping db: %v", err) } } @@ -29,8 +29,8 @@ func TestConnection(t *testing.T) { createDatabases() - for _, dbUrl := range urls { - conFn := db.ConnectionFn(dbUrl, db.WithDriver("sqlite3")) + for _, url := range urls { + conFn := db.ConnectionFn(url, db.WithDriver("sqlite3")) conn, err := conFn() if err != nil { @@ -42,10 +42,11 @@ func TestConnection(t *testing.T) { } } }) + t.Run("Correct - using current exiting connection if already created", func(t *testing.T) { t.Cleanup(func() { - for _, dbUrl := range urls { - if err := db.Drop(dbUrl); err != nil { + for _, url := range urls { + if err := db.Drop(url); err != nil { t.Errorf("error dropping db: %v", err) } } @@ -55,15 +56,15 @@ func TestConnection(t *testing.T) { var currentConn *sql.DB - for _, dbUrl := range urls { - conFn := db.ConnectionFn(dbUrl, db.WithDriver("sqlite3")) + for _, url := range urls { + conFn := db.ConnectionFn(url, db.WithDriver("sqlite3")) conn, err := conFn() if err != nil { t.Errorf("Expected nil, got err %v", err) } - if dbUrl == urls[0] { + if url == urls[0] { currentConn = conn } @@ -83,4 +84,47 @@ func TestConnection(t *testing.T) { t.Errorf("Expected current connection to be %v, got %v", currentConn, conn) } }) + + t.Run("connection params", func(t *testing.T) { + cases := []string{ + ":memory:", + "file::memory:?cache=shared", + t.TempDir() + "memory.db", + t.TempDir() + "memory.db?mode=memory", + } + + for _, tcase := range cases { + t.Run(tcase, func(t *testing.T) { + connFn := db.ConnectionFn( + tcase, + + db.WithDriver("sqlite3"), + db.Params("_cache_size", "54321"), // Applying cache size parameter to check it later. + ) + + conn, err := connFn() + if err != nil { + t.Errorf("Expected nil, got err %v", err) + } + defer conn.Close() + + rows, err := conn.Query("pragma cache_size;") + if err != nil { + t.Errorf("Expected nil, got err %v", err) + } + defer rows.Close() + + var size int + for rows.Next() { + if err := rows.Scan(&size); err != nil { + t.Errorf("Expected nil, got err %v", err) + } + } + + if size != 54321 { + t.Fatalf("Expected cache size to be 54321, got %v", size) + } + }) + } + }) }