Merge branch 'georgysavva-store-connection-string'

conn.go

@@ -29,6 +29,9 @@ type ConnConfig struct {
 	Logger   Logger
 	LogLevel LogLevel
 
+	// Original connection string that was parsed into config.
+	connString string
+
 	// BuildStatementCache creates the stmtcache.Cache implementation for connections created with this config. Set
 	// to nil to disable automatic prepared statements.
 	BuildStatementCache BuildStatementCacheFunc
@@ -44,6 +47,8 @@ type ConnConfig struct {
 	createdByParseConfig bool // Used to enforce created by ParseConfig rule.
 }
 
+func (cc *ConnConfig) ConnString() string { return cc.connString }
+
 // BuildStatementCacheFunc is a function that can be used to create a stmtcache.Cache implementation for connection.
 type BuildStatementCacheFunc func(conn *pgconn.PgConn) stmtcache.Cache
 
@@ -157,6 +162,7 @@ func ParseConfig(connString string) (*ConnConfig, error) {
 		createdByParseConfig: true,
 		LogLevel:             LogLevelInfo,
 		BuildStatementCache:  buildStatementCache,
+		connString:           connString,
 	}
 
 	return connConfig, nil
@@ -418,6 +424,9 @@ func (c *Conn) StatementCache() stmtcache.Cache { return c.stmtcache }
 // ConnInfo returns the connection info used for this connection.
 func (c *Conn) ConnInfo() *pgtype.ConnInfo { return c.connInfo }
 
+// Config returns config that was used to establish this connection.
+func (c *Conn) Config() *ConnConfig { return c.config }
+
 // Exec executes sql. sql can be either a prepared statement name or an SQL string. arguments should be referenced
 // positionally from the sql string as $1, $2, etc.
 func (c *Conn) Exec(ctx context.Context, sql string, arguments ...interface{}) (pgconn.CommandTag, error) {

conn_test.go

@@ -28,6 +28,8 @@ func TestCrateDBConnect(t *testing.T) {
 	require.Nil(t, err)
 	defer closeConn(t, conn)
 
+	assert.Equal(t, connString, conn.Config().ConnString())
+
 	var result int
 	err = conn.QueryRow(context.Background(), "select 1 +1").Scan(&result)
 	if err != nil {
@@ -41,13 +43,16 @@ func TestCrateDBConnect(t *testing.T) {
 func TestConnect(t *testing.T) {
 	t.Parallel()
 
-	config := mustParseConfig(t, os.Getenv("PGX_TEST_DATABASE"))
+	connString := os.Getenv("PGX_TEST_DATABASE")
+	config := mustParseConfig(t, connString)
 
 	conn, err := pgx.ConnectConfig(context.Background(), config)
 	if err != nil {
 		t.Fatalf("Unable to establish connection: %v", err)
 	}
 
+	assert.Equal(t, connString, conn.Config().ConnString())
+
 	var currentDB string
 	err = conn.QueryRow(context.Background(), "select current_database()").Scan(&currentDB)
 	if err != nil {
@@ -104,6 +109,13 @@ func TestConnectConfigRequiresConnConfigFromParseConfig(t *testing.T) {
 	})
 }
 
+func TestConfigContainsConnStr(t *testing.T) {
+	connStr := os.Getenv("PGX_TEST_DATABASE")
+	config, err := pgx.ParseConfig(connStr)
+	require.NoError(t, err)
+	assert.Equal(t, connStr, config.ConnString())
+}
+
 func TestParseConfigExtractsStatementCacheOptions(t *testing.T) {
 	t.Parallel()
 

pgxpool/pool.go

@@ -69,6 +69,7 @@ func (cr *connResource) getPoolRows(c *Conn, r pgx.Rows) *poolRows {
 
 type Pool struct {
 	p                 *puddle.Pool
+	config            *Config
 	afterConnect      func(context.Context, *pgx.Conn) error
 	beforeAcquire     func(context.Context, *pgx.Conn) bool
 	afterRelease      func(*pgx.Conn) bool
@@ -120,6 +121,8 @@ type Config struct {
 	createdByParseConfig bool // Used to enforce created by ParseConfig rule.
 }
 
+func (c *Config) ConnString() string { return c.ConnConfig.ConnString() }
+
 // Connect creates a new Pool and immediately establishes one connection. ctx can be used to cancel this initial
 // connection. See ParseConfig for information on connString format.
 func Connect(ctx context.Context, connString string) (*Pool, error) {
@@ -141,6 +144,7 @@ func ConnectConfig(ctx context.Context, config *Config) (*Pool, error) {
 	}
 
 	p := &Pool{
+		config:            config,
 		afterConnect:      config.AfterConnect,
 		beforeAcquire:     config.BeforeAcquire,
 		afterRelease:      config.AfterRelease,
@@ -369,6 +373,9 @@ func (p *Pool) AcquireAllIdle(ctx context.Context) []*Conn {
 	return conns
 }
 
+// Config returns config that was used to initialize this pool.
+func (p *Pool) Config() *Config { return p.config }
+
 func (p *Pool) Stat() *Stat {
 	return &Stat{s: p.p.Stat()}
 }

pgxpool/pool_test.go

@@ -14,9 +14,21 @@ import (
 
 func TestConnect(t *testing.T) {
 	t.Parallel()
-
-	pool, err := pgxpool.Connect(context.Background(), os.Getenv("PGX_TEST_DATABASE"))
+	connString := os.Getenv("PGX_TEST_DATABASE")
+	pool, err := pgxpool.Connect(context.Background(), connString)
 	require.NoError(t, err)
+	assert.Equal(t, connString, pool.Config().ConnString())
+	pool.Close()
+}
+
+func TestConnectConfig(t *testing.T) {
+	t.Parallel()
+	connString := os.Getenv("PGX_TEST_DATABASE")
+	config, err := pgxpool.ParseConfig(connString)
+	require.NoError(t, err)
+	pool, err := pgxpool.ConnectConfig(context.Background(), config)
+	require.NoError(t, err)
+	assert.Equal(t, config, pool.Config())
 	pool.Close()
 }