From 20c6c44f9f5539205829e36f0babad5c7f24037d Mon Sep 17 00:00:00 2001
From: georgysavva <georgy.savva@gmail.com>
Date: Sat, 16 May 2020 19:53:08 +0300
Subject: [PATCH] Expose Conn.Config() and Pool.Config().

---
 conn.go              | 11 ++++++++---
 conn_test.go         | 10 ++++++----
 pgxpool/pool.go      |  9 ++++++---
 pgxpool/pool_test.go | 17 ++++++++++++++---
 4 files changed, 34 insertions(+), 13 deletions(-)

diff --git a/conn.go b/conn.go
index b367a98c..ffa4e66e 100644
--- a/conn.go
+++ b/conn.go
@@ -30,7 +30,7 @@ type ConnConfig struct {
 	LogLevel LogLevel
 
 	// Original connection string that was parsed into config.
-	ConnString string
+	connString string
 
 	// BuildStatementCache creates the stmtcache.Cache implementation for connections created with this config. Set
 	// to nil to disable automatic prepared statements.
@@ -47,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
 
@@ -160,7 +162,7 @@ func ParseConfig(connString string) (*ConnConfig, error) {
 		createdByParseConfig: true,
 		LogLevel:             LogLevelInfo,
 		BuildStatementCache:  buildStatementCache,
-		ConnString:           connString,
+		connString:           connString,
 	}
 
 	return connConfig, nil
@@ -423,7 +425,10 @@ func (c *Conn) StatementCache() stmtcache.Cache { return c.stmtcache }
 func (c *Conn) ConnInfo() *pgtype.ConnInfo { return c.connInfo }
 
 // ConnString returns the connection string that was used to establish this connection.
-func (c *Conn) ConnString() string { return c.config.ConnString }
+func (c *Conn) ConnString() string { return c.config.ConnString() }
+
+// 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.
diff --git a/conn_test.go b/conn_test.go
index c4710fee..5e1cc6a9 100644
--- a/conn_test.go
+++ b/conn_test.go
@@ -28,6 +28,7 @@ func TestCrateDBConnect(t *testing.T) {
 	require.Nil(t, err)
 	defer closeConn(t, conn)
 
+	assert.Equal(t, connString, conn.Config().ConnString())
 	assert.Equal(t, connString, conn.ConnString())
 
 	var result int
@@ -43,15 +44,16 @@ func TestCrateDBConnect(t *testing.T) {
 func TestConnect(t *testing.T) {
 	t.Parallel()
 
-	connStr := os.Getenv("PGX_TEST_DATABASE")
-	config := mustParseConfig(t, connStr)
+	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, connStr, conn.ConnString())
+	assert.Equal(t, connString, conn.Config().ConnString())
+	assert.Equal(t, connString, conn.ConnString())
 
 	var currentDB string
 	err = conn.QueryRow(context.Background(), "select current_database()").Scan(&currentDB)
@@ -113,7 +115,7 @@ 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)
+	assert.Equal(t, connStr, config.ConnString())
 }
 
 func TestParseConfigExtractsStatementCacheOptions(t *testing.T) {
diff --git a/pgxpool/pool.go b/pgxpool/pool.go
index 3cc41ee9..6b0d3bdc 100644
--- a/pgxpool/pool.go
+++ b/pgxpool/pool.go
@@ -69,7 +69,7 @@ func (cr *connResource) getPoolRows(c *Conn, r pgx.Rows) *poolRows {
 
 type Pool struct {
 	p                 *puddle.Pool
-	connString        string
+	config            *Config
 	afterConnect      func(context.Context, *pgx.Conn) error
 	beforeAcquire     func(context.Context, *pgx.Conn) bool
 	afterRelease      func(*pgx.Conn) bool
@@ -142,7 +142,7 @@ func ConnectConfig(ctx context.Context, config *Config) (*Pool, error) {
 	}
 
 	p := &Pool{
-		connString:        config.ConnConfig.ConnString,
+		config:            config,
 		afterConnect:      config.AfterConnect,
 		beforeAcquire:     config.BeforeAcquire,
 		afterRelease:      config.AfterRelease,
@@ -372,7 +372,10 @@ func (p *Pool) AcquireAllIdle(ctx context.Context) []*Conn {
 }
 
 // ConnString returns the connection string that was used to initialize this pool.
-func (p *Pool) ConnString() string { return p.connString }
+func (p *Pool) ConnString() string { return p.config.ConnConfig.ConnString() }
+
+// 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()}
diff --git a/pgxpool/pool_test.go b/pgxpool/pool_test.go
index a47a8edf..e7132d2c 100644
--- a/pgxpool/pool_test.go
+++ b/pgxpool/pool_test.go
@@ -14,10 +14,21 @@ import (
 
 func TestConnect(t *testing.T) {
 	t.Parallel()
-	connStr := os.Getenv("PGX_TEST_DATABASE")
-	pool, err := pgxpool.Connect(context.Background(), connStr)
+	connString := os.Getenv("PGX_TEST_DATABASE")
+	pool, err := pgxpool.Connect(context.Background(), connString)
 	require.NoError(t, err)
-	assert.Equal(t, connStr, pool.ConnString())
+	assert.Equal(t, connString, pool.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()
 }