-- The following code shows how a provider implements a custom UDF and UDTF in a clean room.
-- 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 = 'Provider UDF example';
SET udf_template_name = 'basic_UDF_template';
SET udtf_template_name = 'basic_UDTF_template';
SET consumer_locator = '<consumer_locator>';
SET consumer_account_name = '<org_name>.<account_name>';
    
--Create the clean room.
CALL samooha_by_snowflake_local_db.provider.cleanroom_init($cleanroom_name, 'INTERNAL');

-- Link data.
CALL samooha_by_snowflake_local_db.provider.link_datasets($cleanroom_name,
['SAMOOHA_SAMPLE_DATABASE.DEMO.CUSTOMERS']);

---- Inline UDF code ----
-- Note the new patch number returned by this procedure. Keep this number for later use.
CALL samooha_by_snowflake_local_db.provider.load_python_into_cleanroom(
    $cleanroom_name, 
    'get_next_status',  -- Name of the UDF. Can be different from the handler.
    ['status VARCHAR'], -- Arguments of the UDF, specified as (variable name, SQL type).
    ['numpy'],          -- Packages needed by UDF.
    'VARCHAR',          -- Return type of UDF.
    'get_next_status',  -- Handler.
    $$
import numpy as np
def get_next_status(status):
  """Return the next higher status, or a random status
  if no matching status found or at the top of the list."""

  statuses = ['MEMBER', 'SILVER', 'GOLD', 'PLATINUM', 'DIAMOND']
  try:
    return statuses[statuses.index(status.upper()) + 1]
  except:
    return 'NO MATCH'
    $$
);

-- Add custom UDF template.
-- Note the new patch number returned by this procedure. Keep this number for later use.
CALL samooha_by_snowflake_local_db.provider.add_custom_sql_template(
    $cleanroom_name, 
    $udf_template_name,
    $$
    SELECT
      p.status,
      cleanroom.get_next_status(p.status) AS NEXT_STATUS,
    FROM SAMOOHA_SAMPLE_DATABASE.DEMO.CUSTOMERS AS p
    LIMIT 100;
    $$
);

---- Add custom UDTF. ----
CALL samooha_by_snowflake_local_db.provider.load_python_into_cleanroom(
    $cleanroom_name, 
    'get_info',  -- Name of the UDTF. Can be different from the handler.
    ['hashed_email VARCHAR', 'days_active INT', 'status VARCHAR', 'income VARCHAR'],   -- Name/Type arguments of the UDTF.
    ['numpy'],         -- Packages used by UDTF.
    'TABLE(hashed_email VARCHAR, months_active INT, level VARCHAR)',  -- Return type of UDTF.
    'GetSomeVals',     -- Handler class name.
$$
class GetSomeVals:
  def __init__(self):
    self.month_days = 30

  def process(self, hashed_email, days_active, status, income):
    '''Change days into rough months, and also return whether we
    think the user's membership status is lower, higher, or equal to 
    what is expected, based on their income.'''
    
    months_active = days_active // self.month_days
    brackets = ['0-50K', '50K-100K', '100K-250K', '250K+']
    statuses = ['MEMBER', 'SILVER', 'GOLD', 'PLATINUM']
    if(statuses.index(status) < brackets.index(income)):
      level = 'low'
    elif(statuses.index(status) > brackets.index(income)):
      level = 'high'
    else:
      level = 'equal'
    
    yield(hashed_email, months_active, level)
$$
);

-- Add custom UDTF template.
CALL samooha_by_snowflake_local_db.provider.add_custom_sql_template(
    $cleanroom_name, 
    $udtf_template_name,
    $$
    SELECT
      p.hashed_email,
      p1.months_active,
      p.status,
      p.income_bracket,
      p1.level
    FROM SAMOOHA_SAMPLE_DATABASE.DEMO.CUSTOMERS AS p
    INNER JOIN TABLE(cleanroom.get_info(p.hashed_email, p.days_active, p.status, p.income_bracket)) AS p1
    WHERE p.hashed_email = p1.hashed_email
    LIMIT 100;
    $$
);

-- Not needed for INTERNAL clean rooms, but adding here as a reminder to use it for EXTERNAL clean rooms.
-- NOT_REVIEWED means it can be used internally.
CALL samooha_by_snowflake_local_db.provider.view_cleanroom_scan_status($cleanroom_name);

-- See the latest patch number, in case you didn't save the version when you
-- uploaded the previous code.
SHOW VERSIONS IN APPLICATION PACKAGE SAMOOHA_CLEANROOM_<cleanroom_name>

-- Set the release directive specified by the last load_python_into_cleanroom call.
CALL samooha_by_snowflake_local_db.provider.set_default_release_directive($cleanroom_name, 'V1_0', <PATCH_NUMBER>);

-- Add consumers.
CALL samooha_by_snowflake_local_db.provider.add_consumers($cleanroom_name, $consumer_locator, $consumer_account_name);
    
-- Publish the clean room.
CALL samooha_by_snowflake_local_db.provider.create_or_update_cleanroom_listing($cleanroom_name);

-- SWITCH TO CONSUMER TO INSTALL THE CLEAN ROOM AND RUN THE TEMPLATE.

-- Clean up afterward.
CALL samooha_by_snowflake_local_db.provider.drop_cleanroom($cleanroom_name);
