-----------------------------
---  USER INPUT REQUIRED  ---
-----------------------------
SET USER_WAREHOUSE = '<ADD WAREHOUSE TO USE>';

-- (Optional) Modify the following if a different name was used to install the app.
SET KUMOAPP_NAME = 'KUMO';
SET KUMOAPP_DB = 'KUMO_APP_DB';

SET KUMOAPP_USR = CONCAT(($KUMOAPP_NAME),'.APP_USER');

USE ROLE ACCOUNTADMIN;
USE WAREHOUSE IDENTIFIER($USER_WAREHOUSE);

-- Grant BIND SERVICE ENDPOINT privilege to the Kumo application to enable 
-- network ingress and access to the Kumo UI for users of Kumo in your account.
-- Details of the privilige can be found here
-- https://other-docs.snowflake.com/LIMITEDACCESS/native-apps/na-spcs-consumer#set-up-access-to-network-objects
GRANT BIND SERVICE ENDPOINT ON ACCOUNT TO APPLICATION IDENTIFIER($KUMOAPP_NAME);

-- Grant CREATE COMPUTE POOL privilege to allow the Kumo application to create compute pools
-- This allows for efficient management of compute pools by Kumo.
GRANT CREATE COMPUTE POOL ON ACCOUNT TO APPLICATION IDENTIFIER($KUMOAPP_NAME);

-- Share events from the application with the provider, KUMO.AI
ALTER APPLICATION IDENTIFIER($KUMOAPP_NAME) SET AUTHORIZE_TELEMETRY_EVENT_SHARING=true;

-- Database and schema to hold network rules for the Kumo application along
-- with usage for the Kumo application.
CREATE DATABASE IF NOT EXISTS IDENTIFIER($KUMOAPP_DB);
USE DATABASE IDENTIFIER($KUMOAPP_DB);
CREATE SCHEMA IF NOT EXISTS KUMO_SCHEMA;
USE SCHEMA KUMO_SCHEMA;

---------------------
-- EGRESS RULES FOR KUMO
---------------------

-- Network rule to allow access to Temporal Cloud.
CREATE OR REPLACE NETWORK RULE kumo_temporal_egress
  MODE = EGRESS
  TYPE = HOST_PORT
  VALUE_LIST = ('preprod.put0h.tmprl.cloud', 'preprod.put0h.tmprl.cloud:7233');

-- Access to AWS resources used by Kumo
CREATE OR REPLACE NETWORK RULE kumo_aws_egress
  MODE = EGRESS
  TYPE = HOST_PORT
  VALUE_LIST = (
    -- AWS Secrets Manager contains secrets Kumo uses to access other AWS services.
    -- No customer data or secrets are written to this.
    'secretsmanager.us-west-2.amazonaws.com'
);
  
-- Network rule to allow access to Mixpanel for logging product usage stats. While optional,
-- enabling Mixpanel will allow Kumo's customer success team provide better support to
-- users, by counting which pages were visited in the UI, and how many errors did they encounter.
-- No snowflake data, metadata (eg. predictive queries, table names, etc), or error logs are
-- written to Mixpanel from the Kumo Native app.
CREATE OR REPLACE NETWORK RULE kumo_mixpanel
  MODE = EGRESS
  TYPE = HOST_PORT
  VALUE_LIST = ('api.mixpanel.com');


-- Create an External access integration with the allowed network rules.
CREATE OR REPLACE EXTERNAL ACCESS INTEGRATION KUMO_EXTERNAL_ACCESS_INTEGRATION
    ALLOWED_NETWORK_RULES = (kumo_temporal_egress, kumo_aws_egress, kumo_mixpanel)
    ENABLED = true;

-- Grant usage on the KUMO_EXTERNAL_ACCESS_INTEGRATION to the Kump application.
GRANT USAGE ON INTEGRATION KUMO_EXTERNAL_ACCESS_INTEGRATION TO APPLICATION IDENTIFIER($KUMOAPP_NAME);

------------------------

-- Create KUMO_USER_ROLE and grant necessary privileges.
CREATE ROLE IF NOT EXISTS KUMO_USER_ROLE;

-- Grant the Application role to KUMO_USER_ROLE.
GRANT APPLICATION ROLE IDENTIFIER($KUMOAPP_USR) TO ROLE KUMO_USER_ROLE;

-- Grant USAGE on the database, schema and warehouse to KUMO_USER_ROLE.
GRANT USAGE ON DATABASE IDENTIFIER($KUMOAPP_DB) TO ROLE KUMO_USER_ROLE;
GRANT USAGE ON ALL SCHEMAS IN DATABASE IDENTIFIER($KUMOAPP_DB) TO ROLE KUMO_USER_ROLE;
GRANT USAGE ON WAREHOUSE IDENTIFIER($USER_WAREHOUSE) TO ROLE KUMO_USER_ROLE;

-----------------------------
---  USER INPUT REQUIRED  ---
-----------------------------
USE ROLE ACCOUNTADMIN;
-- Grant the KUMO_USER_ROLE to the DEFAULT_ROLE of any user who are expected to use the Kumo application.
GRANT ROLE KUMO_USER_ROLE TO ROLE <DEFAULT_USER_ROLE>;

-----------------------------
---  USER INPUT REQUIRED  ---
-----------------------------
-- Note that KUMO_USER_ROLE should be granted usage on this warehouse.
SET USER_WAREHOUSE = '<ADD WAREHOUSE TO USE>';
-- Modify the following if a different name was used to install the app.
SET KUMOAPP_NAME = 'KUMO';

USE ROLE KUMO_USER_ROLE;

USE WAREHOUSE IDENTIFIER($USER_WAREHOUSE);
SET KUMOAPP_COMPUTE = CONCAT(($KUMOAPP_NAME),'.PROCEDURES.START_APP');
CALL IDENTIFIER($KUMOAPP_COMPUTE)('USER_SCHEMA');

SET KUMOAPP_STOP = CONCAT(($KUMOAPP_NAME),'.PROCEDURES.SHUTDOWN_APP');
CALL IDENTIFIER($KUMOAPP_STOP)('USER_SCHEMA');

-- Description: This script is used to reset the account to the initial state before the Kumo app setup.

-- NOTE: All the objects created for the Kumo app will be deleted.
--       Any models trained or data stored in the Kumo app will be lost.

USE ROLE ACCOUNTADMIN;
-----------------------------
---  USER INPUT REQUIRED  ---
-----------------------------

SET USER_WAREHOUSE = '<ADD WAREHOUSE TO USE>';
SET KUMOAPP_NAME = 'KUMO';
SET KUMOAPP_DB = 'KUMO_APP_DB';

USE WAREHOUSE IDENTIFIER($USER_WAREHOUSE);

-- Delete the Kumo application and all associated objects.
DROP APPLICATION IF EXISTS IDENTIFIER($KUMOAPP_NAME) CASCADE;

-- Drop the role created for the user of Kumo app.
DROP ROLE IF EXISTS KUMO_USER_ROLE;

-- Drop the external access rule created for the Kump app.
DROP EXTERNAL ACCESS INTEGRATION IF EXISTS KUMO_EXTERNAL_ACCESS_INTEGRATION;

USE DATABASE IDENTIFIER($KUMOAPP_DB);

CREATE OR REPLACE PROCEDURE DROP_ALL_NETWORK_RULES(db_name varchar)
RETURNS STRING
LANGUAGE PYTHON
RUNTIME_VERSION = '3.8'
PACKAGES = ('snowflake-snowpark-python')
HANDLER = 'main'
AS
$$

def main(session, db_name):
    """
    Method to drop all network egress rules created for Kumo.
    """
    # get allow list and 
    db_names = session.sql(f"SHOW DATABASES like '{db_name}'").collect()
    if len(db_names) == 0:
        return f"No Database found with name {db_name}. No rules to delete"

    rows = session.sql(f"SHOW NETWORK RULES IN DATABASE {db_name}").collect()
    dropped_rules = []
    for row in rows:
        row_dict = row.as_dict()
        if 'name' in row_dict:
            rule = f"{db_name}.{row_dict['schema_name']}.{row_dict['name']}"
            session.sql(f"DROP NETWORK RULE IF EXISTS {rule}").collect()
            dropped_rules.append(rule)
    return f"Successfully deleted {len(rows)} network rules in Database {db_name}: {', '.join(dropped_rules)}"
$$;

CALL DROP_ALL_NETWORK_RULES($KUMOAPP_DB);

-- Drop Database created for the Kumo app
DROP DATABASE IF EXISTS IDENTIFIER($KUMOAPP_DB);

-- Set app name if it's different from the default.
SET KUMOAPP_NAME = 'KUMO';
-- Enables event sharing for the Native app.
ALTER APPLICATION IDENTIFIER($KUMOAPP_NAME) SET AUTHORIZE_TELEMETRY_EVENT_SHARING=true;