From 19de71e2fc31c78ce4fc82de99a97e743e1ac9c2 Mon Sep 17 00:00:00 2001
From: TJ Murphy <teej.murphy@gmail.com>
Date: Wed, 13 Nov 2024 10:28:05 -0800
Subject: [PATCH] [BUGFIX] share ownership (#151)

* Support share with custom owner

---------

Co-authored-by: TJ Murphy <1796+teej@users.noreply.github.com>
---
 tests/integration/test_blueprint.py | 16 ++++++++++++++++
 tests/test_blueprint_ownership.py   |  1 -
 titan/blueprint.py                  |  7 +++++++
 tools/test_account_configs/base.yml |  8 ++++++++
 4 files changed, 31 insertions(+), 1 deletion(-)

diff --git a/tests/integration/test_blueprint.py b/tests/integration/test_blueprint.py
index ad2dba9a..1b9a74f6 100644
--- a/tests/integration/test_blueprint.py
+++ b/tests/integration/test_blueprint.py
@@ -591,3 +591,19 @@ def test_blueprint_split_role_user(cursor):
         cursor.execute("DROP USER IF EXISTS SPLIT_ROLE_USER")
         cursor.execute("DROP ROLE IF EXISTS SPLIT_ROLE_A")
         cursor.execute("DROP ROLE IF EXISTS SPLIT_ROLE_B")
+
+
+def test_blueprint_share_custom_owner(cursor, suffix):
+    session = cursor.connection
+    share_name = f"TEST_SHARE_CUSTOM_OWNER_{suffix}"
+    share = res.Share(name=share_name, owner="TITAN_SHARE_ADMIN")
+
+    try:
+        blueprint = Blueprint(resources=[share])
+        plan = blueprint.plan(session)
+        assert len(plan) == 1
+        assert isinstance(plan[0], CreateResource)
+        assert plan[0].urn.fqn.name == share_name
+        blueprint.apply(session, plan)
+    finally:
+        cursor.execute(f"DROP SHARE IF EXISTS {share_name}")
diff --git a/tests/test_blueprint_ownership.py b/tests/test_blueprint_ownership.py
index c279ec64..11024e11 100644
--- a/tests/test_blueprint_ownership.py
+++ b/tests/test_blueprint_ownership.py
@@ -11,7 +11,6 @@
 )
 from titan.enums import AccountEdition
 from titan.identifiers import parse_URN
-from titan.privs import AccountPriv, GrantedPrivilege
 from titan.resource_name import ResourceName
 
 
diff --git a/titan/blueprint.py b/titan/blueprint.py
index 02b617f1..4349dc41 100644
--- a/titan/blueprint.py
+++ b/titan/blueprint.py
@@ -1052,6 +1052,13 @@ def execution_strategy_for_change(
     elif isinstance(change, CreateResource):
         if isinstance(change.resource_cls.scope, AccountScope):
             create_priv = CREATE_PRIV_FOR_RESOURCE_TYPE[change.urn.resource_type]
+
+            # SHARE ownership cannot be changed
+            if change.urn.resource_type == ResourceType.SHARE:
+                if change_owner is None:
+                    raise RuntimeError
+                return change_owner, False
+
             system_role = system_role_for_priv(create_priv)
             if system_role and system_role in available_roles:
                 transfer_ownership = system_role != change_owner
diff --git a/tools/test_account_configs/base.yml b/tools/test_account_configs/base.yml
index 6677ca4b..66224fa6 100644
--- a/tools/test_account_configs/base.yml
+++ b/tools/test_account_configs/base.yml
@@ -78,6 +78,8 @@ roles:
     comment: This role has every privilege
   - name: TITAN_GRANT_ADMIN
     comment: This role has MANAGE GRANTS privileges
+  - name: TITAN_SHARE_ADMIN
+    comment: This role has CREATE SHARE privilege
 
 databases:
   - name: static_database
@@ -103,6 +105,9 @@ role_grants:
   - role: TITAN_GRANT_ADMIN
     roles:
       - SYSADMIN
+  - role: TITAN_SHARE_ADMIN
+    roles:
+      - SYSADMIN
 
 # database_role_grants:
 #   - role: static_database_role
@@ -134,6 +139,9 @@ grants:
   # TITAN_GRANT_ADMIN grants
   - GRANT MANAGE GRANTS ON ACCOUNT TO ROLE TITAN_GRANT_ADMIN
 
+  # TITAN_SHARE_ADMIN grants
+  - GRANT CREATE SHARE ON ACCOUNT TO ROLE TITAN_SHARE_ADMIN
+
   # CI grants
   - GRANT USAGE ON WAREHOUSE STATIC_WAREHOUSE TO ROLE CI
   - GRANT USAGE ON DATABASE static_database TO ROLE CI