From d288033211e53178811a71fd9c6b08e7bc1d253d Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Juho=20M=C3=A4kinen?= <juho@squareup.com>
Date: Fri, 29 Nov 2024 11:01:06 +1100
Subject: [PATCH] feat: PG Permission + region (#3567)

---
 .../postgres.go                               | 11 +++++++++-
 internal/dsn/dsn.go                           | 20 ++++++++++++++++---
 internal/dsn/dsn_test.go                      | 13 ++++++++++++
 3 files changed, 40 insertions(+), 4 deletions(-)
 create mode 100644 internal/dsn/dsn_test.go

diff --git a/cmd/ftl-provisioner-cloudformation/postgres.go b/cmd/ftl-provisioner-cloudformation/postgres.go
index 774960e2ea..ad14214225 100644
--- a/cmd/ftl-provisioner-cloudformation/postgres.go
+++ b/cmd/ftl-provisioner-cloudformation/postgres.go
@@ -93,7 +93,16 @@ func PostgresPostUpdate(ctx context.Context, secrets *secretsmanager.Client, byN
 					return fmt.Errorf("failed to create database: %w", err)
 				}
 			}
-			if _, err := db.ExecContext(ctx, "GRANT ALL ON SCHEMA public TO ftluser; GRANT ALL PRIVILEGES ON DATABASE "+resourceID+" TO ftluser;"); err != nil {
+			if _, err := db.ExecContext(ctx, fmt.Sprintf(`
+				GRANT CONNECT ON DATABASE %s TO ftluser;
+				GRANT USAGE ON SCHEMA public TO ftluser;
+				GRANT USAGE ON SCHEMA public TO ftluser;
+                GRANT CREATE ON SCHEMA public TO ftluser;
+				GRANT ALL PRIVILEGES ON ALL TABLES IN SCHEMA public TO ftluser;
+				GRANT ALL PRIVILEGES ON ALL SEQUENCES IN SCHEMA public TO ftluser;
+				ALTER DEFAULT PRIVILEGES IN SCHEMA public GRANT ALL ON TABLES TO ftluser;
+				ALTER DEFAULT PRIVILEGES IN SCHEMA public GRANT ALL ON SEQUENCES TO ftluser;
+			`, resourceID)); err != nil {
 				return fmt.Errorf("failed to grant FTL user privileges: %w", err)
 			}
 		}
diff --git a/internal/dsn/dsn.go b/internal/dsn/dsn.go
index ba7e3ad84e..8792c4708c 100644
--- a/internal/dsn/dsn.go
+++ b/internal/dsn/dsn.go
@@ -4,6 +4,7 @@ import (
 	"context"
 	"fmt"
 	"net"
+	"strings"
 
 	"github.com/aws/aws-sdk-go-v2/config"
 	"github.com/aws/aws-sdk-go-v2/feature/rds/auth"
@@ -58,9 +59,12 @@ func ResolvePostgresDSN(ctx context.Context, connector schema.DatabaseConnector)
 			return "", fmt.Errorf("configuration error: %w", err)
 		}
 
-		authenticationToken, err := auth.BuildAuthToken(
-			// TODO: proper region
-			ctx, c.Endpoint, "us-west-2", c.Username, cfg.Credentials)
+		region, err := parseRegionFromEndpoint(c.Endpoint)
+		if err != nil {
+			return "", fmt.Errorf("failed to parse region from endpoint: %w", err)
+		}
+
+		authenticationToken, err := auth.BuildAuthToken(ctx, c.Endpoint, region, c.Username, cfg.Credentials)
 		if err != nil {
 			return "", fmt.Errorf("failed to create authentication token: %w", err)
 		}
@@ -74,6 +78,16 @@ func ResolvePostgresDSN(ctx context.Context, connector schema.DatabaseConnector)
 	}
 }
 
+func parseRegionFromEndpoint(endpoint string) (string, error) {
+	host, _, err := net.SplitHostPort(endpoint)
+	if err != nil {
+		return "", fmt.Errorf("failed to split host and port: %w", err)
+	}
+	host = strings.TrimSuffix(host, ".rds.amazonaws.com")
+	parts := strings.Split(host, ".")
+	return parts[len(parts)-1], nil
+}
+
 func ResolveMySQLDSN(ctx context.Context, connector schema.DatabaseConnector) (string, error) {
 	dsnRuntime, ok := connector.(*schema.DSNDatabaseConnector)
 	if !ok {
diff --git a/internal/dsn/dsn_test.go b/internal/dsn/dsn_test.go
new file mode 100644
index 0000000000..417efd6fe3
--- /dev/null
+++ b/internal/dsn/dsn_test.go
@@ -0,0 +1,13 @@
+package dsn
+
+import (
+	"testing"
+
+	"github.com/alecthomas/assert/v2"
+)
+
+func TestParseRegionFromEndpoint(t *testing.T) {
+	region, err := parseRegionFromEndpoint("ftl-alice-dbxcluster-rms7cnlwyggg.cluster-cr24kso0s7in.us-west-2.rds.amazonaws.com:5432")
+	assert.NoError(t, err)
+	assert.Equal(t, "us-west-2", region)
+}