stage priv execution order

tests/integration/test_blueprint.py

@@ -16,6 +16,7 @@
 )
 from titan.client import reset_cache
 from titan.enums import BlueprintScope, ResourceType
+from titan.exceptions import NotADAGException
 from titan.gitops import collect_blueprint_config
 from titan.resources.database import public_schema_urn
 
@@ -607,3 +608,60 @@ def test_blueprint_share_custom_owner(cursor, suffix):
         blueprint.apply(session, plan)
     finally:
         cursor.execute(f"DROP SHARE IF EXISTS {share_name}")
+
+
+def test_stage_read_write_privilege_execution_order(cursor, suffix, marked_for_cleanup):
+    session = cursor.connection
+
+    role_name = f"STAGE_ACCESS_ROLE_{suffix}"
+
+    blueprint = Blueprint()
+
+    role = res.Role(name=role_name)
+    read_grant = res.Grant(priv="READ", on_stage="STATIC_DATABASE.PUBLIC.STATIC_STAGE", to=role)
+    write_grant = res.Grant(priv="WRITE", on_stage="STATIC_DATABASE.PUBLIC.STATIC_STAGE", to=role)
+
+    # Incorrect order of execution
+    read_grant.requires(write_grant)
+
+    blueprint.add(role, read_grant, write_grant)
+
+    marked_for_cleanup.append(role)
+
+    with pytest.raises(NotADAGException):
+        blueprint.plan(session)
+
+    blueprint = Blueprint()
+
+    role = res.Role(name=role_name)
+    read_grant = res.Grant(priv="READ", on_stage="STATIC_DATABASE.PUBLIC.STATIC_STAGE", to=role)
+    write_grant = res.Grant(priv="WRITE", on_stage="STATIC_DATABASE.PUBLIC.STATIC_STAGE", to=role)
+
+    # Implicitly ordered incorrectly
+    blueprint.add(role, write_grant, read_grant)
+
+    plan = blueprint.plan(session)
+    assert len(plan) == 3
+    blueprint.apply(session, plan)
+
+    blueprint = Blueprint()
+
+    read_on_all = res.GrantOnAll(
+        priv="READ", on_type="STAGE", in_type="SCHEMA", in_name="STATIC_DATABASE.PUBLIC", to=role_name
+    )
+    future_read = res.FutureGrant(
+        priv="READ", on_type="STAGE", in_type="SCHEMA", in_name="STATIC_DATABASE.PUBLIC", to=role_name
+    )
+    write_on_all = res.GrantOnAll(
+        priv="WRITE", on_type="STAGE", in_type="SCHEMA", in_name="STATIC_DATABASE.PUBLIC", to=role_name
+    )
+    future_write = res.FutureGrant(
+        priv="WRITE", on_type="STAGE", in_type="SCHEMA", in_name="STATIC_DATABASE.PUBLIC", to=role_name
+    )
+
+    # Implicitly ordered incorrectly
+    blueprint.add(future_write, future_read, write_on_all, read_on_all)
+
+    plan = blueprint.plan(session)
+    assert len(plan) == 4
+    blueprint.apply(session, plan)

tests/integration/test_lifecycle.py

@@ -214,8 +214,8 @@ def test_task_lifecycle_remove_predecessor(cursor, suffix, marked_for_cleanup):
 
 
 def test_database_role_grants(cursor, suffix, marked_for_cleanup):
-    db = res.Database(name="whatever")
-    role = res.DatabaseRole(name="whatever_role", database=db)
+    db = res.Database(name=f"TEST_DATABASE_ROLE_GRANTS_{suffix}")
+    role = res.DatabaseRole(name=f"TEST_DATABASE_ROLE_GRANTS_{suffix}", database=db)
     grant = res.Grant(priv="USAGE", on_schema=db.public_schema.fqn, to=role)
     future_grant = res.FutureGrant(priv="SELECT", on_future_tables_in=db, to=role)
 

tests/integration/test_resources.py

@@ -139,3 +139,16 @@ def test_grant_database_role_to_database_role(cursor, suffix, marked_for_cleanup
 
     grant = res.DatabaseRoleGrant(database_role=child, to_database_role=parent)
     cursor.execute(grant.create_sql())
+
+
+def test_snowflake_builtin_database_role_grant(cursor, suffix, marked_for_cleanup):
+    drg = res.DatabaseRoleGrant(database_role="SNOWFLAKE.CORTEX_USER", to_role="STATIC_ROLE")
+    marked_for_cleanup.append(drg)
+    cursor.execute(drg.create_sql())
+
+    dbr = res.DatabaseRole(name=f"TEST_GRANT_DATABASE_ROLE_{suffix}", database="STATIC_DATABASE")
+    drg = res.DatabaseRoleGrant(database_role=dbr, to_database_role="STATIC_DATABASE.STATIC_DATABASE_ROLE")
+    marked_for_cleanup.append(dbr)
+    marked_for_cleanup.append(drg)
+    cursor.execute(dbr.create_sql())
+    cursor.execute(drg.create_sql())

titan/blueprint.py

@@ -24,6 +24,7 @@
     MissingPrivilegeException,
     MissingResourceException,
     NonConformingPlanException,
+    NotADAGException,
     OrphanResourceException,
 )
 from .identifiers import URN, parse_identifier, parse_URN, resource_label_for_type
@@ -33,7 +34,7 @@
 )
 from .resource_name import ResourceName
 from .resource_tags import ResourceTags
-from .resources import Database, RoleGrant, Schema
+from .resources import Database, FutureGrant, Grant, GrantOnAll, RoleGrant, Schema
 from .resources.database import public_schema_urn
 from .resources.resource import (
     RESOURCE_SCOPES,
@@ -870,6 +871,46 @@ def _create_grandparent_refs(self) -> None:
             if isinstance(resource.scope, SchemaScope):
                 resource.requires(resource.container.container)
 
+    def _create_stage_privilege_refs(self) -> None:
+        stage_grants: dict[str, list[Grant]] = {}
+        stage_future_grants: dict[ResourceName, list[FutureGrant]] = {}
+        stage_grant_on_all: dict[ResourceName, list[GrantOnAll]] = {}
+
+        for resource in _walk(self._root):
+            if isinstance(resource, Grant):
+                if resource._data.on_type == ResourceType.STAGE:
+                    if resource._data.on not in stage_grants:
+                        stage_grants[resource._data.on] = []
+                    stage_grants[resource._data.on].append(resource)
+            elif isinstance(resource, FutureGrant):
+                if resource._data.on_type == ResourceType.STAGE:
+                    if resource._data.in_name not in stage_future_grants:
+                        stage_future_grants[resource._data.in_name] = []
+                    stage_future_grants[resource._data.in_name].append(resource)
+            elif isinstance(resource, GrantOnAll):
+                if resource._data.on_type == ResourceType.STAGE:
+                    if resource._data.in_name not in stage_grant_on_all:
+                        stage_grant_on_all[resource._data.in_name] = []
+                    stage_grant_on_all[resource._data.in_name].append(resource)
+
+        def _apply_refs(stage_grants):
+            for stage in stage_grants.keys():
+                read_grants = []
+                write_grants = []
+                for grant in stage_grants[stage]:
+                    if grant._data.priv == "READ":
+                        read_grants.append(grant)
+                    elif grant._data.priv == "WRITE":
+                        write_grants.append(grant)
+
+                for w_grant in write_grants:
+                    for r_grant in read_grants:
+                        w_grant.requires(r_grant)
+
+        _apply_refs(stage_grants)
+        _apply_refs(stage_future_grants)
+        _apply_refs(stage_grant_on_all)
+
     def _finalize_resources(self) -> None:
         for resource in _walk(self._root):
             resource._finalized = True
@@ -883,6 +924,7 @@ def _finalize(self, session_ctx: SessionContext) -> None:
         self._create_tag_references()
         self._create_ownership_refs(session_ctx)
         self._create_grandparent_refs()
+        self._create_stage_privilege_refs()
         self._finalize_resources()
 
     def generate_manifest(self, session_ctx: SessionContext) -> Manifest:
@@ -1233,7 +1275,7 @@ def topological_sort(resource_set: set[T], references: set[tuple[T, T]]) -> dict
         outgoing_edges[node].difference_update(empty_neighbors)
     nodes.reverse()
     if len(nodes) != len(resource_set):
-        raise Exception("Graph is not a DAG")
+        raise NotADAGException("Graph is not a DAG")
     return {value: index for index, value in enumerate(nodes)}
 
 

titan/exceptions.py

@@ -48,3 +48,7 @@ class WrongEditionException(Exception):
 
 class ResourceHasContainerException(Exception):
     pass
+
+
+class NotADAGException(Exception):
+    pass