release/v0.11.0 (#174) Bump version

* v0.11.0 release

---------

Co-authored-by: TJ Murphy <[email protected]>

Author: teej

misc/generate_docstring_prompt.txt prompts/generate_docstring_prompt.txt

@@ -0,0 +1,5 @@
+{
+    "database_role": "STATIC_DATABASE.STATIC_DATABASE_ROLE",
+    "to_role": "STATIC_ROLE",
+    "to_database_role": null
+}

misc/generate_new_resource_prompt.txt prompts/generate_new_resource_prompt.txt

@@ -0,0 +1,11 @@
+{
+    "name": "some_masking_policy",
+    "args": [
+        {"name": "val", "data_type": "VARCHAR"}
+    ],
+    "returns": "VARCHAR(16777216)",
+    "body": "CASE WHEN current_role() IN ('ANALYST') THEN VAL ELSE '*********' END",
+    "comment": "Masks email addresses",
+    "exempt_other_policies": false,
+    "owner": "SYSADMIN"
+}

tests/fixtures/json/database_role_grant.json

@@ -744,3 +744,61 @@ def test_fetch_database_role_grant(cursor, suffix, marked_for_cleanup):
     result = clean_resource_data(res.Grant.spec, result)
     data = clean_resource_data(res.Grant.spec, grant.to_dict())
     assert result == data
+
+
+def test_fetch_database_role(cursor, suffix, marked_for_cleanup):
+    role = res.DatabaseRole(
+        name=f"TEST_FETCH_DATABASE_ROLE_{suffix}",
+        database="STATIC_DATABASE",
+        owner=TEST_ROLE,
+    )
+    create(cursor, role)
+    marked_for_cleanup.append(role)
+
+    result = safe_fetch(cursor, role.urn)
+    assert result is not None
+    result = clean_resource_data(res.DatabaseRole.spec, result)
+    data = clean_resource_data(res.DatabaseRole.spec, role.to_dict())
+    assert result == data
+
+
+def test_fetch_grant_of_database_role(cursor, suffix, marked_for_cleanup):
+    db_role = res.DatabaseRole(
+        name=f"TEST_FETCH_GRANT_OF_DATABASE_ROLE_{suffix}",
+        database="STATIC_DATABASE",
+        owner=TEST_ROLE,
+    )
+    create(cursor, db_role)
+    marked_for_cleanup.append(db_role)
+
+    role = res.Role(name=f"TEST_FETCH_GRANT_OF_DATABASE_ROLE_{suffix}", owner=TEST_ROLE)
+    create(cursor, role)
+    marked_for_cleanup.append(role)
+
+    grant = res.DatabaseRoleGrant(database_role=db_role, to_role=role)
+    create(cursor, grant)
+
+    result = safe_fetch(cursor, grant.urn)
+    assert result is not None
+    result = clean_resource_data(res.DatabaseRoleGrant.spec, result)
+    data = clean_resource_data(res.DatabaseRoleGrant.spec, grant.to_dict())
+    assert result == data
+
+
+def test_fetch_masking_policy(cursor, suffix, marked_for_cleanup):
+    policy = res.MaskingPolicy(
+        name=f"TEST_FETCH_MASKING_POLICY_{suffix}",
+        args=[{"name": "val", "data_type": "STRING"}],
+        returns="STRING",
+        body="CASE WHEN current_role() IN ('ANALYST') THEN VAL ELSE '*********' END",
+        comment="Masks email addresses",
+        owner=TEST_ROLE,
+    )
+    create(cursor, policy)
+    marked_for_cleanup.append(policy)
+
+    result = safe_fetch(cursor, policy.urn)
+    assert result is not None
+    result = clean_resource_data(res.MaskingPolicy.spec, result)
+    data = clean_resource_data(res.MaskingPolicy.spec, policy.to_dict())
+    assert result == data

tests/fixtures/json/masking_policy.json

@@ -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,90 @@ 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)
+
+
+def test_grant_database_role_to_database_role(cursor, suffix, marked_for_cleanup):
+    session = cursor.connection
+    bp = Blueprint()
+
+    parent = res.DatabaseRole(name=f"DBR2DBR_PARENT_{suffix}", database="STATIC_DATABASE")
+    child1 = res.DatabaseRole(name=f"DBR2DBR_CHILD_1_{suffix}", database="STATIC_DATABASE")
+    child2 = res.DatabaseRole(name=f"DBR2DBR_CHILD_2_{suffix}", database="STATIC_DATABASE")
+    drg1 = res.DatabaseRoleGrant(database_role=child1, to_database_role=parent)
+    drg2 = res.DatabaseRoleGrant(database_role=child2, to_database_role=parent)
+
+    marked_for_cleanup.append(parent)
+    marked_for_cleanup.append(child1)
+    marked_for_cleanup.append(child2)
+
+    bp.add(parent, child1, child2, drg1, drg2)
+    plan = bp.plan(session)
+    assert len(plan) == 5
+    bp.apply(session, plan)
+
+    grant1 = safe_fetch(cursor, res.DatabaseRoleGrant(database_role=child1, to_database_role=parent).urn)
+    assert grant1 is not None
+    assert grant1["database_role"] == str(child1.fqn)
+    assert grant1["to_database_role"] == str(parent.fqn)
+
+    grant2 = safe_fetch(cursor, res.DatabaseRoleGrant(database_role=child2, to_database_role=parent).urn)
+    assert grant2 is not None
+    assert grant2["database_role"] == str(child2.fqn)
+    assert grant2["to_database_role"] == str(parent.fqn)

tests/integration/data_provider/test_fetch_resource.py

@@ -1,9 +1,22 @@
 import pytest
 
-from titan.operations.export import export_resources
+from titan.identifiers import URN, parse_FQN
+from titan.operations.export import export_resources, _format_resource_config
+from titan.enums import ResourceType
+from titan.data_provider import fetch_resource
 
 pytestmark = pytest.mark.requires_snowflake
 
 
 def test_export_all(cursor):
     assert export_resources(session=cursor.connection)
+
+
+def test_export_schema(cursor):
+    urn = URN(ResourceType.SCHEMA, parse_FQN("STATIC_DATABASE.STATIC_SCHEMA", is_db_scoped=True))
+    resource = fetch_resource(cursor, urn)
+    assert resource
+    resource_cfg = _format_resource_config(urn, resource, ResourceType.SCHEMA)
+    assert resource_cfg
+    assert "database" in resource_cfg
+    assert resource_cfg["database"] == "STATIC_DATABASE"

tests/integration/test_blueprint.py

@@ -55,6 +55,7 @@ def test_create_drop_from_json(resource, cursor, suffix):
         res.FutureGrant,
         res.Grant,
         res.RoleGrant,
+        res.DatabaseRoleGrant,
         res.ScannerPackage,
         res.Service,
     ):
@@ -64,6 +65,7 @@ def test_create_drop_from_json(resource, cursor, suffix):
     database = res.Database(name=lifecycle_db, owner="SYSADMIN")
 
     feature_enabled = True
+    drop_sql = None
 
     try:
         fetch_session.cache_clear()
@@ -213,8 +215,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_export.py

@@ -127,3 +127,16 @@ def test_fetch_warehouse_snowpark_optimized(cursor, suffix, marked_for_cleanup):
     data = safe_fetch(cursor, warehouse.urn)
     assert data is not None
     assert data["warehouse_type"] == "SNOWPARK-OPTIMIZED"
+
+
+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())

tests/integration/test_lifecycle.py

@@ -253,23 +253,23 @@ def test_grant_database_role_to_database_role():
     database = res.Database(name="somedb")
     parent = res.DatabaseRole(name="parent", database=database)
     child = res.DatabaseRole(name="child", database=database)
-    grant = res.RoleGrant(role=child, to_role=parent)
-    assert grant.role.name == "child"
+    grant = res.DatabaseRoleGrant(database_role=child, to_database_role=parent)
+    assert grant.database_role.name == "child"
     assert grant.to.name == "parent"
 
 
 def test_grant_database_role_to_account_role():
     database = res.Database(name="somedb")
     parent = res.Role(name="parent")
     child = res.DatabaseRole(name="child", database=database)
-    grant = res.RoleGrant(role=child, to_role=parent)
-    assert grant.role.name == "child"
+    grant = res.DatabaseRoleGrant(database_role=child, to_role=parent)
+    assert grant.database_role.name == "child"
     assert grant.to.name == "parent"
 
 
 def test_grant_database_role_to_system_role():
     database = res.Database(name="somedb")
     child = res.DatabaseRole(name="child", database=database)
-    grant = res.RoleGrant(role=child, to_role="SYSADMIN")
-    assert grant.role.name == "child"
+    grant = res.DatabaseRoleGrant(database_role=child, to_role="SYSADMIN")
+    assert grant.database_role.name == "child"
     assert grant.to.name == "SYSADMIN"

tests/integration/test_resources.py

@@ -107,6 +107,7 @@ def test_data_identity(resource):
     assert serialized == data
 
 
+@pytest.mark.skip(reason="SQL parsing will be deprecated")
 def test_sql_identity(resource: tuple[type[Resource], dict]):
     resource_cls, data = resource
     if resource_cls.__name__ == "ScannerPackage":

tests/test_grant.py

@@ -344,3 +344,10 @@ def test_user_type_fallback(caplog):
     user = res.User(name="test_user", user_type="SERVICE")
     assert "The 'user_type' parameter is deprecated. Use 'type' instead." in caplog.text
     assert user._data.type == UserType.SERVICE
+
+
+def test_future_grant_alt_syntax():
+    db = res.Database(name="DB")
+    role = res.Role(name="ROLE")
+    fg = res.FutureGrant(priv="SELECT", on_type="table", in_type=db.resource_type, in_name=db.name, to=role)
+    assert fg

tests/test_identities.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)}