Upload and run custom functions in a clean room¶

Overview¶

You can upload custom Python UDFs and UDTFs into your clean room and run them from your templates to perform complex data actions. These actions include machine learning or customized data manipulation within a query, as part of a single-step or multi-step flow. Python is the only coding language supported for custom UDFs.

Your uploaded code can import and use packages from an approved bundle of Python packages and the Snowpark API.

Both providers and consumers can upload custom Python code into a clean room, although the process is different for providers and consumers. Each bundle of uploaded code can define multiple functions that call each other, but a bundle exposes only one handler function. This handler function can be called by templates created or run by anyone who uses the clean room. If the code creates internal tables, these tables can be accessed as described in Designing multi-step flows.

Uploaded code can’t be deleted, but it can be updated.

This page shows you how to upload and run custom Python UDFs and UDTFs as a provider or a consumer.

Provider-submitted code¶

Provider-submitted functions can be uploaded as inline code or from a Snowflake stage. Both techniques are covered here.

Your uploaded code can natively import and use packages from an approved set of Python packages. If you need a non-default package, you must use Snowpark Container Services in a clean room to host your code.

You can’t view uploaded code, even your own code, so be sure to include a copy of exactly what you upload into a clean room.

Here is a high level view of how a provider adds code to a clean room:

1. The provider creates and configures the clean room in the normal way.

2. The provider uploads code by calling provider.load_python_into_cleanroom. You can either upload your code inline directly within that procedure, or upload a code file to a stage, then provide the stage location to that procedure.

   Although your code can include multiple functions, only one handler is exposed for each upload. To expose multiple functions to templates, upload each handler by calling provider.load_python_into_cleanroom.

3. After each successful code upload, a new patch version of the clean room is generated. You must then increase the default version by calling provider.set_default_release_directive with the new patch number. If the clean room is exposed externally, security checks are run before installing your code, and you must call provider.view_cleanroom_scan_status to confirm that security checks have passed before incrementing the default version.

4. You create and upload a custom template that calls your code. The template calls the handler function using the cleanroom scope, that is: cleanroom.my_function. For example, a template that calls a custom simple_add function that you uploaded might look like this:

   SELECT cleanroom.simple_add(1, 2), cleanroom.simple_add({{ price | sqlsafe | int }}, {{ tax | sqlsafe | int }})
   

   Copy

5. The consumer runs your template the same way as any other template.

   Tip

   If the consumer encounters a mount error when they install a clean room with custom code, this can indicate a syntax error in the code.

You can find code examples demonstrating this flow in the provider-written code example section.

CALL samooha_by_snowflake_local_db.provider.load_python_into_cleanroom(
$cleanroom_name,
'simple_add',                         -- Name used to call the UDF from a template.
['first INTEGER', 'second INTEGER'],  -- Arguments of the UDF, specified as '<variable_name> <SQL type>' pairs.
['numpy', 'pandas'],                  -- Packages imported by the UDF.
'INTEGER',                            -- SQL return type of UDF.
'add_two',                            -- Handler function in your code called when external name is called.
$$
import numpy as np   # Not used, but you can load supported packages.
import pandas as pd

def add_two(first, second):
    return first + second
$$
);

-- Save the following code as reverser.py:
--import numpy as np
--def main(some_string):
--  '''Return the reverse of a string plus a random number 1-10'''
--  return some_string[::-1] + str(np.random.randint(1,10))

-- Get the stage for your clean room.
CALL samooha_by_snowflake_local_db.provider.get_stage_for_python_files($cleanroom_name);

-- Save the file to the stage. Here is how to do it by using the Snowflake CLI
PUT file://~/reverser.py <STAGE_NAME> overwrite=True auto_compress=False;

-- Load the code from the stage into the clean room.
CALL samooha_by_snowflake_local_db.provider.load_python_into_cleanroom(
    $cleanroom_name,
    'reverse', -- Name used to call the function
    ['some_string  STRING'], -- Arguments and SQL types
    ['numpy'],               -- Any required packages
    ['/reverser.py'],        -- Relative path to file on stage
    'STRING',                -- Return type
    'reverser.main'          -- <FILE_NAME>.<FUNCTION_NAME>
);

-- Uploading code, even from a stage, increases the patch number.
CALL samooha_by_snowflake_local_db.provider.set_default_release_directive(
  $cleanroom_name, 'V1_0', <NEW_PATCH_NUMBER>);

-- Upload a template that calls the function.
CALL samooha_by_snowflake_local_db.provider.add_custom_sql_template(
    $cleanroom_name,
    $udf_template_name,
    $$
    SELECT
      p.status,
      cleanroom.reverse(p.status)
    FROM SAMOOHA_SAMPLE_DATABASE.DEMO.CUSTOMERS AS p
    LIMIT 100;
    $$
);

-- Switch to the consumer account and run the template to see the results.

Consumer-submitted code¶

A consumer can submit UDF or UDTF code and call it from a custom template. Consumer-uploaded code is bundled into a single procedure along with a custom template, and uploaded in a single procedure call. The consumer code is tied directly to that template and cannot be called by other templates.

To upload code as a consumer, you should understand custom template syntax and how to submit a consumer-defined template.

Note that any code uploaded by a consumer can be seen by the provider when they request permission to upload. The consumer code is also visible whenever a provider or consumer examines the template.

Here is an overview of the steps to upload custom consumer code:

1. The provider creates the clean room in the standard way and then invites the consumer.

2. The consumer installs and configures the clean room in the standard way.

3. The consumer prepares a template. The template calls the UDF or UDTF within the cleanroom namespace. For example, to call the consumer-defined calculate_tax function, a simple template might look like the following snippet:

   SELECT {{ cleanroom.calculate_tax(p.cost) }} AS Tax FROM my_db.my_sch.sales AS p;
   

   Copy

4. The consumer prepares their Python code. We recommend using double quotation marks (" ") rather than single quotation marks (' ') in your code to avoid extra escaping needed later. Your code can reference a bundle of selected Python libraries.

5. The consumer passes their Python code into consumer.generate_python_request_template. The procedure returns the Python code as a stored procedure, with a placeholder for the custom JinjaSQL template. There are several multi-line strings in the template that use $$ as multi-line delimiters.

6. Replace the template placeholder in the output from generate_python_request_template with your JinjaSQL template.

7. In the combined template, escape any single quotes like this: \'. This is because single quotes will be used as the outermost delimiter for the entire multi-line procedure string. For example:

     BEGIN
   
     CREATE OR REPLACE FUNCTION CLEANROOM.custom_compare(min_status STRING, max_status STRING, this_status STRING)
     RETURNS boolean
     LANGUAGE PYTHON
     RUNTIME_VERSION = 3.10
     PACKAGES = (\'numpy\')
   
     HANDLER = \'custom_compare\'
     AS $$
     import numpy as np
   
     def custom_compare(min_status:str, max_status:str, this_status:str):
       statuses = [\'MEMBER\', \'SILVER\', \'GOLD\', \'PLATINUM\']
       return ((statuses.index(this_status) >= statuses.index(min_status)) &
               (statuses.index(this_status) <= statuses.index(max_status)))
     $$;
   
     -- Custom template
     LET SQL_TEXT varchar := $$
     SELECT
       c.status,
       c.hashed_email
     FROM IDENTIFIER( {{ my_table[0] }} ) as c
     WHERE cleanroom.custom_compare({{ min_status }}, {{ max_status }}, c.status);
     $$;
   
     LET RES resultset := (EXECUTE IMMEDIATE :SQL_TEXT);
     RETURN TABLE(RES);
   
     END;
   

   Copy

8. Call consumer.create_template_request with your combined template. Use single quotation marks (' ') instead of double dollar sign delimiters ($$...$$) around the code you provide for stored procedure in the template_definition argument. For example:

   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.custom_compare(min_status STRING, max_status STRING, this_status STRING)
   RETURNS boolean
   LANGUAGE PYTHON
   RUNTIME_VERSION = 3.10
   PACKAGES = (\'numpy\')
   
   HANDLER = \'custom_compare\'
   AS $$
   import numpy as np
   
   def custom_compare(min_status:str, max_status:str, this_status:str):
     statuses = [\'MEMBER\', \'SILVER\', \'GOLD\', \'PLATINUM\']
     return ((statuses.index(this_status) >= statuses.index(min_status)) &
             (statuses.index(this_status) <= statuses.index(max_status)))
       $$;
   
   -- Then define and execute the SQL query.
   LET SQL_TEXT varchar := $$
   SELECT
     c.status,
     c.hashed_email
   FROM IDENTIFIER( {{ my_table[0] }} ) as c
   WHERE cleanroom.custom_compare({{ min_status }}, {{ max_status }}, c.status);
   $$;
   
   -- Execute the query and then return the result.
   LET RES resultset := (EXECUTE IMMEDIATE :SQL_TEXT);
   RETURN TABLE(RES);
   
   END;
   ');
   

   Copy

9. The consumer and provider continue with the standard consumer-defined template flow:

   1. The provider views the template request (provider.list_pending_template_requests) and then approves it by calling approve_template_request. In the request, the provider can see the template and the bundled code.

   2. The consumer checks the request status (consumer.list_template_requests), and when the status is APPROVED, runs the template (consumer.run_analysis).