Verify PR #317 (SNOW-460) grants in Snowflake

verify_pr317_grants.py

"""
Verify that all grants introduced in PR #317 (SNOW-460) are applied.

Checks:
  1. Roles exist (new roles added in roles.sql)
  2. Schemas exist (schema init migrations)
  3. Schema ownership (ownership_grants migrations)
  4. Database-level privileges (USAGE, MONITOR on SAGE)
  5. Schema-level privileges (USAGE, MONITOR on SAGE.<schema>)
  6. Role-to-role grants (analyst roles → DATA_ENGINEER, admin roles → SAGE_ADMIN, etc.)
  7. Warehouse grants (USAGE on STREAMLIT_XSMALL)
  8. User-to-role grants
  9. Future grants (USAGE ON FUTURE STREAMLITS; full AD schema future grants)

Usage:
    venv/snowflake/bin/python verify_pr317_grants.py
"""

import sys
import snowflake.connector


PASS = "PASS"
FAIL = "FAIL"

results: list[tuple[str, str, str]] = []  # (check, expected, status)


def check(label: str, condition: bool, detail: str = "") -> None:
    status = PASS if condition else FAIL
    results.append((label, detail, status))


# ---------------------------------------------------------------------------
# Connection
# ---------------------------------------------------------------------------

print("Connecting to Snowflake...")
conn = snowflake.connector.connect()
cur = conn.cursor()

# Use ACCOUNTADMIN so we can see all grants.
cur.execute("USE ROLE ACCOUNTADMIN")
print("Connected.\n")


# ---------------------------------------------------------------------------
# Helpers
# ---------------------------------------------------------------------------

def fetch_set(sql: str) -> set[str]:
    cur.execute(sql)
    return {str(row[0]).upper() for row in cur.fetchall()}


def show_grants_to_role(role: str) -> list[dict]:
    """Return list of dicts with keys privilege, granted_on, name."""
    cur.execute(f"SHOW GRANTS TO ROLE {role}")
    cols = [d[0].lower() for d in cur.description]
    rows = cur.fetchall()
    return [dict(zip(cols, row)) for row in rows]


def role_has_privilege_on(role: str, privilege: str, object_type: str, object_name: str) -> bool:
    grants = show_grants_to_role(role)
    priv_up = privilege.upper()
    type_up = object_type.upper()
    name_up = object_name.upper()
    return any(
        str(g.get("privilege", "")).upper() == priv_up
        and str(g.get("granted_on", "")).upper() == type_up
        and str(g.get("name", "")).upper() == name_up
        for g in grants
    )


def role_inherits_role(grantee_role: str, granted_role: str) -> bool:
    """Check that granted_role is granted TO grantee_role."""
    cur.execute(f"SHOW GRANTS TO ROLE {grantee_role}")
    cols = [d[0].lower() for d in cur.description]
    rows = cur.fetchall()
    for row in rows:
        g = dict(zip(cols, row))
        if (
            str(g.get("privilege", "")).upper() == "USAGE"
            and str(g.get("granted_on", "")).upper() == "ROLE"
            and str(g.get("name", "")).upper() == granted_role.upper()
        ):
            return True
    return False


def role_granted_to_role(granted_role: str, grantee_role: str) -> bool:
    """Check that granted_role is granted to grantee_role via SHOW GRANTS OF ROLE."""
    cur.execute(f"SHOW GRANTS OF ROLE {granted_role}")
    cols = [d[0].lower() for d in cur.description]
    rows = cur.fetchall()
    for row in rows:
        g = dict(zip(cols, row))
        if (
            str(g.get("granted_to", "")).upper() == "ROLE"
            and str(g.get("grantee_name", "")).upper() == grantee_role.upper()
        ):
            return True
    return False


def role_granted_to_user(role: str, user: str) -> bool:
    cur.execute(f"SHOW GRANTS OF ROLE {role}")
    cols = [d[0].lower() for d in cur.description]
    rows = cur.fetchall()
    for row in rows:
        g = dict(zip(cols, row))
        if (
            str(g.get("granted_to", "")).upper() == "USER"
            and str(g.get("grantee_name", "")).upper() == user.upper()
        ):
            return True
    return False


def schema_owner(database: str, schema: str) -> str:
    cur.execute(f"SHOW SCHEMAS LIKE '{schema}' IN DATABASE {database}")
    cols = [d[0].lower() for d in cur.description]
    rows = cur.fetchall()
    for row in rows:
        g = dict(zip(cols, row))
        if str(g.get("name", "")).upper() == schema.upper():
            return str(g.get("owner", "")).upper()
    return ""


def future_grant_exists(schema_db: str, schema: str, object_type: str, privilege: str, grantee: str) -> bool:
    """Check SHOW FUTURE GRANTS IN SCHEMA for a specific grant."""
    cur.execute(f"SHOW FUTURE GRANTS IN SCHEMA {schema_db}.{schema}")
    cols = [d[0].lower() for d in cur.description]
    rows = cur.fetchall()
    for row in rows:
        g = dict(zip(cols, row))
        if (
            str(g.get("privilege", "")).upper() == privilege.upper()
            and str(g.get("grant_on", "")).upper() == object_type.upper()
            and str(g.get("grantee_name", "")).upper() == grantee.upper()
        ):
            return True
    return False


# ---------------------------------------------------------------------------
# 1. Roles exist
# ---------------------------------------------------------------------------

print("=" * 60)
print("1. ROLES EXIST")
print("=" * 60)

cur.execute("SHOW ROLES")
cols = [d[0].lower() for d in cur.description]
existing_roles = {
    str(dict(zip(cols, row)).get("name", "")).upper()
    for row in cur.fetchall()
}

new_roles = [
    "SAGE_NF_ADMIN", "SAGE_NF_ANALYST",
    "SAGE_GENIE_ADMIN", "SAGE_GENIE_ANALYST",
    "SAGE_ARK_ADMIN", "SAGE_ARK_ANALYST",
    "SAGE_AMP_ALS_ADMIN", "SAGE_AMP_ALS_ANALYST",
    "SAGE_PRODUCT_ADMIN", "SAGE_PRODUCT_ANALYST",
    "SAGE_CCKP_ADMIN", "SAGE_CCKP_ANALYST",
    "SAGE_AD_ADMIN", "SAGE_AD_ANALYST",
    "SAGE_ELITE_ADMIN", "SAGE_ELITE_ANALYST",
]

for role in new_roles:
    exists = role in existing_roles
    check(f"Role exists: {role}", exists, role)

# NF_ADMIN should have been dropped
nf_admin_dropped = "NF_ADMIN" not in existing_roles
check("Role dropped: NF_ADMIN", nf_admin_dropped, "NF_ADMIN")


# ---------------------------------------------------------------------------
# 2. Schemas exist
# ---------------------------------------------------------------------------

print("\n" + "=" * 60)
print("2. SCHEMAS EXIST")
print("=" * 60)

cur.execute("SHOW SCHEMAS IN DATABASE SAGE")
cols = [d[0].lower() for d in cur.description]
existing_schemas = {
    str(dict(zip(cols, row)).get("name", "")).upper()
    for row in cur.fetchall()
}

new_schemas = ["NF", "GENIE", "ARK", "AMP_ALS", "PRODUCT", "CCKP", "ELITE", "AD"]
for schema in new_schemas:
    exists = schema in existing_schemas
    check(f"Schema exists: SAGE.{schema}", exists, schema)


# ---------------------------------------------------------------------------
# 3. Schema ownership
# ---------------------------------------------------------------------------

print("\n" + "=" * 60)
print("3. SCHEMA OWNERSHIP")
print("=" * 60)

schema_owners = {
    "NF": "SAGE_NF_ADMIN",
    "GENIE": "SAGE_GENIE_ADMIN",
    "ARK": "SAGE_ARK_ADMIN",
    "AMP_ALS": "SAGE_AMP_ALS_ADMIN",
    "PRODUCT": "SAGE_PRODUCT_ADMIN",
    "CCKP": "SAGE_CCKP_ADMIN",
    "AD": "SAGE_AD_ADMIN",
    "ELITE": "SAGE_ELITE_ADMIN",
}

for schema, expected_owner in schema_owners.items():
    actual_owner = schema_owner("SAGE", schema)
    ok = actual_owner == expected_owner.upper()
    check(
        f"Schema owner: SAGE.{schema} → {expected_owner}",
        ok,
        f"actual={actual_owner}",
    )


# ---------------------------------------------------------------------------
# 4. Database privileges: USAGE, MONITOR on SAGE
# ---------------------------------------------------------------------------

print("\n" + "=" * 60)
print("4. DATABASE PRIVILEGES (USAGE, MONITOR on SAGE)")
print("=" * 60)

db_priv_roles = [
    "SAGE_NF_ADMIN", "SAGE_NF_ANALYST",
    "SAGE_GENIE_ADMIN", "SAGE_GENIE_ANALYST",
    "SAGE_ARK_ADMIN", "SAGE_ARK_ANALYST",
    "SAGE_AMP_ALS_ADMIN", "SAGE_AMP_ALS_ANALYST",
    "SAGE_PRODUCT_ADMIN", "SAGE_PRODUCT_ANALYST",
    "SAGE_CCKP_ADMIN", "SAGE_CCKP_ANALYST",
    "SAGE_AD_ADMIN", "SAGE_AD_ANALYST",
    "SAGE_ELITE_ADMIN", "SAGE_ELITE_ANALYST",
]

for role in db_priv_roles:
    for priv in ("USAGE", "MONITOR"):
        ok = role_has_privilege_on(role, priv, "DATABASE", "SAGE")
        check(f"DB priv: GRANT {priv} ON DATABASE SAGE TO ROLE {role}", ok, "")


# ---------------------------------------------------------------------------
# 5. Schema privileges: USAGE, MONITOR on SAGE.<schema>
# ---------------------------------------------------------------------------

print("\n" + "=" * 60)
print("5. SCHEMA PRIVILEGES (USAGE, MONITOR on SAGE.<schema>)")
print("=" * 60)

analyst_schema_map = {
    "SAGE_NF_ANALYST": "SAGE.NF",
    "SAGE_GENIE_ANALYST": "SAGE.GENIE",
    "SAGE_ARK_ANALYST": "SAGE.ARK",
    "SAGE_AMP_ALS_ANALYST": "SAGE.AMP_ALS",
    "SAGE_PRODUCT_ANALYST": "SAGE.PRODUCT",
    "SAGE_CCKP_ANALYST": "SAGE.CCKP",
    "SAGE_AD_ANALYST": "SAGE.AD",
    "SAGE_ELITE_ANALYST": "SAGE.ELITE",
}

for role, schema_fqn in analyst_schema_map.items():
    for priv in ("USAGE", "MONITOR"):
        ok = role_has_privilege_on(role, priv, "SCHEMA", schema_fqn)
        check(f"Schema priv: GRANT {priv} ON SCHEMA {schema_fqn} TO ROLE {role}", ok, "")


# ---------------------------------------------------------------------------
# 6. Role-to-role grants
# ---------------------------------------------------------------------------

print("\n" + "=" * 60)
print("6. ROLE-TO-ROLE GRANTS")
print("=" * 60)

role_to_role_grants = [
    # (granted_role, grantee_role)
    ("SAGE_NF_ANALYST", "DATA_ENGINEER"),
    ("SAGE_GENIE_ANALYST", "DATA_ENGINEER"),
    ("SAGE_ARK_ANALYST", "DATA_ENGINEER"),
    ("SAGE_AMP_ALS_ANALYST", "DATA_ENGINEER"),
    ("SAGE_PRODUCT_ANALYST", "DATA_ENGINEER"),
    ("SAGE_PRODUCT_ANALYST", "TECH_PRODUCT"),
    ("SAGE_CCKP_ANALYST", "DATA_ENGINEER"),
    ("SAGE_AD_ANALYST", "DATA_ENGINEER"),
    ("SAGE_AD_ANALYST", "AD"),
    ("SAGE_AD_ADMIN", "AD"),
    ("SAGE_ELITE_ANALYST", "DATA_ENGINEER"),
    # Admin roles → SAGE_ADMIN
    ("SAGE_NF_ADMIN", "SAGE_ADMIN"),
    ("SAGE_GENIE_ADMIN", "SAGE_ADMIN"),
    ("SAGE_ARK_ADMIN", "SAGE_ADMIN"),
    ("SAGE_AMP_ALS_ADMIN", "SAGE_ADMIN"),
    ("SAGE_PRODUCT_ADMIN", "SAGE_ADMIN"),
    ("SAGE_CCKP_ADMIN", "SAGE_ADMIN"),
    ("SAGE_AD_ADMIN", "SAGE_ADMIN"),
    ("SAGE_ELITE_ADMIN", "SAGE_ADMIN"),
    # NF_ADMIN (new) grants from roles.sql
    ("SAGE_NF_ADMIN", "SYSADMIN"),
    ("DATA_ANALYTICS", "SAGE_NF_ADMIN"),  # GRANT ROLE DATA_ANALYTICS TO ROLE SAGE_NF_ADMIN
]

for granted, grantee in role_to_role_grants:
    ok = role_granted_to_role(granted, grantee)
    check(f"Role grant: {granted} → {grantee}", ok, "")


# ---------------------------------------------------------------------------
# 7. Warehouse grants: USAGE on STREAMLIT_XSMALL
# ---------------------------------------------------------------------------

print("\n" + "=" * 60)
print("7. WAREHOUSE GRANTS (USAGE on STREAMLIT_XSMALL)")
print("=" * 60)

warehouse_roles = [
    "SAGE_ARK_ANALYST", "SAGE_ARK_ADMIN",
    "SAGE_NF_ANALYST", "SAGE_NF_ADMIN",
    "SAGE_AMP_ALS_ANALYST", "SAGE_AMP_ALS_ADMIN",
    "SAGE_GENIE_ANALYST", "SAGE_GENIE_ADMIN",
    "SAGE_PRODUCT_ANALYST", "SAGE_PRODUCT_ADMIN",
    "SAGE_CCKP_ANALYST", "SAGE_CCKP_ADMIN",
    "SAGE_AD_ANALYST", "SAGE_AD_ADMIN",
    "SAGE_ELITE_ANALYST", "SAGE_ELITE_ADMIN",
]

for role in warehouse_roles:
    ok = role_has_privilege_on(role, "USAGE", "WAREHOUSE", "STREAMLIT_XSMALL")
    check(f"WH grant: USAGE ON STREAMLIT_XSMALL TO ROLE {role}", ok, "")


# ---------------------------------------------------------------------------
# 8. User-to-role grants
# ---------------------------------------------------------------------------

print("\n" + "=" * 60)
print("8. USER-TO-ROLE GRANTS")
print("=" * 60)

user_role_grants = [
    # role, user (login name / username as stored in Snowflake)
    ("SAGE_ARK_ANALYST", "jessica.vera@sagebase.org"),
    ("SAGE_ARK_ANALYST", "bishoy.kamel@sagebase.org"),
    ("SAGE_NF_ANALYST", "james.moon@sagebase.org"),
    ("SAGE_NF_ANALYST", "robert.allaway@sagebase.org"),
    ("SAGE_NF_ANALYST", "belinda.garana@sagebase.org"),
    ("SAGE_AMP_ALS_ANALYST", "jessica.vera@sagebase.org"),
    ("SAGE_AMP_ALS_ANALYST", "ram.ayyala@sagebase.org"),
    ("SAGE_AMP_ALS_ANALYST", "vanessa.barone@sagebase.org"),
    ("SAGE_GENIE_ANALYST", "adam.taylor@sagebase.org"),
    ("SAGE_GENIE_ANALYST", "ashley.clayton@sagebase.org"),
    ("SAGE_ELITE_ANALYST", "milan.vu@sagebase.org"),
    ("SAGE_ELITE_ANALYST", "melissa.klein@sagebase.org"),
    # SAGE_NF_ADMIN user grants (added in this PR)
    ("SAGE_NF_ADMIN", "anh.nguyet.vu@sagebase.org"),
    ("SAGE_NF_ADMIN", "thomas.yu@sagebase.org"),
]

for role, user in user_role_grants:
    ok = role_granted_to_user(role, user)
    check(f"User grant: {role} → {user}", ok, "")


# ---------------------------------------------------------------------------
# 9. Future grants
# ---------------------------------------------------------------------------

print("\n" + "=" * 60)
print("9. FUTURE GRANTS")
print("=" * 60)

# Each new schema gets USAGE ON FUTURE STREAMLITS to analyst role
streamlit_future_grants = {
    "NF": "SAGE_NF_ANALYST",
    "GENIE": "SAGE_GENIE_ANALYST",
    "ARK": "SAGE_ARK_ANALYST",
    "AMP_ALS": "SAGE_AMP_ALS_ANALYST",
    "PRODUCT": "SAGE_PRODUCT_ANALYST",
    "CCKP": "SAGE_CCKP_ANALYST",
    "ELITE": "SAGE_ELITE_ANALYST",
    "AD": "SAGE_AD_ANALYST",
}

for schema, analyst_role in streamlit_future_grants.items():
    ok = future_grant_exists("SAGE", schema, "STREAMLIT", "USAGE", analyst_role)
    check(f"Future grant: USAGE ON FUTURE STREAMLITS IN SAGE.{schema} TO {analyst_role}", ok, "")

# AD: full set of future analyst grants (V1.32.14)
ad_analyst_future = [
    ("TABLE", "SELECT"),
    ("TABLE", "REFERENCES"),
    ("VIEW", "SELECT"),
    ("VIEW", "REFERENCES"),
    ("STAGE", "USAGE"),
    ("STAGE", "READ"),
    ("FILE_FORMAT", "USAGE"),
    ("FUNCTION", "USAGE"),
    ("PROCEDURE", "USAGE"),
]

for obj_type, priv in ad_analyst_future:
    ok = future_grant_exists("SAGE", "AD", obj_type, priv, "SAGE_AD_ANALYST")
    check(f"Future grant (AD analyst): {priv} ON FUTURE {obj_type}S IN SAGE.AD TO SAGE_AD_ANALYST", ok, "")

# AD: future ownership grants for admin role (V1.32.14)
ad_admin_future_objects = ["TABLE", "VIEW", "STAGE", "FILE_FORMAT", "FUNCTION", "PROCEDURE"]
for obj_type in ad_admin_future_objects:
    ok = future_grant_exists("SAGE", "AD", obj_type, "OWNERSHIP", "SAGE_AD_ADMIN")
    check(f"Future grant (AD admin): OWNERSHIP ON FUTURE {obj_type}S IN SAGE.AD TO SAGE_AD_ADMIN", ok, "")


# ---------------------------------------------------------------------------
# Summary
# ---------------------------------------------------------------------------

conn.close()

passes = [r for r in results if r[2] == PASS]
failures = [r for r in results if r[2] == FAIL]

print("\n" + "=" * 60)
print("SUMMARY")
print("=" * 60)

if failures:
    print(f"\nFAILED ({len(failures)}):")
    for label, detail, _ in failures:
        suffix = f"  [{detail}]" if detail else ""
        print(f"  FAIL  {label}{suffix}")

print(f"\nPassed: {len(passes)} / {len(results)}")
print(f"Failed: {len(failures)} / {len(results)}")

if failures:
    sys.exit(1)
else:
    print("\nAll checks passed.")