Support setting up proxy via provider attribute proxy_url

postgresql/config.go

@@ -172,6 +172,7 @@ type Config struct {
 	ExpectedVersion   semver.Version
 	SSLClientCert     *ClientCertificateConfig
 	SSLRootCertPath   string
+	ProxyURL          string
 }
 
 // Client struct holding connection string
@@ -279,7 +280,10 @@ func (c *Client) Connect() (*DBConnection, error) {
 		var db *sql.DB
 		var err error
 		if c.config.Scheme == "postgres" {
-			db, err = sql.Open(proxyDriverName, dsn)
+			db = sql.OpenDB(proxyConnector{
+				dsn:      dsn,
+				proxyURL: c.config.ProxyURL,
+			})
 		} else {
 			db, err = postgres.Open(context.Background(), dsn)
 		}

postgresql/provider.go

@@ -182,6 +182,12 @@ func Provider() *schema.Provider {
 				Description:  "Specify the expected version of PostgreSQL.",
 				ValidateFunc: validateExpectedVersion,
 			},
+			"proxy_url": {
+				Type:         schema.TypeString,
+				Optional:     true,
+				Description:  "SOCKS5 proxy URL.",
+				ValidateFunc: validation.IsURLWithScheme([]string{"socks5", "socks5h"}),
+			},
 		},
 
 		ResourcesMap: map[string]*schema.Resource{
@@ -336,6 +342,7 @@ func providerConfigure(d *schema.ResourceData) (interface{}, error) {
 		MaxConns:          d.Get("max_connections").(int),
 		ExpectedVersion:   version,
 		SSLRootCertPath:   d.Get("sslrootcert").(string),
+		ProxyURL:          d.Get("proxy_url").(string),
 	}
 
 	if value, ok := d.GetOk("clientcert"); ok {

postgresql/proxy_driver.go

@@ -4,7 +4,10 @@ import (
 	"context"
 	"database/sql"
 	"database/sql/driver"
+	"fmt"
 	"net"
+	"net/url"
+	"os"
 	"time"
 
 	"github.com/lib/pq"
@@ -13,21 +16,87 @@ import (
 
 const proxyDriverName = "postgresql-proxy"
 
-type proxyDriver struct{}
+type proxyDriver struct {
+	proxyURL string
+}
 
 func (d proxyDriver) Open(name string) (driver.Conn, error) {
 	return pq.DialOpen(d, name)
 }
 
 func (d proxyDriver) Dial(network, address string) (net.Conn, error) {
-	dialer := proxy.FromEnvironment()
+	dialer, err := d.dialer()
+	if err != nil {
+		return nil, err
+	}
 	return dialer.Dial(network, address)
 }
 
 func (d proxyDriver) DialTimeout(network, address string, timeout time.Duration) (net.Conn, error) {
 	ctx, cancel := context.WithTimeout(context.TODO(), timeout)
 	defer cancel()
-	return proxy.Dial(ctx, network, address)
+
+	dialer, err := d.dialer()
+	if err != nil {
+		return nil, err
+	}
+
+	if xd, ok := dialer.(proxy.ContextDialer); ok {
+		return xd.DialContext(ctx, network, address)
+	} else {
+		return nil, fmt.Errorf("unexpected protocol error")
+	}
+}
+
+func (d proxyDriver) dialer() (proxy.Dialer, error) {
+	proxyURL := d.proxyURL
+	if proxyURL == "" {
+		proxyURL = os.Getenv("PGPROXY")
+	}
+	if proxyURL == "" {
+		return proxy.FromEnvironment(), nil
+	}
+
+	u, err := url.Parse(proxyURL)
+	if err != nil {
+		return nil, err
+	}
+
+	dialer, err := proxy.FromURL(u, proxy.Direct)
+	if err != nil {
+		return nil, err
+	}
+
+	noProxy := ""
+	if v := os.Getenv("NO_PROXY"); v != "" {
+		noProxy = v
+	}
+	if v := os.Getenv("no_proxy"); noProxy == "" && v != "" {
+		noProxy = v
+	}
+	if noProxy != "" {
+		perHost := proxy.NewPerHost(dialer, proxy.Direct)
+		perHost.AddFromString(noProxy)
+
+		dialer = perHost
+	}
+
+	return dialer, nil
+}
+
+type proxyConnector struct {
+	dsn      string
+	proxyURL string
+}
+
+var _ driver.Connector = (*proxyConnector)(nil)
+
+func (c proxyConnector) Connect(ctx context.Context) (driver.Conn, error) {
+	return c.Driver().Open(c.dsn)
+}
+
+func (c proxyConnector) Driver() driver.Driver {
+	return proxyDriver{c.proxyURL}
 }
 
 func init() {

postgresql/resource_postgresql_database_test.go

@@ -270,7 +270,7 @@ func checkUserMembership(
 	t *testing.T, dsn, member, role string, shouldHaveRole bool,
 ) resource.TestCheckFunc {
 	return func(s *terraform.State) error {
-		db, err := sql.Open("postgres", dsn)
+		db, err := sql.Open(proxyDriverName, dsn)
 		if err != nil {
 			t.Fatalf("could to create connection pool: %v", err)
 		}

postgresql/resource_postgresql_grant_role_test.go

@@ -141,7 +141,7 @@ func TestAccPostgresqlGrantRole(t *testing.T) {
 
 func checkGrantRole(t *testing.T, dsn, role string, grantRole string, withAdmin bool) resource.TestCheckFunc {
 	return func(s *terraform.State) error {
-		db, err := sql.Open("postgres", dsn)
+		db, err := sql.Open(proxyDriverName, dsn)
 		if err != nil {
 			t.Fatalf("could to create connection pool: %v", err)
 		}

postgresql/resource_postgresql_grant_test.go

@@ -1365,7 +1365,7 @@ func TestAccPostgresqlGrantOwnerPG15(t *testing.T) {
 			// Change the owner to the new pg_database_owner role
 			func() {
 				config := getTestConfig(t)
-				db, err := sql.Open("postgres", config.connStr(dbName))
+				db, err := sql.Open(proxyDriverName, config.connStr(dbName))
 				if err != nil {
 					t.Fatalf("could not connect to database %s: %v", dbName, err)
 				}

postgresql/resource_postgresql_role_test.go

@@ -300,7 +300,7 @@ func testAccCheckRoleCanLogin(t *testing.T, role, password string) resource.Test
 		config := getTestConfig(t)
 		config.Username = role
 		config.Password = password
-		db, err := sql.Open("postgres", config.connStr("postgres"))
+		db, err := sql.Open(proxyDriverName, config.connStr("postgres"))
 		if err != nil {
 			return fmt.Errorf("could not open SQL connection: %v", err)
 		}

postgresql/utils_test.go

@@ -77,7 +77,7 @@ func skipIfNotSuperuser(t *testing.T) {
 
 // dbExecute is a test helper to create a pool, execute one query then close the pool
 func dbExecute(t *testing.T, dsn, query string, args ...interface{}) {
-	db, err := sql.Open("postgres", dsn)
+	db, err := sql.Open(proxyDriverName, dsn)
 	if err != nil {
 		t.Fatalf("could to create connection pool: %v", err)
 	}
@@ -147,7 +147,7 @@ func createTestTables(t *testing.T, dbSuffix string, tables []string, owner stri
 	dbName, _ := getTestDBNames(dbSuffix)
 	adminUser := config.getDatabaseUsername()
 
-	db, err := sql.Open("postgres", config.connStr(dbName))
+	db, err := sql.Open(proxyDriverName, config.connStr(dbName))
 	if err != nil {
 		t.Fatalf("could not open connection pool for db %s: %v", dbName, err)
 	}
@@ -182,7 +182,7 @@ func createTestTables(t *testing.T, dbSuffix string, tables []string, owner stri
 
 	// In this case we need to drop table after each test.
 	return func() {
-		db, err := sql.Open("postgres", config.connStr(dbName))
+		db, err := sql.Open(proxyDriverName, config.connStr(dbName))
 		if err != nil {
 			t.Fatalf("could not open connection pool for db %s: %v", dbName, err)
 		}
@@ -213,7 +213,7 @@ func createTestSchemas(t *testing.T, dbSuffix string, schemas []string, owner st
 	dbName, _ := getTestDBNames(dbSuffix)
 	adminUser := config.getDatabaseUsername()
 
-	db, err := sql.Open("postgres", config.connStr(dbName))
+	db, err := sql.Open(proxyDriverName, config.connStr(dbName))
 	if err != nil {
 		t.Fatalf("could not open connection pool for db %s: %v", dbName, err)
 	}
@@ -248,7 +248,7 @@ func createTestSchemas(t *testing.T, dbSuffix string, schemas []string, owner st
 
 	// In this case we need to drop schema after each test.
 	return func() {
-		db, err := sql.Open("postgres", config.connStr(dbName))
+		db, err := sql.Open(proxyDriverName, config.connStr(dbName))
 		if err != nil {
 			t.Fatalf("could not open connection pool for db %s: %v", dbName, err)
 		}
@@ -278,7 +278,7 @@ func createTestSequences(t *testing.T, dbSuffix string, sequences []string, owne
 	dbName, _ := getTestDBNames(dbSuffix)
 	adminUser := config.getDatabaseUsername()
 
-	db, err := sql.Open("postgres", config.connStr(dbName))
+	db, err := sql.Open(proxyDriverName, config.connStr(dbName))
 	if err != nil {
 		t.Fatalf("could not open connection pool for db %s: %v", dbName, err)
 	}
@@ -312,7 +312,7 @@ func createTestSequences(t *testing.T, dbSuffix string, sequences []string, owne
 	}
 
 	return func() {
-		db, err := sql.Open("postgres", config.connStr(dbName))
+		db, err := sql.Open(proxyDriverName, config.connStr(dbName))
 		if err != nil {
 			t.Fatalf("could not open connection pool for db %s: %v", dbName, err)
 		}
@@ -362,7 +362,7 @@ func connectAsTestRole(t *testing.T, role, dbName string) *sql.DB {
 	config.Username = role
 	config.Password = testRolePassword
 
-	db, err := sql.Open("postgres", config.connStr(dbName))
+	db, err := sql.Open(proxyDriverName, config.connStr(dbName))
 	if err != nil {
 		t.Fatalf("could not open connection pool for db %s: %v", dbName, err)
 	}

tests/docker-compose.yml

@@ -13,9 +13,14 @@ services:
         environment:
             POSTGRES_PASSWORD: ${PGPASSWORD}
         ports:
-            - 25432:5432
+            - "25432:5432"
         healthcheck:
           test: [ "CMD-SHELL", "pg_isready" ]
           interval: 10s
           timeout: 5s
           retries: 5
+
+    proxy:
+        image: ghcr.io/httptoolkit/docker-socks-tunnel
+        ports:
+            - "11080:1080"

tests/switch_proxy.sh

@@ -0,0 +1,10 @@
+#!/bin/bash
+
+export TF_ACC=true
+export PGHOST=postgres
+export PGPORT=5432
+export PGUSER=postgres
+export PGPASSWORD=postgres
+export PGSSLMODE=disable
+export PGSUPERUSER=true
+export PGPROXY=socks5://127.0.0.1:11080

tests/switch_rds.sh

@@ -30,3 +30,4 @@ export PGUSER=rds
 export PGPASSWORD=rds
 export PGSSLMODE=disable
 export PGSUPERUSER=false
+export PGPROXY=""

tests/switch_superuser.sh

@@ -7,3 +7,4 @@ export PGUSER=postgres
 export PGPASSWORD=postgres
 export PGSSLMODE=disable
 export PGSUPERUSER=true
+export PGPROXY=""

tests/testacc_full.sh

@@ -13,7 +13,7 @@ setup() {
 
 run() {
   go test -count=1 ./postgresql -v -timeout 120m
-  
+
   # keep the return value for the scripts to fail and clean properly
   return $?
 }
@@ -31,4 +31,5 @@ run_suite() {
 }
 
 run_suite "superuser"
+run_suite "proxy"
 run_suite "rds"

website/docs/index.html.markdown

@@ -185,6 +185,7 @@ The following arguments are supported:
 * `aws_rds_iam_region` - (Optional) The AWS region to use while using AWS RDS IAM Auth.
 * `azure_identity_auth` - (Optional) If set to `true`, call the Azure OAuth token endpoint for temporary token
 * `azure_tenant_id` - (Optional) (Required if `azure_identity_auth` is `true`) Azure tenant ID [read more](https://registry.terraform.io/providers/hashicorp/azurerm/latest/docs/data-sources/client_config.html)
+* `proxy_url` - (Optional) SOCKS5 proxy URL. Must be a valid URL with schema `socks5` or `socks5h`.
 
 ## GoCloud
 
@@ -310,7 +311,9 @@ provider "postgresql" {
 
 ### SOCKS5 Proxy Support
 
-The provider supports connecting via a SOCKS5 proxy, but when the `postgres` scheme is used. It can be configured by setting the `ALL_PROXY` or `all_proxy` environment variable to a value like `socks5://127.0.0.1:1080`.
+The provider supports connecting via a SOCKS5 proxy, but only when the `postgres` scheme is used. It can be configured
+by setting the `proxy_url` provider attribute, or `PGPROXY`, `ALL_PROXY` or `all_proxy` environment variable to a value
+like `socks5://127.0.0.1:1080`.
 
 The `NO_PROXY` or `no_proxy` environment can also be set to opt out of proxying for specific hostnames or ports.