addressing review comments

Author: tsivaprasad

server/internal/orchestrator/swarm/orchestrator.go

@@ -541,12 +541,10 @@ func (o *Orchestrator) generateMCPInstanceResources(spec *database.ServiceInstan
 func (o *Orchestrator) generateRAGInstanceResources(spec *database.ServiceInstanceSpec) (*database.ServiceInstanceResources, error) {
 	// RAG service user role (per-host, not replicated by Spock)
 	ragUserRole := &RAGServiceUserRole{
-		ServiceInstanceID: spec.ServiceInstanceID,
-		ServiceID:         spec.ServiceSpec.ServiceID,
-		DatabaseID:        spec.DatabaseID,
-		DatabaseName:      spec.DatabaseName,
-		HostID:            spec.HostID,
-		NodeName:          spec.NodeName,
+		ServiceID:    spec.ServiceSpec.ServiceID,
+		DatabaseID:   spec.DatabaseID,
+		DatabaseName: spec.DatabaseName,
+		NodeName:     spec.NodeName,
 	}
 	if spec.Credentials != nil {
 		ragUserRole.Username = spec.Credentials.Username

server/internal/orchestrator/swarm/rag_service_user_role.go

@@ -2,17 +2,12 @@ package swarm
 
 import (
 	"context"
-	"errors"
 	"fmt"
-	"time"
 
-	"github.com/jackc/pgx/v5"
 	"github.com/rs/zerolog"
 	"github.com/samber/do"
 
-	"github.com/pgEdge/control-plane/server/internal/certificates"
 	"github.com/pgEdge/control-plane/server/internal/database"
-	"github.com/pgEdge/control-plane/server/internal/patroni"
 	"github.com/pgEdge/control-plane/server/internal/postgres"
 	"github.com/pgEdge/control-plane/server/internal/resource"
 	"github.com/pgEdge/control-plane/server/internal/utils"
@@ -22,24 +17,24 @@ var _ resource.Resource = (*RAGServiceUserRole)(nil)
 
 const ResourceTypeRAGServiceUserRole resource.Type = "swarm.rag_service_user_role"
 
-func RAGServiceUserRoleIdentifier(serviceInstanceID string) resource.Identifier {
+func RAGServiceUserRoleIdentifier(serviceID string) resource.Identifier {
 	return resource.Identifier{
-		ID:   serviceInstanceID,
+		ID:   serviceID,
 		Type: ResourceTypeRAGServiceUserRole,
 	}
 }
 
+// RAGServiceUserRole manages the Postgres role for a RAG service.
 // The role is created on the primary of the co-located Postgres instance
-// (same HostID) and granted the pgedge_application_read_only built-in role.
+// and granted the pgedge_application_read_only built-in role.
+// Spock replicates the role to every other node because we connect via r.DatabaseName.
 type RAGServiceUserRole struct {
-	ServiceInstanceID string `json:"service_instance_id"`
-	ServiceID         string `json:"service_id"`
-	DatabaseID        string `json:"database_id"`
-	DatabaseName      string `json:"database_name"`
-	HostID            string `json:"host_id"`   // Used to find the co-located Postgres instance
-	NodeName          string `json:"node_name"` // Database node name for PrimaryExecutor routing
-	Username          string `json:"username"`
-	Password          string `json:"password"` // Generated on Create, persisted in state
+	ServiceID    string `json:"service_id"`
+	DatabaseID   string `json:"database_id"`
+	DatabaseName string `json:"database_name"`
+	NodeName     string `json:"node_name"` // Database node name for PrimaryExecutor routing
+	Username     string `json:"username"`
+	Password     string `json:"password"` // Generated on Create, persisted in state
 }
 
 func (r *RAGServiceUserRole) ResourceVersion() string {
@@ -55,7 +50,7 @@ func (r *RAGServiceUserRole) DiffIgnore() []string {
 }
 
 func (r *RAGServiceUserRole) Identifier() resource.Identifier {
-	return RAGServiceUserRoleIdentifier(r.ServiceInstanceID)
+	return RAGServiceUserRoleIdentifier(r.ServiceID)
 }
 
 func (r *RAGServiceUserRole) Executor() resource.Executor {
@@ -80,28 +75,26 @@ func (r *RAGServiceUserRole) Refresh(ctx context.Context, rc *resource.Context)
 		return err
 	}
 	logger = logger.With().
-		Str("service_instance_id", r.ServiceInstanceID).
+		Str("service_id", r.ServiceID).
 		Str("database_id", r.DatabaseID).
 		Logger()
 
-	conn, err := r.connectToColocatedPrimary(ctx, rc, logger, r.DatabaseName)
+	primary, err := database.GetPrimaryInstance(ctx, rc, r.NodeName)
 	if err != nil {
-		logger.Warn().Err(err).Msg("could not connect to verify RAG role existence, assuming it exists")
-		return nil
+		return fmt.Errorf("failed to get primary instance: %w", err)
+	}
+	conn, err := primary.Connection(ctx, rc, r.DatabaseName)
+	if err != nil {
+		return fmt.Errorf("failed to connect to database: %w", err)
 	}
 	defer conn.Close(ctx)
 
-	var exists bool
-	err = conn.QueryRow(ctx,
-		"SELECT EXISTS(SELECT 1 FROM pg_catalog.pg_roles WHERE rolname = $1)",
-		r.Username,
-	).Scan(&exists)
+	needsCreate, err := postgres.UserRoleNeedsCreate(r.Username).Scalar(ctx, conn)
 	if err != nil {
-		// On query failure, assume it exists
-		logger.Warn().Err(err).Msg("pg_roles query failed, assuming RAG role exists")
-		return nil
+		logger.Warn().Err(err).Msg("pg_roles query failed")
+		return fmt.Errorf("pg_roles query failed: %w", err)
 	}
-	if !exists {
+	if needsCreate {
 		return resource.ErrNotFound
 	}
 	return nil
@@ -113,30 +106,36 @@ func (r *RAGServiceUserRole) Create(ctx context.Context, rc *resource.Context) e
 		return err
 	}
 	logger = logger.With().
-		Str("service_instance_id", r.ServiceInstanceID).
+		Str("service_id", r.ServiceID).
 		Str("database_id", r.DatabaseID).
 		Logger()
 	logger.Info().Msg("creating RAG service user role")
 
-	r.Username = database.GenerateServiceUsername(r.ServiceInstanceID)
-	password, err := utils.RandomString(32)
-	if err != nil {
-		return fmt.Errorf("failed to generate password: %w", err)
+	r.Username = database.GenerateServiceUsername(r.ServiceID)
+	if r.Password == "" {
+		password, err := utils.RandomString(32)
+		if err != nil {
+			return fmt.Errorf("failed to generate password: %w", err)
+		}
+		r.Password = password
 	}
-	r.Password = password
 
-	if err := r.createRole(ctx, rc, logger); err != nil {
+	if err := r.createRole(ctx, rc); err != nil {
 		return fmt.Errorf("failed to create RAG service user role: %w", err)
 	}
 
 	logger.Info().Str("username", r.Username).Msg("RAG service user role created successfully")
 	return nil
 }
 
-func (r *RAGServiceUserRole) createRole(ctx context.Context, rc *resource.Context, logger zerolog.Logger) error {
-	conn, err := r.connectToColocatedPrimary(ctx, rc, logger, r.DatabaseName)
+func (r *RAGServiceUserRole) createRole(ctx context.Context, rc *resource.Context) error {
+	primary, err := database.GetPrimaryInstance(ctx, rc, r.NodeName)
 	if err != nil {
-		return err
+		return fmt.Errorf("failed to get primary instance: %w", err)
+	}
+	conn, err := primary.Connection(ctx, rc, r.DatabaseName)
+	if err != nil {
+		return fmt.Errorf("failed to connect to database: %w", err)
 	}
 	defer conn.Close(ctx)
 
@@ -169,16 +168,22 @@ func (r *RAGServiceUserRole) Delete(ctx context.Context, rc *resource.Context) e
 		return err
 	}
 	logger = logger.With().
-		Str("service_instance_id", r.ServiceInstanceID).
+		Str("service_id", r.ServiceID).
 		Str("database_id", r.DatabaseID).
 		Str("username", r.Username).
 		Logger()
 	logger.Info().Msg("deleting RAG service user from database")
 
-	conn, err := r.connectToColocatedPrimary(ctx, rc, logger, "postgres")
+	primary, err := database.GetPrimaryInstance(ctx, rc, r.NodeName)
+	if err != nil {
+		// During deletion the database may already be gone or unreachable.
+		logger.Warn().Err(err).Msg("failed to get primary instance, skipping RAG user deletion")
+		return nil
+	}
+	conn, err := primary.Connection(ctx, rc, r.DatabaseName)
 	if err != nil {
 		// During deletion the database may already be gone or unreachable.
-		logger.Warn().Err(err).Msg("failed to connect to co-located primary, skipping RAG user deletion")
+		logger.Warn().Err(err).Msg("failed to connect to database, skipping RAG user deletion")
 		return nil
 	}
 	defer conn.Close(ctx)
@@ -192,94 +197,3 @@ func (r *RAGServiceUserRole) Delete(ctx context.Context, rc *resource.Context) e
 	logger.Info().Msg("RAG service user deleted successfully")
 	return nil
 }
-
-// connectToColocatedPrimary finds the primary Postgres instance on the same
-// host as this RAG service instance and returns an authenticated connection.
-// Filtering by HostID ensures the role is created on the correct node, since
-// CREATE ROLE is not replicated by Spock in a multi-active setup.
-func (r *RAGServiceUserRole) connectToColocatedPrimary(ctx context.Context, rc *resource.Context, logger zerolog.Logger, dbName string) (*pgx.Conn, error) {
-	dbSvc, err := do.Invoke[*database.Service](rc.Injector)
-	if err != nil {
-		return nil, err
-	}
-
-	primaryInstanceID, err := r.resolveColocatedPrimary(ctx, dbSvc, logger)
-	if err != nil {
-		return nil, err
-	}
-
-	connInfo, err := dbSvc.GetInstanceConnectionInfo(ctx, r.DatabaseID, primaryInstanceID)
-	if err != nil {
-		return nil, fmt.Errorf("failed to get instance connection info: %w", err)
-	}
-
-	certSvc, err := do.Invoke[*certificates.Service](rc.Injector)
-	if err != nil {
-		return nil, fmt.Errorf("failed to get certificate service: %w", err)
-	}
-
-	tlsConfig, err := certSvc.PostgresUserTLS(ctx, primaryInstanceID, connInfo.InstanceHostname, "pgedge")
-	if err != nil {
-		return nil, fmt.Errorf("failed to create TLS config: %w", err)
-	}
-
-	conn, err := database.ConnectToInstance(ctx, &database.ConnectionOptions{
-		DSN: connInfo.AdminDSN(dbName),
-		TLS: tlsConfig,
-	})
-	if err != nil {
-		return nil, fmt.Errorf("failed to connect to database: %w", err)
-	}
-
-	return conn, nil
-}
-
-// resolveColocatedPrimary fetches the database, selects co-located instances,
-// and returns the primary instance ID via Patroni.
-func (r *RAGServiceUserRole) resolveColocatedPrimary(ctx context.Context, dbSvc *database.Service, logger zerolog.Logger) (string, error) {
-	db, err := dbSvc.GetDatabase(ctx, r.DatabaseID)
-	if err != nil {
-		if errors.Is(err, database.ErrDatabaseNotFound) {
-			return "", fmt.Errorf("database not found: %w", err)
-		}
-		return "", fmt.Errorf("failed to get database: %w", err)
-	}
-	if len(db.Instances) == 0 {
-		return "", fmt.Errorf("database has no instances")
-	}
-	candidates := r.colocatedInstances(db.Instances, logger)
-	return r.findPrimaryAmong(ctx, dbSvc, candidates, logger), nil
-}
-
-// colocatedInstances returns the subset of instances that share r.HostID.
-// Falls back to all instances if none are co-located.
-func (r *RAGServiceUserRole) colocatedInstances(all []*database.Instance, logger zerolog.Logger) []*database.Instance {
-	candidates := make([]*database.Instance, 0, len(all))
-	for _, inst := range all {
-		if inst.HostID == r.HostID {
-			candidates = append(candidates, inst)
-		}
-	}
-	if len(candidates) == 0 {
-		logger.Warn().Str("host_id", r.HostID).Msg("no co-located Postgres instances found, falling back to all instances")
-		return all
-	}
-	return candidates
-}
-
-// findPrimaryAmong queries Patroni for each candidate and returns the primary
-// instance ID. Falls back to the first candidate if none can be determined.
-func (r *RAGServiceUserRole) findPrimaryAmong(ctx context.Context, dbSvc *database.Service, candidates []*database.Instance, logger zerolog.Logger) string {
-	for _, inst := range candidates {
-		connInfo, err := dbSvc.GetInstanceConnectionInfo(ctx, r.DatabaseID, inst.InstanceID)
-		if err != nil {
-			continue
-		}
-		primaryID, err := database.GetPrimaryInstanceID(ctx, patroni.NewClient(connInfo.PatroniURL(), nil), 10*time.Second)
-		if err == nil && primaryID != "" {
-			return primaryID
-		}
-	}
-	logger.Warn().Msg("could not determine primary instance, using first co-located instance")
-	return candidates[0].InstanceID
-}

server/internal/orchestrator/swarm/rag_service_user_role_test.go

@@ -16,10 +16,10 @@ func TestRAGServiceUserRole_ResourceVersion(t *testing.T) {
 }
 
 func TestRAGServiceUserRole_Identifier(t *testing.T) {
-	r := &RAGServiceUserRole{ServiceInstanceID: "db1-rag-host1"}
+	r := &RAGServiceUserRole{ServiceID: "rag"}
 	id := r.Identifier()
-	if id.ID != "db1-rag-host1" {
-		t.Errorf("Identifier().ID = %q, want %q", id.ID, "db1-rag-host1")
+	if id.ID != "rag" {
+		t.Errorf("Identifier().ID = %q, want %q", id.ID, "rag")
 	}
 	if id.Type != ResourceTypeRAGServiceUserRole {
 		t.Errorf("Identifier().Type = %q, want %q", id.Type, ResourceTypeRAGServiceUserRole)
@@ -65,9 +65,9 @@ func TestRAGServiceUserRole_RefreshEmptyCredentials(t *testing.T) {
 	for _, tt := range tests {
 		t.Run(tt.name, func(t *testing.T) {
 			r := &RAGServiceUserRole{
-				ServiceInstanceID: "inst1",
-				Username:          tt.username,
-				Password:          tt.password,
+				ServiceID: "rag",
+				Username:  tt.username,
+				Password:  tt.password,
 			}
 			// Refresh with nil rc — the empty-credential guard fires before any
 			// injection call, so no injector is needed.

server/internal/postgres/roles.go

@@ -11,6 +11,17 @@ import (
 var defaultSchemas = []string{"public", "spock", "pg_catalog", "information_schema"}
 var builtinRoles = []string{"pgedge_application", "pgedge_application_read_only", "pgedge_superuser"}
 
+// UserRoleNeedsCreate returns a query that evaluates to true when the named
+// role does not yet exist in pg_catalog.pg_roles.
+func UserRoleNeedsCreate(name string) Query[bool] {
+	return Query[bool]{
+		SQL: "SELECT NOT EXISTS (SELECT 1 FROM pg_catalog.pg_roles WHERE rolname = @name);",
+		Args: pgx.NamedArgs{
+			"name": name,
+		},
+	}
+}
+
 type UserRoleOptions struct {
 	Name       string
 	Password   string