CXH-2379: fix grant/revoke idempotency for DDL-based engines

.github/workflows/ci.yaml

@@ -41,14 +41,14 @@ jobs:
       - name: Build baton-sql
         run: go build ./cmd/baton-sql
       - name: Run sync tests
-        uses: ConductorOne/github-workflows/actions/sync-test@v2
+        uses: ConductorOne/github-workflows/actions/sync-test@v4
         with:
           connector: ./baton-sql
           baton-entitlement: 'role:admin:member'
           baton-principal: john.smith
           baton-principal-type: user
       - name: Run account provisioning tests
-        uses: ConductorOne/github-workflows/actions/account-provisioning@v3
+        uses: ConductorOne/github-workflows/actions/account-provisioning@v4
         with:
           connector: ./baton-sql
           account-email: robert.tables2@example.com

pkg/bsql/config.go

@@ -422,7 +422,15 @@ type EntitlementProvisioningQueries struct {
 	// NoTransaction indicates whether the provisioning queries should be executed without a transaction.
 	NoTransaction bool `yaml:"no_transaction,omitempty" json:"no_transaction,omitempty"`
 
-	// ValidationQueries is a list of SQL statements to execute for validating the provisioning operation before execution.
+	// ValidationQueries is a list of SQL statements run before the provisioning queries.
+	// On engines that report rows-affected, a query returning no rows fails the operation
+	// (an existence precondition). On DDL-based engines (Db2) that don't report rows-affected,
+	// a query returning no rows instead means the state is already as desired, so the operation
+	// is reported as an idempotent success (GrantAlreadyExists / GrantAlreadyRevoked).
+	//
+	// Warning: on DDL-based engines, do NOT use these as existence preconditions
+	// (e.g. "does this user/role exist?"). A no-rows result is reported as idempotent
+	// success, so a missing or mistyped principal is silently swallowed instead of erroring.
 	ValidationQueries []string `yaml:"validation_queries,omitempty" json:"validation_queries,omitempty"`
 
 	// Queries is a list of SQL statements to execute for the provisioning operation.

pkg/bsql/provisioning.go

@@ -88,7 +88,8 @@ func (s *SQLSyncer) Grant(ctx context.Context, principal *v2.Resource, entitleme
 	if err != nil {
 		if errors.Is(err, ErrQueryAffectedZeroRows) {
 			l.Debug("entitlement is already granted", zap.String("entitlement_id", entitlement.GetId()))
-			anno := annotations.Annotations{}
+			// Reuse the returned annotations so a GrantReplaced from an already-committed
+			// grant_replace revoke survives; a fresh set would drop it.
 			anno.Update(&v2.GrantAlreadyExists{})
 			return anno, nil
 		}

pkg/bsql/provisioning_validation_idempotency_test.go

@@ -0,0 +1,165 @@
+package bsql
+
+import (
+	"testing"
+
+	v2 "github.com/conductorone/baton-sdk/pb/c1/connector/v2"
+	sdkGrant "github.com/conductorone/baton-sdk/pkg/types/grant"
+	"github.com/conductorone/baton-sql/pkg/database"
+	"github.com/stretchr/testify/require"
+)
+
+// grantValidationQuery returns a row only while the membership is absent, mirroring
+// the DDL-dialect pattern where the validation query is the "is there work to do?" gate.
+const grantValidationQuery = `SELECT 1 FROM users u WHERE u.id = ?<principal_id> AND NOT EXISTS (SELECT 1 FROM user_roles WHERE user_id = ?<principal_id> AND role = ?<role>)`
+
+// revokeValidationQuery returns a row only while the membership is present.
+const revokeValidationQuery = `SELECT 1 FROM user_roles WHERE user_id = ?<principal_id> AND role = ?<role>`
+
+func withValidationQueryConfig(s *SQLSyncer) {
+	s.config = ResourceType{
+		StaticEntitlements: []*EntitlementMapping{
+			{
+				Id: "member",
+				Provisioning: &EntitlementProvisioning{
+					Vars: map[string]string{
+						"principal_id": "principal.ID",
+						"role":         "resource.ID",
+					},
+					Grant: &GrantEntitlementProvisioningQueries{
+						EntitlementProvisioningQueries: EntitlementProvisioningQueries{
+							ValidationQueries: []string{grantValidationQuery},
+							Queries:           []string{`INSERT INTO user_roles (user_id, role) VALUES (?<principal_id>, ?<role>)`},
+						},
+					},
+					Revoke: &RevokeEntitlementProvisioningQueries{
+						EntitlementProvisioningQueries: EntitlementProvisioningQueries{
+							ValidationQueries: []string{revokeValidationQuery},
+							Queries:           []string{`DELETE FROM user_roles WHERE user_id = ?<principal_id> AND role = ?<role>`},
+						},
+					},
+				},
+			},
+		},
+	}
+}
+
+func memberEntitlementFor(role string) *v2.Entitlement {
+	roleResource := &v2.Resource{Id: &v2.ResourceId{ResourceType: "role", Resource: role}}
+	principal := &v2.Resource{Id: &v2.ResourceId{ResourceType: "user", Resource: "unused"}}
+	return sdkGrant.NewGrant(roleResource, "member", principal).GetEntitlement()
+}
+
+func userPrincipal(userID string) *v2.Resource {
+	return &v2.Resource{Id: &v2.ResourceId{ResourceType: "user", Resource: userID}}
+}
+
+func TestGrant_ValidationNoRowsReportsAlreadyExists(t *testing.T) {
+	s, db := newRevokeProvisioningTestSyncer(t)
+	withValidationQueryConfig(s)
+	// validation "no rows" only signals idempotency on DDL engines (Db2)
+	s.dbEngine = database.DB2
+	// membership already present: the grant validation query returns no rows
+	seedUserWithRoles(t, db, "user-1", "admin")
+
+	annos, err := s.Grant(t.Context(), userPrincipal("user-1"), memberEntitlementFor("admin"))
+	require.NoError(t, err)
+
+	ok, err := annos.Pick(&v2.GrantAlreadyExists{})
+	require.NoError(t, err)
+	require.True(t, ok)
+
+	// the INSERT never ran, so no duplicate row was created
+	require.Equal(t, 1, countRows(t, db, `SELECT COUNT(*) FROM user_roles WHERE user_id = ? AND role = ?`, "user-1", "admin"))
+}
+
+// On a non-DDL engine, validation "no rows" is a failed precondition, not idempotency:
+// Grant must return an error rather than reporting GrantAlreadyExists.
+func TestGrant_ValidationNoRowsOnNonDDLEngineFailsLoudly(t *testing.T) {
+	s, db := newRevokeProvisioningTestSyncer(t)
+	withValidationQueryConfig(s)
+	// membership already present: the grant validation query returns no rows
+	seedUserWithRoles(t, db, "user-1", "admin")
+
+	annos, err := s.Grant(t.Context(), userPrincipal("user-1"), memberEntitlementFor("admin"))
+	require.Error(t, err)
+	require.Nil(t, annos)
+
+	// the INSERT never ran, so no duplicate row was created
+	require.Equal(t, 1, countRows(t, db, `SELECT COUNT(*) FROM user_roles WHERE user_id = ? AND role = ?`, "user-1", "admin"))
+}
+
+func TestGrant_ValidationRowsAppliesGrant(t *testing.T) {
+	s, db := newRevokeProvisioningTestSyncer(t)
+	withValidationQueryConfig(s)
+	// user exists without the role: validation returns a row, grant proceeds
+	seedUserWithRoles(t, db, "user-1")
+
+	annos, err := s.Grant(t.Context(), userPrincipal("user-1"), memberEntitlementFor("admin"))
+	require.NoError(t, err)
+
+	ok, err := annos.Pick(&v2.GrantAlreadyExists{})
+	require.NoError(t, err)
+	require.False(t, ok)
+
+	require.Equal(t, 1, countRows(t, db, `SELECT COUNT(*) FROM user_roles WHERE user_id = ? AND role = ?`, "user-1", "admin"))
+}
+
+func TestRevoke_ValidationNoRowsReportsAlreadyRevoked(t *testing.T) {
+	s, _ := newRevokeProvisioningTestSyncer(t)
+	withValidationQueryConfig(s)
+	// validation "no rows" only signals idempotency on DDL engines (Db2)
+	s.dbEngine = database.DB2
+	// nothing seeded: the revoke validation query returns no rows
+
+	annos, err := s.Revoke(t.Context(), revokeGrantFor("user-1", "admin"))
+	require.NoError(t, err)
+
+	ok, err := annos.Pick(&v2.GrantAlreadyRevoked{})
+	require.NoError(t, err)
+	require.True(t, ok)
+}
+
+// On a non-DDL engine, validation "no rows" is a failed precondition, not idempotency:
+// Revoke must return an error rather than reporting GrantAlreadyRevoked.
+func TestRevoke_ValidationNoRowsOnNonDDLEngineFailsLoudly(t *testing.T) {
+	s, _ := newRevokeProvisioningTestSyncer(t)
+	withValidationQueryConfig(s)
+	// nothing seeded: the revoke validation query returns no rows
+
+	annos, err := s.Revoke(t.Context(), revokeGrantFor("user-1", "admin"))
+	require.Error(t, err)
+	require.Nil(t, annos)
+}
+
+func TestRevoke_ValidationRowsAppliesRevoke(t *testing.T) {
+	s, db := newRevokeProvisioningTestSyncer(t)
+	withValidationQueryConfig(s)
+	seedUserWithRoles(t, db, "user-1", "admin")
+
+	annos, err := s.Revoke(t.Context(), revokeGrantFor("user-1", "admin"))
+	require.NoError(t, err)
+
+	ok, err := annos.Pick(&v2.GrantAlreadyRevoked{})
+	require.NoError(t, err)
+	require.False(t, ok)
+
+	require.Equal(t, 0, countRows(t, db, `SELECT COUNT(*) FROM user_roles WHERE user_id = ? AND role = ?`, "user-1", "admin"))
+}
+
+// The revoke helper must map validation "no rows" onto the sentinel so the caller
+// can detect idempotency with errors.Is.
+func TestRunProvisioningQueriesWithExecutor_ValidationNoRowsWrapsSentinel(t *testing.T) {
+	s, db := newRevokeProvisioningTestSyncer(t)
+	// validation "no rows" only signals idempotency on DDL engines (Db2)
+	s.dbEngine = database.DB2
+
+	err := s.RunProvisioningQueriesWithExecutor(
+		t.Context(),
+		[]string{`DELETE FROM user_roles WHERE user_id = ?<principal_id>`},
+		[]string{revokeValidationQuery},
+		map[string]any{"principal_id": "user-1", "role": "admin"},
+		db,
+	)
+	require.ErrorIs(t, err, ErrQueryAffectedZeroRows)
+}

pkg/bsql/query.go

@@ -602,9 +602,18 @@ func (s *SQLSyncer) runPrincipalExistsCheck(
 	return exists, nil
 }
 
-func (s *SQLSyncer) RunProvisioningQueriesWithExecutor(
+// validationNoRowsMeansIdempotent reports whether a validation query returning no
+// rows should be treated as "already in the desired state" rather than a failed
+// precondition. Only DDL-based engines (Db2) need this: their GRANT/REVOKE statements
+// don't report rows-affected, so the validation query is the only zero-effect signal
+// available. Engines that report rows-affected keep using validation queries as
+// existence preconditions that fail loudly.
+func (s *SQLSyncer) validationNoRowsMeansIdempotent() bool {
+	return s.dbEngine == database.DB2
+}
+
+func (s *SQLSyncer) runValidationQueries(
 	ctx context.Context,
-	queries,
 	validationQueries []string,
 	vars map[string]any,
 	executor executor,
@@ -632,19 +641,38 @@ func (s *SQLSyncer) RunProvisioningQueriesWithExecutor(
 		valid := result.Next()
 
 		if err := result.Err(); err != nil {
+			_ = result.Close()
 			return fmt.Errorf("failed to read validation query result: %w", err)
 		}
 
-		err = result.Close()
-		if err != nil {
+		if err := result.Close(); err != nil {
 			return fmt.Errorf("failed to close validation query result: %w", err)
 		}
 
 		if !valid {
+			if s.validationNoRowsMeansIdempotent() {
+				return fmt.Errorf("validation query returned no rows: %w", ErrQueryAffectedZeroRows)
+			}
 			return fmt.Errorf("validation query returned no rows")
 		}
 	}
 
+	return nil
+}
+
+func (s *SQLSyncer) RunProvisioningQueriesWithExecutor(
+	ctx context.Context,
+	queries,
+	validationQueries []string,
+	vars map[string]any,
+	executor executor,
+) error {
+	l := ctxzap.Extract(ctx)
+
+	if err := s.runValidationQueries(ctx, validationQueries, vars, executor); err != nil {
+		return err
+	}
+
 	zeroRowCount := 0
 
 	for idx, q := range queries {
@@ -1047,32 +1075,8 @@ func (s *SQLSyncer) RunGrantProvisioning(
 		}
 	}
 
-	for _, q := range validationQueries {
-		q, qArgs, err := s.prepareProvisioningQuery(q, vars)
-		if err != nil {
-			return anno, fmt.Errorf("failed to prepare validation query: %w", err)
-		}
-
-		result, err := executor.QueryContext(ctx, q, qArgs...)
-		if err != nil {
-			return anno, fmt.Errorf("failed to execute validation query: %w", err)
-		}
-
-		valid := result.Next()
-
-		if err := result.Err(); err != nil {
-			_ = result.Close()
-			return anno, fmt.Errorf("failed to read validation query result: %w", err)
-		}
-
-		err = result.Close()
-		if err != nil {
-			return anno, fmt.Errorf("failed to close validation query result: %w", err)
-		}
-
-		if !valid {
-			return anno, fmt.Errorf("grant provisioning: validation query returned no rows")
-		}
+	if err := s.runValidationQueries(ctx, validationQueries, vars, executor); err != nil {
+		return anno, err
 	}
 
 	zeroRowCount := 0