Setup generators

-- Permission setup for {{app}}
USE ROLE {{privRole}};

-- Let the application see the database and its schemas
GRANT USAGE ON DATABASE {{db}} TO APPLICATION {{app}};
GRANT USAGE ON SCHEMA {{db}}.{{readSchema}} TO APPLICATION {{app}};
GRANT USAGE ON SCHEMA {{db}}.{{writeSchema}} TO APPLICATION {{app}};

-- Required to read tabular data into a graph
GRANT SELECT ON ALL TABLES IN SCHEMA {{db}}.{{readSchema}} TO APPLICATION {{app}};
GRANT SELECT ON ALL VIEWS IN SCHEMA {{db}}.{{readSchema}} TO APPLICATION {{app}};

-- Required to write computation results into a table
GRANT CREATE TABLE ON SCHEMA {{db}}.{{writeSchema}} TO APPLICATION {{app}};

-- Access to tables and views created later
CREATE DATABASE ROLE IF NOT EXISTS {{db}}.{{dbRole}};
GRANT SELECT ON FUTURE TABLES IN SCHEMA {{db}}.{{readSchema}} TO DATABASE ROLE {{db}}.{{dbRole}};
GRANT SELECT ON FUTURE VIEWS IN SCHEMA {{db}}.{{readSchema}} TO DATABASE ROLE {{db}}.{{dbRole}};
GRANT DATABASE ROLE {{db}}.{{dbRole}} TO APPLICATION {{app}};

-- Consumer roles, to be granted to your users
CREATE ROLE IF NOT EXISTS {{userRole}};
GRANT APPLICATION ROLE {{app}}.app_user TO ROLE {{userRole}};
CREATE ROLE IF NOT EXISTS {{adminRole}};
GRANT APPLICATION ROLE {{app}}.app_admin TO ROLE {{adminRole}};

-- Let users read the result tables
GRANT USAGE ON DATABASE {{db}} TO ROLE {{userRole}};
GRANT USAGE ON SCHEMA {{db}}.{{writeSchema}} TO ROLE {{userRole}};
GRANT SELECT ON FUTURE TABLES IN SCHEMA {{db}}.{{writeSchema}} TO ROLE {{userRole}};

-- Agent setup for {{agent}} on {{app}}

-- Part 1: privileges and caller grants. Run as ACCOUNTADMIN.
USE ROLE ACCOUNTADMIN;

-- The role that creates the agent and converses with it
CREATE ROLE IF NOT EXISTS {{role}};
GRANT APPLICATION ROLE {{app}}.app_admin TO ROLE {{role}};
GRANT APPLICATION ROLE {{app}}.app_user TO ROLE {{role}};
GRANT DATABASE ROLE SNOWFLAKE.CORTEX_AGENT_USER TO ROLE {{role}};
GRANT USAGE ON WAREHOUSE {{wh}} TO ROLE {{role}};
SET agent_setup_user = (SELECT CURRENT_USER());
GRANT ROLE {{role}} TO USER IDENTIFIER($agent_setup_user);

-- Working area for the views and result tables the agent creates
CREATE DATABASE IF NOT EXISTS {{workDb}};
CREATE SCHEMA IF NOT EXISTS {{workDb}}.{{workSchema}};

-- Application access to the source schema
CREATE DATABASE ROLE IF NOT EXISTS {{srcDb}}.{{srcDbRole}};
GRANT USAGE ON DATABASE {{srcDb}} TO DATABASE ROLE {{srcDb}}.{{srcDbRole}};
GRANT USAGE ON SCHEMA {{srcDb}}.{{srcSchema}} TO DATABASE ROLE {{srcDb}}.{{srcDbRole}};
GRANT SELECT ON ALL TABLES IN SCHEMA {{srcDb}}.{{srcSchema}} TO DATABASE ROLE {{srcDb}}.{{srcDbRole}};
GRANT SELECT ON ALL VIEWS IN SCHEMA {{srcDb}}.{{srcSchema}} TO DATABASE ROLE {{srcDb}}.{{srcDbRole}};
GRANT SELECT ON FUTURE TABLES IN SCHEMA {{srcDb}}.{{srcSchema}} TO DATABASE ROLE {{srcDb}}.{{srcDbRole}};
GRANT SELECT ON FUTURE VIEWS IN SCHEMA {{srcDb}}.{{srcSchema}} TO DATABASE ROLE {{srcDb}}.{{srcDbRole}};
GRANT DATABASE ROLE {{srcDb}}.{{srcDbRole}} TO APPLICATION {{app}};
GRANT DATABASE ROLE {{srcDb}}.{{srcDbRole}} TO ROLE {{role}};

-- The working database needs a database role of its own. If both schemas are in
-- the same database, drop this block and grant its statements to the role above.
CREATE DATABASE ROLE IF NOT EXISTS {{workDb}}.{{workDbRole}};
GRANT USAGE ON DATABASE {{workDb}} TO DATABASE ROLE {{workDb}}.{{workDbRole}};
GRANT USAGE ON SCHEMA {{workDb}}.{{workSchema}} TO DATABASE ROLE {{workDb}}.{{workDbRole}};
GRANT SELECT ON ALL TABLES IN SCHEMA {{workDb}}.{{workSchema}} TO DATABASE ROLE {{workDb}}.{{workDbRole}};
GRANT SELECT ON ALL VIEWS IN SCHEMA {{workDb}}.{{workSchema}} TO DATABASE ROLE {{workDb}}.{{workDbRole}};
GRANT SELECT ON FUTURE TABLES IN SCHEMA {{workDb}}.{{workSchema}} TO DATABASE ROLE {{workDb}}.{{workDbRole}};
GRANT SELECT ON FUTURE VIEWS IN SCHEMA {{workDb}}.{{workSchema}} TO DATABASE ROLE {{workDb}}.{{workDbRole}};
GRANT CREATE TABLE ON SCHEMA {{workDb}}.{{workSchema}} TO DATABASE ROLE {{workDb}}.{{workDbRole}};
GRANT CREATE VIEW ON SCHEMA {{workDb}}.{{workSchema}} TO DATABASE ROLE {{workDb}}.{{workDbRole}};
GRANT DATABASE ROLE {{workDb}}.{{workDbRole}} TO APPLICATION {{app}};
GRANT DATABASE ROLE {{workDb}}.{{workDbRole}} TO ROLE {{role}};

-- The consumer role reads the input data and the results itself
GRANT USAGE ON DATABASE {{srcDb}} TO ROLE {{role}};
GRANT USAGE ON DATABASE {{workDb}} TO ROLE {{role}};
GRANT USAGE ON SCHEMA {{srcDb}}.{{srcSchema}} TO ROLE {{role}};
GRANT SELECT ON ALL TABLES IN SCHEMA {{srcDb}}.{{srcSchema}} TO ROLE {{role}};
GRANT SELECT ON ALL VIEWS IN SCHEMA {{srcDb}}.{{srcSchema}} TO ROLE {{role}};
GRANT SELECT ON FUTURE TABLES IN SCHEMA {{srcDb}}.{{srcSchema}} TO ROLE {{role}};
GRANT USAGE ON SCHEMA {{workDb}}.{{workSchema}} TO ROLE {{role}};
GRANT SELECT ON FUTURE TABLES IN SCHEMA {{workDb}}.{{workSchema}} TO ROLE {{role}};
GRANT SELECT ON FUTURE VIEWS IN SCHEMA {{workDb}}.{{workSchema}} TO ROLE {{role}};

-- The agent's _explore_graph_data tool runs under Restricted Caller's Rights and
-- cannot read these schemas without the grants below. They require ACCOUNTADMIN
-- or a role with MANAGE CALLER GRANTS.
GRANT CALLER USAGE ON DATABASE {{srcDb}} TO APPLICATION {{app}};
GRANT CALLER USAGE ON DATABASE {{workDb}} TO APPLICATION {{app}};
GRANT CALLER USAGE ON SCHEMA {{srcDb}}.{{srcSchema}} TO APPLICATION {{app}};
GRANT INHERITED CALLER SELECT ON ALL TABLES IN SCHEMA {{srcDb}}.{{srcSchema}} TO APPLICATION {{app}};
GRANT INHERITED CALLER SELECT ON ALL VIEWS IN SCHEMA {{srcDb}}.{{srcSchema}} TO APPLICATION {{app}};
GRANT CALLER USAGE ON SCHEMA {{workDb}}.{{workSchema}} TO APPLICATION {{app}};
GRANT INHERITED CALLER SELECT ON ALL TABLES IN SCHEMA {{workDb}}.{{workSchema}} TO APPLICATION {{app}};
GRANT INHERITED CALLER SELECT ON ALL VIEWS IN SCHEMA {{workDb}}.{{workSchema}} TO APPLICATION {{app}};

-- The caller grants are scoped to the application and the schema pair, not to one
-- agent, so they survive agent re-creation and application upgrades. Reprint them
-- at any time with CALL {{app}}.GRAPH.GET_AGENT_CALLER_GRANTS('{{agent}}');

-- Part 2: create the agent. Run as the consumer role that will own it, not as
-- ACCOUNTADMIN.
USE ROLE {{role}};
USE WAREHOUSE {{wh}};
USE SCHEMA {{workDb}}.{{workSchema}};

CALL {{app}}.GRAPH.CREATE_AGENT(
    '{{agent}}',
    '{{srcDb}}.{{srcSchema}}',
    '{{workDb}}.{{workSchema}}'
);

-- The return value repeats the caller grants from part 1.

-- Part 1: enable PAT auth and mint the token. Run as ACCOUNTADMIN.
USE ROLE ACCOUNTADMIN;

-- Prerequisite: let the user mint and use a PAT without a network policy
CREATE AUTHENTICATION POLICY IF NOT EXISTS {{authPolicy}}
    PAT_POLICY = (NETWORK_POLICY_EVALUATION = ENFORCED_NOT_REQUIRED);
ALTER USER {{user}} SET AUTHENTICATION POLICY {{authPolicy}};

-- Step 1: enable PAT auth on the install (once per account)
CALL {{app}}.preview.set_enable_custom_credentials(TRUE);

-- Step 2a: mint a PAT bound to the role
ALTER USER {{user}} ADD PROGRAMMATIC ACCESS TOKEN {{patName}}
    DAYS_TO_EXPIRY = {{days}}
    ROLE_RESTRICTION = '{{role}}';

-- ---------------------------------------------------------------------------
-- STOP. Run everything above first, then copy the token_secret value from the
-- result, because Snowflake will not show it again, and paste it into
-- SECRET_STRING below before running the rest.
-- ---------------------------------------------------------------------------

-- Part 2: store the token and register the user.
USE ROLE ACCOUNTADMIN;

-- Step 2b: stash the token in a secret you own
CREATE OR REPLACE SECRET {{secretDb}}.{{secretSchema}}.{{secretName}}
    TYPE = GENERIC_STRING
    SECRET_STRING = '<paste_token_secret_from_part_1_here>';

-- Step 2c: let the application read the secret
GRANT USAGE ON DATABASE {{secretDb}} TO APPLICATION {{app}};
GRANT USAGE ON SCHEMA {{secretDb}}.{{secretSchema}} TO APPLICATION {{app}};
GRANT READ ON SECRET {{secretDb}}.{{secretSchema}}.{{secretName}} TO APPLICATION {{app}};

-- Step 2d: register the user with the application. Run this one in a session
-- opened as {{user}}, so that CURRENT_USER() resolves to them.
CALL {{app}}.preview.register_user_role(
    '{{role}}',
    '{{secretDb}}.{{secretSchema}}.{{secretName}}'
);

-- Step 2e: caller grants for the data the jobs read and write
USE ROLE ACCOUNTADMIN;
GRANT CALLER USAGE ON DATABASE {{dataDb}} TO APPLICATION {{app}};
GRANT CALLER USAGE ON SCHEMA {{dataDb}}.{{dataSchema}} TO APPLICATION {{app}};
GRANT CALLER CREATE TABLE ON SCHEMA {{dataDb}}.{{dataSchema}} TO APPLICATION {{app}};
GRANT INHERITED CALLER INSERT ON ALL TABLES IN SCHEMA {{dataDb}}.{{dataSchema}} TO APPLICATION {{app}};
GRANT INHERITED CALLER SELECT ON ALL TABLES IN SCHEMA {{dataDb}}.{{dataSchema}} TO APPLICATION {{app}};
GRANT INHERITED CALLER SELECT ON ALL VIEWS IN SCHEMA {{dataDb}}.{{dataSchema}} TO APPLICATION {{app}};

-- Step 3: after the user has run a job, check whose identity it used
SELECT user_name, role_name, LEFT(query_text, 120) AS sql, start_time
FROM SNOWFLAKE.ACCOUNT_USAGE.QUERY_HISTORY
WHERE start_time > DATEADD(minute, -30, CURRENT_TIMESTAMP())
  AND role_name = '{{role}}'
ORDER BY start_time DESC LIMIT 20;