{
 "metadata": {
  "kernelspec": {
   "display_name": "Streamlit Notebook",
   "name": "streamlit"
  },
  "lastEditStatus": {
   "notebookId": "ftjxqtlw7f6leelzj5rk",
   "authorId": "432459825768",
   "authorName": "BMHESS",
   "authorEmail": "brian.hess@snowflake.com",
   "sessionId": "e4333e43-bbc3-4740-a6a6-59a97ff3859a",
   "lastEditTime": 1745848877367
  }
 },
 "nbformat_minor": 5,
 "nbformat": 4,
 "cells": [
  {
   "cell_type": "markdown",
   "id": "7a643901-62cc-4336-9209-71d344cd9650",
   "metadata": {
    "name": "cell4",
    "collapsed": false
   },
   "source": "# Using Row Access Policies based on the Current User in Streamlit in Snowflake\nThis Notebook will help set up a Row Access Policy (RAP) to restrict data access based on the current user when using Streamlit in Snowflake.\n\nThe key to this is the `READ SESSION` permission. If the owner role of the Streamlit has _not_ been granted the `READ SESSION` permission, then the session context functions (e.g., `CURRENT_USER()`, `CURRENT_AVAILABLE_ROLES()`, etc) return `NULL`. But if the owner role of the Streamlit _has_ been granted this permission, then those context functions return the proper values. \n\n:orange[Note:] `CURRENT_ROLE()` will always return the owner role of the Streamlit.\n\nIn this Notebook we:\n1. create a role to own the Streamlit and a role that will be allowed to use the Streamlit, and create 3 users and grant each the role that will be allowed to use the Streamlit\n2. create a data table with a column we will use to restrict row access (i.e., the `REGION_ID` column)\n3. Create an entitlements table that maps which users can see which region(s)\n4. create a UDF and RAP based on the current user, and apply that to the data table\n5. test data access\n6. create a Streamlit in Snowflake and grant usage to the Streamlit user role\n"
  },
  {
   "cell_type": "markdown",
   "id": "7df9bfc7-1363-40e3-bfa4-f0b2b3749c4b",
   "metadata": {
    "name": "cell12",
    "collapsed": false
   },
   "source": "## 1. Create the necessary roles and users\nFirst we create the role that will own the Streamlit. We also grant the role to the current user. This role also needs the `READ SESSION` permission.\n\nWe will also create 3 users (`USER1`, `USER2`, and `USER3`). We create another role that will get usage permission on the Streamlit, and grant that role `RL_APP_USER` to the current user and the 3 new users."
  },
  {
   "cell_type": "code",
   "id": "b338097f-1ff5-45ac-b7cb-d9f7cfb77441",
   "metadata": {
    "language": "sql",
    "name": "cell1"
   },
   "outputs": [],
   "source": "USE ROLE securityadmin;\nSET cur_user = CURRENT_USER();\n\nCREATE ROLE IF NOT EXISTS rl_app;\nGRANT ROLE rl_app TO USER IDENTIFIER($cur_user);\nGRANT READ SESSION ON ACCOUNT TO ROLE rl_app;\n\nCREATE ROLE IF NOT EXISTS rl_app_user;\nGRANT ROLE rl_app_user TO USER IDENTIFIER($cur_user);\n\nCREATE USER IF NOT EXISTS user1 PASSWORD='abc123' DEFAULT_ROLE = rl_app_user DEFAULT_SECONDARY_ROLES = ('ALL') MUST_CHANGE_PASSWORD = TRUE;\nCREATE USER IF NOT EXISTS user2 PASSWORD='abc123' DEFAULT_ROLE = rl_app_user DEFAULT_SECONDARY_ROLES = ('ALL') MUST_CHANGE_PASSWORD = TRUE;\nCREATE USER IF NOT EXISTS user2 PASSWORD='abc123' DEFAULT_ROLE = rl_app_user DEFAULT_SECONDARY_ROLES = ('ALL') MUST_CHANGE_PASSWORD = TRUE;\nGRANT ROLE rl_app_user TO USER user1;\nGRANT ROLE rl_app_user TO USER user2;\nGRANT ROLE rl_app_user TO USER user3;",
   "execution_count": null
  },
  {
   "cell_type": "markdown",
   "id": "3c5cdd13-5687-4746-b64e-302d980f94e3",
   "metadata": {
    "name": "cell6",
    "collapsed": false
   },
   "source": "## 2. Create the Data Table\nNext we create the data table that we will later add a Row Access Policy to restrict data access based on the current user. This data is accessible by the role that owns the Streamlit."
  },
  {
   "cell_type": "code",
   "id": "41726195-35e9-44f7-a1f0-96878cae1b63",
   "metadata": {
    "language": "sql",
    "name": "cell7"
   },
   "outputs": [],
   "source": "USE ROLE accountadmin;\nCREATE DATABASE IF NOT EXISTS sandbox;\nCREATE SCHEMA IF NOT EXISTS sandbox.idea;\nGRANT USAGE ON DATABASE sandbox TO ROLE rl_app;\nGRANT ALL ON SCHEMA sandbox.idea TO ROLE rl_app;\nGRANT USAGE ON DATABASE sandbox TO ROLE rl_app_user;\nGRANT USAGE ON SCHEMA sandbox.idea TO ROLE rl_app_user;\n\nUSE ROLE rl_app;\nUSE SCHEMA sandbox.idea;\nDROP TABLE IF EXISTS app_data;\nCREATE TABLE app_data(a INT, b DOUBLE, c TEXT, region_id INT);\nINSERT INTO app_data(a, b, c, region_id) VALUES \n    (101, 101.101, 'one hundred one', 1),\n    (102, 102.102, 'one hundred two', 1),\n    (103, 103.103, 'one hundred three', 1),\n    (201, 101.101, 'two hundred one', 2),\n    (202, 202.202, 'two hundred two', 2),\n    (203, 203.203, 'two hundred three', 2),\n    (301, 301.301, 'three hundred one', 3),\n    (302, 302.302, 'three hundred two', 3),\n    (303, 303.303, 'three hundred three', 3),\n    (901, 901.901, 'nine hundred one', 9),\n    (902, 902.902, 'nine hundred two', 9),\n    (903, 903.903, 'nine hundred three', 9)\n;\nGRANT SELECT ON TABLE app_data TO ROLE rl_app_user;",
   "execution_count": null
  },
  {
   "cell_type": "markdown",
   "id": "310d4e81-8ad5-4c5d-b752-8453b5b5c1ea",
   "metadata": {
    "name": "cell3",
    "collapsed": false
   },
   "source": "## 3. Create the Entitlements Table\nNext we will create the entitlements table that will map which users have access to which region IDs. Additionally, we will allow the `CORTEX_USER_ROLE` to have access to all of the data, so we add a row in the entitlement table for each region ID."
  },
  {
   "cell_type": "code",
   "id": "709d3c18-36c7-4a51-96ad-a4de150faef6",
   "metadata": {
    "language": "sql",
    "name": "cell2"
   },
   "outputs": [],
   "source": "USE ROLE rl_app;\nUSE SCHEMA sandbox.idea;\n\nDROP TABLE IF EXISTS ent_region;\nCREATE TABLE ent_region(key STRING, value INT);\n\n-- Let's you to access all of the data\nINSERT INTO ent_region\n    SELECT CURRENT_USER() AS key, region_id AS value \n    FROM app_data \n    GROUP BY region_id\n;\n\n-- Now, let's restrict tenants to just their region\nINSERT INTO ent_region(key,value) VALUES \n    ('USER1', 1), \n    ('USER2', 2), \n    ('USER3', 3)\n;\n\nSELECT * FROM ent_region ORDER BY key, value;",
   "execution_count": null
  },
  {
   "cell_type": "markdown",
   "id": "d3e67956-3a32-4dca-985a-d71d1366f720",
   "metadata": {
    "name": "cell13",
    "collapsed": false
   },
   "source": "## 4. Set up the Row Access Policy\nIn order to make the Row Access Policy performant, we will create a memoizable function that collects all `VALUE`s that a given `KEY` has access to. Then, we create a Row Access Policy that uses that function, passing in the `CURRENT_USER()` as the argument. Since this is contant for the length of the query, the UDF will run once and the result reused (as is the behavior of memoizable functions)."
  },
  {
   "cell_type": "code",
   "id": "e31e059e-2614-4f5d-ba28-32195f00a080",
   "metadata": {
    "language": "sql",
    "name": "cell8"
   },
   "outputs": [],
   "source": "USE ROLE rl_app;\nUSE SCHEMA sandbox.idea;\n\nDROP FUNCTION IF EXISTS udf_regions(STRING);\nCREATE FUNCTION udf_regions(ikey STRING)\n    RETURNS ARRAY\n    MEMOIZABLE\n    AS 'SELECT ARRAY_AGG(value) FROM ent_region WHERE key = ikey';\n\nDROP ROW ACCESS POLICY IF EXISTS rap_region_user;\nCREATE ROW ACCESS POLICY rap_region_user\n    AS (region_id INT)\n    RETURNS BOOLEAN ->\n        ARRAY_CONTAINS(region_id, udf_regions(CURRENT_USER()))\n;",
   "execution_count": null
  },
  {
   "cell_type": "markdown",
   "id": "d540ce0d-2e1a-4179-acd1-6fb6872abab3",
   "metadata": {
    "name": "cell14",
    "collapsed": false
   },
   "source": "Now we add this Row Access Policy to the `REGION_ID` column of the `APP_DATA` table."
  },
  {
   "cell_type": "code",
   "id": "3fecc7b4-e635-44de-861b-65e58549bd23",
   "metadata": {
    "language": "sql",
    "name": "cell10"
   },
   "outputs": [],
   "source": "USE ROLE rl_app;\nUSE SCHEMA sandbox.idea;\n\nALTER TABLE app_data DROP ALL ROW ACCESS POLICIES; -- Just to make sure there aren't any others\nALTER TABLE app_data ADD ROW ACCESS POLICY rap_region_user ON (region_id);",
   "execution_count": null
  },
  {
   "cell_type": "markdown",
   "id": "ccecfd4e-24b3-4425-bb4f-13671ca9f135",
   "metadata": {
    "name": "cell15",
    "collapsed": false
   },
   "source": "## 5. Testing\nNow we can test the access to the `APP_DATA` table. First, let's use the current user and ensure that all data is accessible."
  },
  {
   "cell_type": "code",
   "id": "35e4aca7-9258-4ebf-931a-5fccac199036",
   "metadata": {
    "language": "sql",
    "name": "cell9"
   },
   "outputs": [],
   "source": "USE ROLE rl_app;\nUSE SCHEMA sandbox.idea;\nSELECT region_id, COUNT(*) FROM app_data GROUP BY region_id;",
   "execution_count": null
  },
  {
   "cell_type": "markdown",
   "id": "ee2e8660-1291-405a-9d47-de8b70b4ca3f",
   "metadata": {
    "name": "cell11",
    "collapsed": false
   },
   "source": "Next, let's use user `USER1` to see that only data from region `1` is accessible. To do this, you will need to log out of Snowsight and log in as `USER1`. In a worksheet run the following:"
  },
  {
   "cell_type": "code",
   "id": "d1950638-dddf-4dde-9688-d3e8fb96b1e2",
   "metadata": {
    "language": "sql",
    "name": "cell16"
   },
   "outputs": [],
   "source": "USE ROLE rl_app_user;\nUSE SCHEMA sandbox.idea;\nSELECT region_id, COUNT(*) FROM app_data GROUP BY region_id;",
   "execution_count": null
  },
  {
   "cell_type": "markdown",
   "id": "39a1a6a6-3abb-4629-b795-efd840b0e329",
   "metadata": {
    "name": "cell17",
    "collapsed": false
   },
   "source": "Feel free to try this with user `USER2` (to see data from region `2`), and `USER3` (to see data from region `3`)."
  },
  {
   "cell_type": "markdown",
   "id": "e92f8512-58f0-4003-9880-3a3b509a35f9",
   "metadata": {
    "name": "cell18",
    "collapsed": false
   },
   "source": "## 6. Application\nAt this point, we can turn our attention to a Streamlit in Snowflake that will only allow users to access the data in their allowed region(s). We will do this using the Snowsight UI.\n\n1. Open a new tab/window (or clone this tab), log into Snowflake, and navigate to the \"Projects\" sidebar and the \"Streamlit\" page.\n2. Make sure you have selected the `RL_APP` role as your current role. You can do this by clicking your initials in the lower left, choose the \"Switch Role\" option, and select the `RL_APP` role.\n3. Click the \"+ Streamlit App\" button to create a new app. You can name it whatever you would like (might I suggest \"RAP Current User\"). Choose the `SANDBOX` database and `IDEA` schema, and choose an app warehouse.\n4. In the editor pane, delete all of the code and replace it with the following:\n```python\n# Import python packages\nimport streamlit as st\nfrom snowflake.snowpark.context import get_active_session\n\n# Get the current credentials\nsession = get_active_session()\n\n# Write directly to the app\nst.title(f\"Using RAP with SIS\")\n\nuser = session.sql(\"SELECT current_user() AS user\").collect()[0]['USER']\nst.markdown(f\"Hello. Your current user is: :blue[{user}]\")\n\nst.markdown(\"Let's access the `APP_DATA` table\")\ndf = session.table(\"sandbox.idea.app_data\").to_pandas()\nst.dataframe(df)\n```\n\nWe now have a Streamlit that will show the current user and select data from the `APP_DATA` table. Running it as yourself you will see that you can see all the data.\n\nNow we need to share the Streamlit app with the `RL_APP_USER` role so that our other users can access it.\n\n5. Click the \"Share\" button in the upper right. In the dialog box, click the selectbox named \"Add role to share with\", enter `RL_APP_USER` and select it. Click \"Done\".\n\nAt this point, you can visit the Streamlit as `USER1`, `USER2`, or `USER3` and see that only some of the data is visible, specifically the data they have been granted entitlement to in our `ENT_REGION` table. Do this by logging out and back in as one of those users, visit the \"Projects\" sidebar and \"Streamlit\" page. Ensure that you have selected the `RL_APP_USER` as your current role (if not, follow the same instructions as step 2 above to change roles). Click on the Streamlit that you created and shared and you should see that the Stremalit identifies your current user and only shows you the data for the region that you have been entitled to see."
  }
 ]
}