-- This code demonstrates running both a provider-run and a consumer-run analysis in a clean room.
-- Both analyses use code uploaded by the provider and consumer.
-- This is the consumer code.
-- For more information, see docs.snowflake.com/user-guide/cleanrooms/demo-flows/custom-code.

USE WAREHOUSE app_wh;
USE ROLE SAMOOHA_APP_ROLE;

SET cleanroom_name = 'Custom consumer UDF';
SET template_name = 'basic_template';

-- Install clean room.
CALL samooha_by_snowflake_local_db.consumer.install_cleanroom($cleanroom_name, $provider_locator);

-- Link your data into the clean room.
CALL samooha_by_snowflake_local_db.consumer.link_datasets(
  $cleanroom_name, ['SAMOOHA_SAMPLE_DATABASE.DEMO.CUSTOMERS']);

-- Create the template and the custom code.
CALL samooha_by_snowflake_local_db.consumer.generate_python_request_template(
  'say_hello',         -- Template uses this name to call your function
  ['name VARCHAR'],    -- Arguments and types for the function
  [],                  -- Libraries used
  [],                  -- Reserved
  'VARCHAR',           -- Return type
  'say_hello',         -- Standard main handler
                       -- Python implementation as UDF
  $$

def say_hello(name:str):
  return f'Hello, {name}!'
    $$
);

-- Here is a basic template that calls your function:

--SELECT
--  c.status,
--  c.hashed_email
--FROM IDENTIFIER( {{ my_table[0] }} ) as c
--WHERE cleanroom.custom_compare({{ min_status }}, {{ max_status }}, c.status);

-- Paste this template into the value returned by generate_python_request_template.
-- Escape all ' marks in the template, then send the template in a request
-- to the provider. We have already done the hard work for you :)
CALL samooha_by_snowflake_local_db.consumer.create_template_request(
  $cleanroom_name, 
  $template_name, 
  '
BEGIN


-- First define the Python UDF
CREATE OR REPLACE FUNCTION CLEANROOM.say_hello(name VARCHAR)
RETURNS VARCHAR
LANGUAGE PYTHON
RUNTIME_VERSION = 3.10
PACKAGES = ()

HANDLER = \'say_hello\'
AS $$

def say_hello(name:str):
  return f\'Hello, {name}!\'
    $$;
        

-- Then define and execute the SQL query
LET SQL_TEXT varchar := $$
SELECT cleanroom.say_hello(c.status)
FROM SAMOOHA_SAMPLE_DATABASE.DEMO.CUSTOMERS AS c
LIMIT 10;
$$;

-- Execute the query and return the result
LET RES resultset := (EXECUTE IMMEDIATE :SQL_TEXT);
RETURN TABLE(RES);

END;
  ');

-- SWITCH TO THE PROVIDER TO APPROVE THE REQUEST

-- Check status of the request.
CALL samooha_by_snowflake_local_db.consumer.list_template_requests($cleanroom_name);
CALL samooha_by_snowflake_local_db.consumer.view_added_templates($cleanroom_name);

-- Run the template.
CALL samooha_by_snowflake_local_db.consumer.run_analysis(
  $cleanroom_name,
  $template_name,
  ['SAMOOHA_SAMPLE_DATABASE.DEMO.CUSTOMERS'],
  [],
  {}
);

-- Enable provider-run analysis for a template. ---
-- Examine the provider's template and then approve it for provider run.
CALL samooha_by_snowflake_local_db.consumer.describe_cleanroom($cleanroom_name);
CALL samooha_by_snowflake_local_db.consumer.view_template_definition(
  $cleanroom_name,
  'say_hello');
CALL samooha_by_snowflake_local_db.consumer.approve_template(
  $cleanroom_name,
  'say_hello'
);

-- SWITCH BACK TO THE PROVIDER TO RUN THEIR OWN TEMPLATE.

-- Clean up when done.
CALL samooha_by_snowflake_local_db.consumer.uninstall_cleanroom($cleanroom_name);
