From 51d6d1a3a6ebb367a44ad795090f9111a6d97d2f Mon Sep 17 00:00:00 2001
From: Jack Christensen <jack@jackchristensen.com>
Date: Sat, 23 May 2015 12:46:36 -0500
Subject: [PATCH] Add ParseEnvLibpq with support for basic envvars

---
 conn.go      | 31 +++++++++++++++++++++++
 conn_test.go | 70 ++++++++++++++++++++++++++++++++++++++++++++++++++++
 2 files changed, 101 insertions(+)

diff --git a/conn.go b/conn.go
index c9937f6e..5b5bda5e 100644
--- a/conn.go
+++ b/conn.go
@@ -331,6 +331,37 @@ func ParseDSN(s string) (ConnConfig, error) {
 	return cp, nil
 }
 
+// ParseEnvLibpq parses the environment like libpq does into a ConnConfig
+//
+// See http://www.postgresql.org/docs/9.4/static/libpq-envars.html for details
+// on the meaning of environment variables.
+//
+// ParseEnvLibpq currently recognizes the following environment variables:
+// PGHOST
+// PGPORT
+// PGDATABASE
+// PGUSER
+// PGPASSWORD
+func ParseEnvLibpq() (ConnConfig, error) {
+	var cc ConnConfig
+
+	cc.Host = os.Getenv("PGHOST")
+
+	if pgport := os.Getenv("PGPORT"); pgport != "" {
+		if port, err := strconv.ParseUint(pgport, 10, 16); err == nil {
+			cc.Port = uint16(port)
+		} else {
+			return cc, err
+		}
+	}
+
+	cc.Database = os.Getenv("PGDATABASE")
+	cc.User = os.Getenv("PGUSER")
+	cc.Password = os.Getenv("PGPASSWORD")
+
+	return cc, nil
+}
+
 // Prepare creates a prepared statement with name and sql. sql can contain placeholders
 // for bound parameters. These placeholders are referenced positional as $1, $2, etc.
 func (c *Conn) Prepare(name, sql string) (ps *PreparedStatement, err error) {
diff --git a/conn_test.go b/conn_test.go
index 2a89c3a7..005347d3 100644
--- a/conn_test.go
+++ b/conn_test.go
@@ -5,6 +5,7 @@ import (
 	"fmt"
 	"github.com/jackc/pgx"
 	"net"
+	"os"
 	"reflect"
 	"strconv"
 	"strings"
@@ -364,6 +365,75 @@ func TestParseDSN(t *testing.T) {
 	}
 }
 
+func TestParseEnvLibpq(t *testing.T) {
+	pgEnvvars := []string{"PGHOST", "PGPORT", "PGDATABASE", "PGUSER", "PGPASSWORD"}
+
+	savedEnv := make(map[string]string)
+	for _, n := range pgEnvvars {
+		savedEnv[n] = os.Getenv(n)
+	}
+	defer func() {
+		for k, v := range savedEnv {
+			err := os.Setenv(k, v)
+			if err != nil {
+				t.Fatalf("Unable to restore environment:", err)
+			}
+		}
+	}()
+
+	tests := []struct {
+		envvars map[string]string
+		config  pgx.ConnConfig
+	}{
+		{
+			envvars: map[string]string{},
+			config:  pgx.ConnConfig{},
+		},
+		{
+			envvars: map[string]string{
+				"PGHOST":     "123.123.123.123",
+				"PGPORT":     "7777",
+				"PGDATABASE": "foo",
+				"PGUSER":     "bar",
+				"PGPASSWORD": "baz",
+			},
+			config: pgx.ConnConfig{
+				Host:     "123.123.123.123",
+				Port:     7777,
+				Database: "foo",
+				User:     "bar",
+				Password: "baz",
+			},
+		},
+	}
+
+	for i, tt := range tests {
+		for _, n := range pgEnvvars {
+			err := os.Unsetenv(n)
+			if err != nil {
+				t.Fatalf("%d. Unable to clear environment:", i, err)
+			}
+		}
+
+		for k, v := range tt.envvars {
+			err := os.Setenv(k, v)
+			if err != nil {
+				t.Fatalf("%d. Unable to set environment:", i, err)
+			}
+		}
+
+		config, err := pgx.ParseEnvLibpq()
+		if err != nil {
+			t.Errorf("%d. Unexpected error from pgx.ParseLibpq() => %v", i, err)
+			continue
+		}
+
+		if !reflect.DeepEqual(config, tt.config) {
+			t.Errorf("%d. expected %#v got %#v", i, tt.config, config)
+		}
+	}
+}
+
 func TestExec(t *testing.T) {
 	t.Parallel()