-----------------------------
---  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);

-- 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', 
    -- The following bucket contains library dependencies used by Kumo. The Kumo Native app downloads
    -- these dependencies on startup and does not write any data to this bucket.
    'kumo-pyspark-venv.s3.amazonaws.com',
    'kumo-pyspark-venv.s3.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 network rules to allow access to customer's Snowflake account from Kumo running as a Native app.
-- Even when running as a Native app, the Snowflake connectors used by Kumo to access objects in customer's
-- account require using certain hostnames and these hostnames will need to be added to the firewall for access.
-- See https://docs.snowflake.com/en/sql-reference/functions/system_allowlist for more details.
CREATE OR REPLACE PROCEDURE CREATE_SNOWFLAKE_ALLOW_LIST()
RETURNS STRING
LANGUAGE PYTHON
RUNTIME_VERSION = '3.8'
PACKAGES = ('snowflake-snowpark-python')
HANDLER = 'main'
AS
$$
import json

def main(session):
    """
    Method to create network egress rules for SYSTEM$ALLOWLIST in the Snowflake account.
    """
    # get allow list
    allow_list = session.sql("SELECT SYSTEM$ALLOWLIST();").collect()[0].as_dict()

    allow_list_hosts = []
    if len(allow_list.values()) > 0:
        data = json.loads(list(allow_list.values())[0])

        # Iterate over each object in the array
        for entry in data:
            # Extract host, port, and type from each object
            if 'host' in entry and 'port' in entry:
                host = entry['host']
                port = entry['port']
                allow_list_hosts.append(f"'{host}:{port}'")
    
    session.sql(
        f"""CREATE OR REPLACE NETWORK RULE snowflake_allow_list
                MODE = EGRESS
                TYPE = HOST_PORT
                VALUE_LIST = ({', '.join(allow_list_hosts)})"""
    ).collect()

    return "Successfully created network rule snowflake_allow_list"
$$;

-- Call the procedure above to create a network rule to access Snowflake endpoints.
CALL CREATE_SNOWFLAKE_ALLOW_LIST();

-- 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, snowflake_allow_list)
    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;

-- GRANT Create Compute Pool privilege to KUMO_USER_ROLE
-- A compute pool is required to start the Kumo application.
GRANT CREATE COMPUTE POOL ON ACCOUNT 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>';
SET KUMOAPP_NAME = 'KUMO';
SET INSTANCE_TYPE = 'GPU_NV_M'; -- use 'GPU_NV_S' only for tables with less than 50M rows.

USE ROLE KUMO_USER_ROLE;

CREATE COMPUTE POOL IF NOT EXISTS KUMO_COMPUTE_POOL
  FOR APPLICATION IDENTIFIER($KUMOAPP_NAME)
  min_nodes = 1
  max_nodes = 1
  instance_family = $INSTANCE_TYPE;

GRANT USAGE ON COMPUTE POOL KUMO_COMPUTE_POOL TO APPLICATION IDENTIFIER($KUMOAPP_NAME);

DESCRIBE COMPUTE POOL KUMO_COMPUTE_POOL;

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

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

DROP COMPUTE POOL IF EXISTS KUMO_COMPUTE_POOL;

-- 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;