diff --git a/sqldb/postgres.go b/sqldb/postgres.go index f5bf79d65..cf52f7a72 100644 --- a/sqldb/postgres.go +++ b/sqldb/postgres.go @@ -6,6 +6,9 @@ import ( "database/sql" "encoding/hex" "fmt" + "net/url" + "path" + "strings" "testing" "time" @@ -15,10 +18,6 @@ import ( "github.com/stretchr/testify/require" ) -const ( - dsnTemplate = "postgres://%v:%v@%v:%d/%v?sslmode=%v" -) - var ( // DefaultPostgresFixtureLifetime is the default maximum time a Postgres // test fixture is being kept alive. After that time the docker @@ -32,31 +31,65 @@ var ( // //nolint:lll type PostgresConfig struct { - SkipMigrations bool `long:"skipmigrations" description:"Skip applying migrations on startup."` - Host string `long:"host" description:"Database server hostname."` - Port int `long:"port" description:"Database server port."` - User string `long:"user" description:"Database user."` - Password string `long:"password" description:"Database user's password."` - DBName string `long:"dbname" description:"Database name to use."` - MaxOpenConnections int `long:"maxconnections" description:"Max open connections to keep alive to the database server."` - RequireSSL bool `long:"requiressl" description:"Whether to require using SSL (mode: require) when connecting to the server."` + Dsn string `long:"dsn" description:"Database connection string."` + Timeout time.Duration `long:"timeout" description:"Database connection timeout. Set to zero to disable."` + MaxConnections int `long:"maxconnections" description:"The maximum number of open connections to the database. Set to zero for unlimited."` + SkipMigrations bool `long:"skipmigrations" description:"Skip applying migrations on startup."` } -// DSN returns the dns to connect to the database. -func (s *PostgresConfig) DSN(hidePassword bool) string { - var sslMode = "disable" - if s.RequireSSL { - sslMode = "require" +func (p *PostgresConfig) Validate() error { + if p.Dsn == "" { + return fmt.Errorf("DSN is required") } - password := s.Password - if hidePassword { - // Placeholder used for logging the DSN safely. - password = "****" + // Parse the DSN as a URL. + _, err := url.Parse(p.Dsn) + if err != nil { + return fmt.Errorf("invalid DSN: %w", err) } - return fmt.Sprintf(dsnTemplate, s.User, password, s.Host, s.Port, - s.DBName, sslMode) + return nil +} + +// replacePasswordInDSN takes a DSN string and returns it with the password +// replaced by "***". +func replacePasswordInDSN(dsn string) (string, error) { + // Parse the DSN as a URL + u, err := url.Parse(dsn) + if err != nil { + return "", err + } + + // Check if the URL has a user info part + if u.User != nil { + username := u.User.Username() + + // Reconstruct user info with "***" as password + userInfo := username + ":***@" + + // Rebuild the DSN with the modified user info + sanitizeDSN := strings.Replace( + dsn, u.User.String()+"@", userInfo, 1, + ) + + return sanitizeDSN, nil + } + + // Return the original DSN if no user info is present + return dsn, nil +} + +// getDatabaseNameFromDSN extracts the database name from a DSN string. +func getDatabaseNameFromDSN(dsn string) (string, error) { + // Parse the DSN as a URL + u, err := url.Parse(dsn) + if err != nil { + return "", err + } + + // The database name is the last segment of the path. Trim leading slash + // and return the last segment. + return path.Base(u.Path), nil } // PostgresStore is a database store implementation that uses a Postgres @@ -70,16 +103,25 @@ type PostgresStore struct { // NewPostgresStore creates a new store that is backed by a Postgres database // backend. func NewPostgresStore(cfg *PostgresConfig) (*PostgresStore, error) { - log.Infof("Using SQL database '%s'", cfg.DSN(true)) + sanitizedDSN, err := replacePasswordInDSN(cfg.Dsn) + if err != nil { + return nil, err + } + log.Infof("Using SQL database '%s'", sanitizedDSN) - rawDB, err := sql.Open("pgx", cfg.DSN(false)) + dbName, err := getDatabaseNameFromDSN(cfg.Dsn) + if err != nil { + return nil, err + } + + rawDB, err := sql.Open("pgx", cfg.Dsn) if err != nil { return nil, err } maxConns := defaultMaxConns - if cfg.MaxOpenConnections > 0 { - maxConns = cfg.MaxOpenConnections + if cfg.MaxConnections > 0 { + maxConns = cfg.MaxConnections } rawDB.SetMaxOpenConns(maxConns) @@ -108,7 +150,7 @@ func NewPostgresStore(cfg *PostgresConfig) (*PostgresStore, error) { }) err = applyMigrations( - postgresFS, driver, "sqlc/migrations", cfg.DBName, + postgresFS, driver, "sqlc/migrations", dbName, ) if err != nil { return nil, err @@ -149,9 +191,7 @@ func NewTestPostgresDB(t *testing.T, fixture *TestPgFixture) *PostgresStore { t.Fatal(err) } - cfg := fixture.GetConfig() - cfg.DBName = dbName - + cfg := fixture.GetConfig(dbName) store, err := NewPostgresStore(cfg) require.NoError(t, err) diff --git a/sqldb/postgres_fixture.go b/sqldb/postgres_fixture.go index 313684f96..be897657b 100644 --- a/sqldb/postgres_fixture.go +++ b/sqldb/postgres_fixture.go @@ -1,7 +1,6 @@ package sqldb import ( - "context" "database/sql" "fmt" "strconv" @@ -75,7 +74,7 @@ func NewTestPgFixture(t *testing.T, expiry time.Duration) *TestPgFixture { host: host, port: int(port), } - databaseURL := fixture.GetDSN() + databaseURL := fixture.GetConfig(testPgDBName).Dsn log.Infof("Connecting to Postgres fixture: %v\n", databaseURL) // Tell docker to hard kill the container in "expiry" seconds. @@ -104,20 +103,13 @@ func NewTestPgFixture(t *testing.T, expiry time.Duration) *TestPgFixture { return fixture } -// GetDSN returns the DSN (Data Source Name) for the started Postgres node. -func (f *TestPgFixture) GetDSN() string { - return f.GetConfig().DSN(false) -} - // GetConfig returns the full config of the Postgres node. -func (f *TestPgFixture) GetConfig() *PostgresConfig { +func (f *TestPgFixture) GetConfig(dbName string) *PostgresConfig { return &PostgresConfig{ - Host: f.host, - Port: f.port, - User: testPgUser, - Password: testPgPass, - DBName: testPgDBName, - RequireSSL: false, + Dsn: fmt.Sprintf( + "postgres://%v:%v@%v:%v/%v?sslmode=disable", + testPgUser, testPgPass, f.host, f.port, dbName, + ), } } @@ -126,16 +118,3 @@ func (f *TestPgFixture) TearDown(t *testing.T) { err := f.pool.Purge(f.resource) require.NoError(t, err, "Could not purge resource") } - -// ClearDB clears the database. -func (f *TestPgFixture) ClearDB(t *testing.T) { - dbConn, err := sql.Open("postgres", f.GetDSN()) - require.NoError(t, err) - - _, err = dbConn.ExecContext( - context.Background(), - `DROP SCHEMA IF EXISTS public CASCADE; - CREATE SCHEMA public;`, - ) - require.NoError(t, err) -} diff --git a/sqldb/sqlite.go b/sqldb/sqlite.go index 4b9895999..9b2d08487 100644 --- a/sqldb/sqlite.go +++ b/sqldb/sqlite.go @@ -40,13 +40,11 @@ const ( // //nolint:lll type SqliteConfig struct { - // SkipMigrations if true, then all the tables will be created on start - // up if they don't already exist. - SkipMigrations bool `long:"skipmigrations" description:"Skip applying migrations on startup."` - - // DatabaseFileName is the full file path where the database file can be - // found. - DatabaseFileName string `long:"dbfile" description:"The full path to the database."` + Timeout time.Duration `long:"timeout" description:"The time after which a database query should be timed out."` + BusyTimeout time.Duration `long:"busytimeout" description:"The maximum amount of time to wait for a database connection to become available for a query."` + MaxConnections int `long:"maxconnections" description:"The maximum number of open connections to the database. Set to zero for unlimited."` + PragmaOptions []string `long:"pragmaoptions" description:"A list of pragma options to set on a database connection. For example, 'auto_vacuum=incremental'. Note that the flag must be specified multiple times if multiple options are to be set."` + SkipMigrations bool `long:"skipmigrations" description:"Skip applying migrations on startup."` } // SqliteStore is a database store implementation that uses a sqlite backend. @@ -58,7 +56,7 @@ type SqliteStore struct { // NewSqliteStore attempts to open a new sqlite database based on the passed // config. -func NewSqliteStore(cfg *SqliteConfig) (*SqliteStore, error) { +func NewSqliteStore(cfg *SqliteConfig, dbPath string) (*SqliteStore, error) { // The set of pragma options are accepted using query options. For now // we only want to ensure that foreign key constraints are properly // enforced. @@ -107,7 +105,7 @@ func NewSqliteStore(cfg *SqliteConfig) (*SqliteStore, error) { // details on the formatting here, see the modernc.org/sqlite docs: // https://pkg.go.dev/modernc.org/sqlite#Driver.Open. dsn := fmt.Sprintf( - "%v?%v&%v", cfg.DatabaseFileName, sqliteOptions.Encode(), + "%v?%v&%v", dbPath, sqliteOptions.Encode(), sqliteTxLockImmediate, ) db, err := sql.Open("sqlite", dsn) @@ -171,9 +169,8 @@ func NewTestSqliteDB(t *testing.T) *SqliteStore { // an in mem version to speed up tests dbFileName := filepath.Join(t.TempDir(), "tmp.db") sqlDB, err := NewSqliteStore(&SqliteConfig{ - DatabaseFileName: dbFileName, - SkipMigrations: false, - }) + SkipMigrations: false, + }, dbFileName) require.NoError(t, err) t.Cleanup(func() {