Bugfix/agent incorrect revoke (#65)

* Change tests to show problems.

* Revised test.

* Fix agent bug

* Simplify logic.

* Rename objectType field to sqlQueryObjectType to be more descriptive.

* Removed account level test

* Reverted rename to objectType and objectTypePlural to avoid merge conflicts

Author: zoom-dustin-moriarty

src/main/java/us/zoom/data/dfence/providers/snowflake/SnowflakeWildcardAllGrantFilter.java

@@ -60,11 +60,11 @@ public static List<SnowflakeGrantBuilder> containerAllGrantBuilders(
                         playbookPrivilegeGrant.databaseName(),
                         playbookPrivilegeGrant.schemaName()));
             }
-            String objectName = String.format("%s.<%s>", containerName, objectType.getObjectType().replace(" ", "_"));
+            String objectName = String.format("%s.<%s>", containerName, objectType.name());
             List<SnowflakeGrantModel> grants = playbookPrivilegeGrant.privileges().stream()
                     .map(p -> new SnowflakeGrantModel(
                             p,
-                            objectType.getObjectType().replace(" ", "_"),
+                            objectType.name(),
                             objectName,
                             "ROLE",
                             roleName,

src/main/java/us/zoom/data/dfence/providers/snowflake/grant/builder/SnowflakeObjectType.java

@@ -5,6 +5,29 @@
 import java.util.Map;
 import java.util.Objects;
 
+/**
+ * Snowflake object types that can be granted on (databases, schemas, tables, agents, etc.).
+ * <p>
+ * Each type has a qualification level ({@link #getQualLevel()}) indicating how many name parts
+ * it has: 0 = account-level, 1 = one part (e.g. database name), 2 = two parts (e.g. db.schema),
+ * 3 = three parts (e.g. db.schema.object). The {@link #getObjectType()} and {@link
+ * #getObjectTypePlural()} strings are used in SQL (e.g. "SHOW AGENTS IN DATABASE"). Use {@link
+ * #name()} when building grant names so desired state matches what Snowflake
+ * returns for the object type in SHOW GRANTS statements. {@link #name()} is also used for matching the object name as
+ * it appears in future grant object names such as MY_DB.MY_SCHEMA.&ltTABLE&gt..
+ * <p>
+ * {@link #getAliasFor()} is used only when building hash keys for grants (e.g. in {@code
+ * SnowflakeGrantBuilder.getKey()} and the revoke index). It is not user-facing and does not
+ * allow the user to specify a different name in the playbook. It exists to align on differences
+ * that come from Snowflake: the same logical object type may appear under different names in
+ * different Snowflake APIs (e.g. MATERIALIZED_VIEW in one place, VIEW in another). By having
+ * multiple enum values return the same alias, they are treated as the same when calculating
+ * which grants to add or remove. For playbook input names (e.g. "AGENT" instead of
+ * "CORTEX_AGENT"), see {@link #overrideObjectTypes}.
+ *
+ * @see #fromString(String) for resolving playbook or Snowflake strings to this enum
+ * @see #overrideObjectTypes for playbook input mapping (e.g. AGENT → CORTEX_AGENT)
+ */
 public enum SnowflakeObjectType {
     ACCOUNT(0, null),
     ALERT(3, null),
@@ -51,30 +74,49 @@ public enum SnowflakeObjectType {
     IMAGE_REPOSITORY(3, null);
 
 
+    /** Number of name parts (0=account, 1=e.g. db, 2=e.g. db.schema, 3=e.g. db.schema.object). */
     @Getter
     private final Integer qualLevel;
 
-    // objectType is used for SQL statements about the object.
+    /** Singular form used in SQL (e.g. "TABLE" will be used in "GRANT SELECT ON TABLE TO ROLE FOO"). */
     @Getter
     private final String objectType;
 
-    // objectTypePlural is used for SQL statements about multiple objects.
+    /** Plural form used in SQL (e.g. TABLES WILL BE USED IN "SHOW GRANTS ON TABLES"). */
     @Getter
     private final String objectTypePlural;
 
-    // Alias for is used for considering grants on different object types equivalent for the sake of hashing and comparison.
+    /**
+     * When non-null, used only in hash keys so that this type and another (e.g. VIEW) are
+     * treated as the same when computing grants to add or remove. Aligns Snowflake variations
+     * (e.g. MATERIALIZED_VIEW may appear as VIEW in some output). Not user-facing; see {@link
+     * #overrideObjectTypes} for playbook input names. However, aliasFor does not change the
+     * underlying enum type. It is only changes how the objectType is used for calculating
+     * grants to add or remove.
+     */
+    @Getter
     private final String aliasFor;
 
-    // Main constructor with all parameters - contains the actual logic
-    SnowflakeObjectType(Integer qualLevel, String aliasFor, String objectType, String objectTypePlural) {
+    /**
+     * Main constructor with all parameters.
+     *
+     * @param qualLevel       number of name parts (0=account, 1=e.g. db, 2=e.g. db.schema,
+     *                        3=e.g. db.schema.object)
+     * @param aliasFor        when non-null, value used in hash keys so this type matches another
+     *                        (e.g. MATERIALIZED_VIEW uses "VIEW"); null if no alias
+     * @param objectType      singular form for SQL (e.g. "AGENT", "TABLE"); null to infer from
+     *                        enum name (underscores to spaces)
+     * @param sqlQueryObjectTypePlural plural form for SQL (e.g. "AGENTS", "TABLES"); null to infer from
+     *                         objectType
+     */
+    SnowflakeObjectType(Integer qualLevel, String aliasFor, String objectType, String sqlQueryObjectTypePlural) {
         this.qualLevel = qualLevel;
-        this.aliasFor = aliasFor;
         // If objectType is provided, use it; otherwise infer from enum name
         String computedObjectType = (objectType != null) ? objectType : this.name().replace("_", " ");
         this.objectType = computedObjectType;
         // If objectTypePlural is provided, use it; otherwise infer from objectType
-        if (objectTypePlural != null) {
-            this.objectTypePlural = objectTypePlural;
+        if (sqlQueryObjectTypePlural != null) {
+            this.objectTypePlural = sqlQueryObjectTypePlural;
         } else {
             // Hooked on phonics works for me.
             if (computedObjectType.endsWith("Y")) {
@@ -83,26 +125,43 @@ public enum SnowflakeObjectType {
                 this.objectTypePlural = computedObjectType + "S";
             }
         }
+        this.aliasFor = Objects.requireNonNullElseGet(aliasFor, this::name);
     }
 
-    // Constructor that infers objectType and objectTypePlural from enum name
+    /**
+     * Constructor that infers objectType and objectTypePlural from enum name (e.g. TABLE → "TABLE",
+     * "TABLES").
+     *
+     * @param qualLevel number of name parts (0=account, 1=e.g. db, 2=e.g. db.schema,
+     *                 3=e.g. db.schema.object)
+     * @param aliasFor when non-null, value used in hash keys so this type matches another; null
+     *                 if no alias
+     */
     SnowflakeObjectType(Integer qualLevel, String aliasFor) {
         this(qualLevel, aliasFor, null, null);
     }
 
-    public String getAliasFor() {
-        if (this.aliasFor == null) {
-            return this.name();
-        }
-        return this.aliasFor;
-    }
-
+    /**
+     * Resolves a string (from playbook, Snowflake grant output, etc.) to this enum.
+     * Applies {@link #overrideObjectTypes} so that e.g. "AGENT" maps to {@link #CORTEX_AGENT}.
+     *
+     * @param objectType the object type string (case-insensitive)
+     * @return the corresponding SnowflakeObjectType
+     * @throws IllegalArgumentException if the string does not match any enum and is not in the
+     *         override map
+     */
     public static SnowflakeObjectType fromString(String objectType) {
         String normalizedObjectType = objectType.toUpperCase();
         String overrideValue = overrideObjectTypes.get(normalizedObjectType);
         return SnowflakeObjectType.valueOf(Objects.requireNonNullElse(overrideValue, normalizedObjectType));
     }
 
+    /**
+     * Map of playbook (or Snowflake) input names to enum names. Used by {@link #fromString(String)}
+     * so that e.g. "AGENT" in a playbook resolves to {@link #CORTEX_AGENT}. This is what allows
+     * the user to specify a different name in the playbook; it is separate from {@link
+     * #getAliasFor()}, which is only for hash keys.
+     */
     public static Map<String, String> overrideObjectTypes = Map.of(
             "AGENT", "CORTEX_AGENT"
     );

src/main/java/us/zoom/data/dfence/providers/snowflake/grant/builder/SnowflakePermissionGrantBuilder.java

@@ -148,7 +148,7 @@ public class SnowflakePermissionGrantBuilder extends SnowflakeGrantBuilder {
                 List.of("MODIFY", "MONITOR", "OPERATE", "USAGE"),
                 List.of(SnowflakeObjectType.COMPUTE_POOL)));
         add(new GrantValidationDefinition(List.of("READ", "WRITE"), List.of(SnowflakeObjectType.IMAGE_REPOSITORY)));
-        add(new GrantValidationDefinition(List.of("SELECT", "REFERENCES"), List.of(SnowflakeObjectType.SEMANTIC_VIEW)));
+        add(new GrantValidationDefinition(List.of("SELECT", "REFERENCES", "MONITOR"), List.of(SnowflakeObjectType.SEMANTIC_VIEW)));
         add(new GrantValidationDefinition(List.of("USAGE", "MODIFY", "MONITOR"), List.of(SnowflakeObjectType.CORTEX_AGENT)));
     }});
 

src/main/java/us/zoom/data/dfence/providers/snowflake/grant/desired/create/internal/AllGrantsCompiler.java

@@ -50,7 +50,7 @@ private List<SnowflakeGrantModel> expandAllGrants(
         results.add(
             new SnowflakeGrantModel(
                 privilege,
-                objectType.getObjectType().replace(" ", "_"),
+                objectType.name(),
                 objectName,
                 "ROLE",
                 roleName,

src/main/java/us/zoom/data/dfence/providers/snowflake/grant/desired/create/internal/FutureGrantsCompiler.java

@@ -67,15 +67,13 @@ private List<SnowflakeGrantModel> createFutureGrantsOnContainer(
       List<String> privileges,
       String roleName,
       Boolean grantOption) {
-    String objectName =
-        String.format(
-            "%s.<%s>", containerName, objectType.getObjectType().replace(" ", "_").toUpperCase());
+    String objectName = String.format("%s.<%s>", containerName, objectType.name());
     return privileges.stream()
         .map(
             p ->
                 new SnowflakeGrantModel(
                     p,
-                    objectType.getObjectType().replace(" ", "_"),
+                    objectType.name(),
                     objectName,
                     "ROLE",
                     roleName,

src/test/java/us/zoom/data/dfence/providers/snowflake/grant/builder/SnowflakeGrantBuilderTest.java

@@ -13,10 +13,7 @@
 import java.util.List;
 import java.util.stream.Stream;
 
-import static org.junit.jupiter.api.Assertions.assertEquals;
-import static org.junit.jupiter.api.Assertions.assertInstanceOf;
-import static org.junit.jupiter.api.Assertions.assertNull;
-import static org.junit.jupiter.api.Assertions.assertThrows;
+import static org.junit.jupiter.api.Assertions.*;
 
 /**
  * Tests for SnowflakeGrantBuilder.
@@ -142,6 +139,22 @@ void fromGrant(GrantTestDataLoader.FixtureGrantTestData testData) throws RbacDat
                 "Expected builder class mismatch for: " + testData.name());
     }
 
+    @Test
+    void agentGrantBuilderFromSnowflakeGrantModel() {
+        SnowflakeGrantModel grantModel = new SnowflakeGrantModel(
+                "USAGE",
+                "AGENT",
+                "MOCK_DB.MOCK_SCHEMA.<CORTEX_AGENT>",
+                "ROLE",
+                "MOCK_ROLE",
+                false,
+                true,
+                false
+        );
+        SnowflakeGrantBuilder builder = SnowflakeGrantBuilder.fromGrant(grantModel);
+        assertNotNull(builder);
+    }
+
     /**
      * Tests conversion from SnowflakeGrantModel to PlaybookPrivilegeGrant.
      * Test cases are defined in playbook-privilege-grants.yml

src/test/java/us/zoom/data/dfence/providers/snowflake/revoke/SnowflakeRevokeGrantsCompilerTest.java

@@ -366,38 +366,6 @@ void compileRevokeGrants_whenInvalidObjectTypeInPlaybook_shouldThrow() {
                 List.of(playbookGrant), currentGrants));
   }
 
-  @Test
-  void compileRevokeGrants_whenInvalidObjectTypeInGrant_cannotCreateBuilder() {
-    // Given
-    PlaybookPrivilegeGrant playbookGrant =
-        createPlaybookGrant("TABLE", "TEST_DB", "TEST_SCHEMA", "TEST_TABLE", List.of("SELECT"));
-    SnowflakeGrantModel invalidGrant =
-        new SnowflakeGrantModel(
-            "SELECT",
-            "INVALID_TYPE",
-            "TEST_DB.TEST_SCHEMA.TEST_TABLE",
-            "ROLE",
-            "TEST_ROLE",
-            false,
-            false,
-            false);
-    SnowflakeGrantModel validGrant =
-        createGrant("UPDATE", "TABLE", "TEST_DB.TEST_SCHEMA.TEST_TABLE", "ROLE", "TEST_ROLE");
-    Map<String, SnowflakeGrantBuilder> currentGrants =
-        createCurrentGrants(invalidGrant, validGrant);
-
-    // When
-    List<SnowflakeGrantBuilder> actualRevokes =
-        SnowflakeRevokeGrantsCompiler.compileRevokeGrants(List.of(playbookGrant), currentGrants);
-
-    // Then
-    assertEquals(1, currentGrants.values().size());
-    assertEquals(1, actualRevokes.size());
-    SnowflakeGrantModel revokedGrant = actualRevokes.get(0).getGrant();
-    assertEquals("UPDATE", revokedGrant.privilege());
-    assertEquals("TABLE", revokedGrant.grantedOn());
-  }
-
   @Test
   void compileRevokeGrants_whenDatabaseAndSchemaWildcards_shouldNotRevoke() {
     // Given
@@ -656,14 +624,10 @@ private SnowflakeGrantModel createAllGrant(
   private Map<String, SnowflakeGrantBuilder> createCurrentGrants(SnowflakeGrantModel... grants) {
     Map<String, SnowflakeGrantBuilder> currentGrants = new HashMap<>();
     for (SnowflakeGrantModel grant : grants) {
-      try {
         SnowflakeGrantBuilder builder = SnowflakeGrantBuilder.fromGrant(grant);
         if (builder != null) {
           currentGrants.put(builder.getKey(), builder);
         }
-      } catch (Exception e) {
-        // Skip grants that cannot be built
-      }
     }
     return currentGrants;
   }
@@ -813,6 +777,23 @@ void compileRevokeGrants_whenPlaybookHasAgentWithWildcard_shouldNotRevoke() {
         "Should not revoke when playbook has AGENT with wildcard object name");
   }
 
+  @Test
+  void compileRevokeGrantsFuture_whenPlaybookHasAgentWithWildcard_shouldNotRevoke() {
+    // Critical: Wildcard object name should match any agent name
+    PlaybookPrivilegeGrant playbookGrant =
+            createPlaybookGrant("AGENT", "TEST_DB", "TEST_SCHEMA", "*", List.of("USAGE"));
+    SnowflakeGrantModel currentGrant =
+            new SnowflakeGrantModel("USAGE", "CORTEX_AGENT", "TEST_DB.TEST_SCHEMA.<CORTEX_AGENT>", "ROLE", "TEST_ROLE", false, true, false);
+    Map<String, SnowflakeGrantBuilder> currentGrants = createCurrentGrants(currentGrant);
+
+    List<SnowflakeGrantBuilder> actualRevokes =
+            SnowflakeRevokeGrantsCompiler.compileRevokeGrants(List.of(playbookGrant), currentGrants);
+
+    assertTrue(
+            actualRevokes.isEmpty(),
+            "Should not revoke when playbook has AGENT with wildcard object name");
+  }
+
   @Test
   void compileRevokeGrants_whenPlaybookHasAgentCaseInsensitive_shouldNotRevoke() {
     // Critical: Case should not matter for AGENT override

src/test/java/us/zoom/data/systemtest/TestRevokeGrantsSystemTest.java

@@ -127,7 +127,7 @@ public void testApplyInitialGrants() throws SQLException {
 
         // Ensure that there are no more changes.
         ChangesSummary newChanges = playbookService.compileChanges(false);
-        Assert.assertTrue(newChanges.changes().isEmpty(), "Expected no remaining changes after apply");
+        Assert.assertTrue(newChanges.changes().isEmpty(), String.format("Expected no remaining changes after apply. Changes: %s", newChanges.changes()));
 
         // Validate that the role has the expected grants from playbook.
         List<Grant> grantsExpected = new ArrayList<>(List.of(
@@ -145,10 +145,6 @@ public void testCreateExtraGrants() throws SQLException {
         // Create extra grants NOT covered by the playbook
         try (Connection connection = securityadminSnowflakeConnectionProvider.getConnection()) {
             // Grant MONITOR on database (not in playbook)
-            connection.createStatement().execute(
-                    String.format("GRANT MONITOR ON DATABASE %s TO ROLE %s",
-                            fixture.databaseName(), fixture.roleName())
-            );
             // Grant USAGE on schema (not in playbook)
             connection.createStatement().execute(
                     String.format("GRANT USAGE ON SCHEMA \"%s\".\"%s\" TO ROLE %s",

src/test/resources/test-data/grant-builder/grant-revoke-statements.yml

@@ -169,6 +169,15 @@ tests:
     expectedGrantStatement: "GRANT SELECT ON SEMANTIC VIEW \"MOCK_DB\".\"MOCK_SCHEMA\".\"MOCK_VIEW\" TO ROLE MOCK_ROLE_1;"
     expectedRevokeStatement: "REVOKE SELECT ON SEMANTIC VIEW \"MOCK_DB\".\"MOCK_SCHEMA\".\"MOCK_VIEW\" FROM ROLE MOCK_ROLE_1;"
 
+  - name: "MONITOR on SEMANTIC_VIEW"
+    grant:
+      privilege: MONITOR
+      objectType: SEMANTIC_VIEW
+      objectName: MOCK_DB.MOCK_SCHEMA.MOCK_VIEW
+      role: MOCK_ROLE_1
+    expectedGrantStatement: "GRANT MONITOR ON SEMANTIC VIEW \"MOCK_DB\".\"MOCK_SCHEMA\".\"MOCK_VIEW\" TO ROLE MOCK_ROLE_1;"
+    expectedRevokeStatement: "REVOKE MONITOR ON SEMANTIC VIEW \"MOCK_DB\".\"MOCK_SCHEMA\".\"MOCK_VIEW\" FROM ROLE MOCK_ROLE_1;"
+
   - name: "REFERENCES on SEMANTIC_VIEW"
     grant:
       privilege: REFERENCES

src/test/resources/test-data/system-test/create-grants/roles.yml

@@ -39,10 +39,6 @@ roles:
   full_patterns_role:
     name: ${var.role2}
     grants:
-      # Qual 0: account (global object)
-      - object-type: account
-        privileges:
-          - monitor
       # Qual 1: normal/account-level role, database
       - object-type: role
         object-name: ${var.role}