fix(pg): add tls host support for pgdump

Author: Amogh-Bharadwaj

flow/connectors/postgres/pgdump_schema.go

@@ -233,6 +233,13 @@ func runPipeline(
 	return nil
 }
 
+func tlsHostOrHost(config *protos.PostgresConfig) string {
+	if config.TlsHost != "" {
+		return config.TlsHost
+	}
+	return config.Host
+}
+
 func buildPgDumpArgs(config *protos.PostgresConfig) []string {
 	port := config.Port
 	if port == 0 {
@@ -243,7 +250,7 @@ func buildPgDumpArgs(config *protos.PostgresConfig) []string {
 		"--schema-only",
 		"--no-owner",
 		"--no-privileges",
-		"-h", config.Host,
+		"-h", tlsHostOrHost(config),
 		"-p", strconv.FormatUint(uint64(port), 10),
 		"-d", config.Database,
 	}
@@ -260,7 +267,7 @@ func buildPsqlArgs(config *protos.PostgresConfig) []string {
 	}
 
 	args := []string{
-		"-h", config.Host,
+		"-h", tlsHostOrHost(config),
 		"-p", strconv.FormatUint(uint64(port), 10),
 		"-d", config.Database,
 		// Wrap the entire dump in a single transaction so partial failures
@@ -282,8 +289,6 @@ func buildPsqlArgs(config *protos.PostgresConfig) []string {
 
 func appendTLSEnv(ctx context.Context, cmd *exec.Cmd, config *protos.PostgresConfig) {
 	hasRootCA := config.RootCa != nil && *config.RootCa != ""
-	// Match the regular connector: enable TLS when explicitly required OR when
-	// a root CA is provided (which implies the user wants a verified connection).
 	if !internal.PGMustUseTlsConnection(config) && !hasRootCA {
 		return
 	}
@@ -294,13 +299,10 @@ func appendTLSEnv(ctx context.Context, cmd *exec.Cmd, config *protos.PostgresCon
 	case hasRootCA:
 		cmd.Env = append(cmd.Env, "PGSSLMODE=verify-ca")
 	default:
-		// TLS required but no CA provided — encrypt the connection without
-		// attempting certificate verification (matches pgx "require" behaviour).
 		cmd.Env = append(cmd.Env, "PGSSLMODE=require")
 	}
 
 	if hasRootCA {
-		// write root CA to a temp file
 		tmpFile, err := os.CreateTemp("", "peerdb-root-ca-*.pem")
 		if err != nil {
 			slog.WarnContext(ctx, "failed to create temp file for root CA, skipping sslrootcert", slog.Any("error", err))
@@ -314,6 +316,9 @@ func appendTLSEnv(ctx context.Context, cmd *exec.Cmd, config *protos.PostgresCon
 		}
 		tmpFile.Close()
 		cmd.Env = append(cmd.Env, "PGSSLROOTCERT="+tmpFile.Name())
-		// note: temp file is cleaned up when the process exits
+	}
+
+	if config.TlsHost != "" {
+		cmd.Env = append(cmd.Env, "PGHOSTADDR="+config.Host)
 	}
 }

flow/connectors/postgres/pgdump_schema_test.go

@@ -319,6 +319,20 @@ func TestBuildPgDumpArgs(t *testing.T) {
 				"-h", "remote", "-p", "6543", "-d", "prod",
 			},
 		},
+		{
+			name: "TlsHost overrides -h",
+			config: &protos.PostgresConfig{
+				Host:     "10.0.0.1",
+				Port:     5432,
+				Database: "mydb",
+				User:     "admin",
+				TlsHost:  "db.example.com",
+			},
+			wantArgs: []string{
+				"--schema-only", "--no-owner", "--no-privileges",
+				"-h", "db.example.com", "-p", "5432", "-d", "mydb", "-U", "admin",
+			},
+		},
 	}
 
 	for _, tt := range tests {
@@ -367,6 +381,23 @@ func TestBuildPsqlArgs(t *testing.T) {
 				"--quiet",
 			},
 		},
+		{
+			name: "TlsHost overrides -h",
+			config: &protos.PostgresConfig{
+				Host:     "10.0.0.1",
+				Port:     5432,
+				Database: "db",
+				User:     "u",
+				TlsHost:  "db.example.com",
+			},
+			wantArgs: []string{
+				"-h", "db.example.com", "-p", "5432", "-d", "db",
+				"--single-transaction",
+				"-v", "ON_ERROR_STOP=1",
+				"--quiet",
+				"-U", "u",
+			},
+		},
 	}
 
 	for _, tt := range tests {
@@ -515,6 +546,68 @@ func TestAppendTLSEnv(t *testing.T) {
 	}
 }
 
+func TestAppendTLSEnv_TlsHost(t *testing.T) {
+	ctx := t.Context()
+
+	t.Run("PGHOSTADDR set when TlsHost configured", func(t *testing.T) {
+		cmd := exec.CommandContext(ctx, "echo") //nolint:gosec
+		cmd.Env = os.Environ()
+
+		config := &protos.PostgresConfig{
+			Host: "10.0.0.1", Port: 5432, Database: "d",
+			RequireTls: true,
+			TlsHost:    "db.example.com",
+		}
+		appendTLSEnv(ctx, cmd, config)
+
+		if got := envValue(cmd, "PGHOSTADDR"); got != "10.0.0.1" {
+			t.Errorf("PGHOSTADDR = %q, want %q", got, "10.0.0.1")
+		}
+		if got := envValue(cmd, "PGSSLMODE"); got != "require" {
+			t.Errorf("PGSSLMODE = %q, want %q", got, "require")
+		}
+	})
+
+	t.Run("PGHOSTADDR not set when TlsHost empty", func(t *testing.T) {
+		cmd := exec.CommandContext(ctx, "echo") //nolint:gosec
+		cmd.Env = os.Environ()
+
+		config := &protos.PostgresConfig{
+			Host: "10.0.0.1", Port: 5432, Database: "d",
+			RequireTls: true,
+		}
+		appendTLSEnv(ctx, cmd, config)
+
+		if got := envValue(cmd, "PGHOSTADDR"); got != "" {
+			t.Errorf("PGHOSTADDR should be empty, got %q", got)
+		}
+	})
+
+	t.Run("PGHOSTADDR with root CA and TlsHost", func(t *testing.T) {
+		cmd := exec.CommandContext(ctx, "echo") //nolint:gosec
+		cmd.Env = os.Environ()
+
+		config := &protos.PostgresConfig{
+			Host: "10.0.0.1", Port: 5432, Database: "d",
+			RequireTls: true,
+			TlsHost:    "db.example.com",
+			RootCa:     strPtr("-----BEGIN CERTIFICATE-----\ntest\n-----END CERTIFICATE-----"),
+		}
+		appendTLSEnv(ctx, cmd, config)
+
+		if got := envValue(cmd, "PGHOSTADDR"); got != "10.0.0.1" {
+			t.Errorf("PGHOSTADDR = %q, want %q", got, "10.0.0.1")
+		}
+		if got := envValue(cmd, "PGSSLMODE"); got != "verify-ca" {
+			t.Errorf("PGSSLMODE = %q, want %q", got, "verify-ca")
+		}
+		// clean up temp file
+		if f := envValue(cmd, "PGSSLROOTCERT"); f != "" {
+			os.Remove(f)
+		}
+	})
+}
+
 func TestIncompatibleLineRegex(t *testing.T) {
 	tests := []struct {
 		line  string