Snowflake UDF example to trigger Immuta API calls

Implementation Workflow

1. Create a lambda function on AWS Lambda

2. Create REST API endpoint on AWS API Gateway

3. Create AWS IAM role for Snowflake connectivity

4. (optional) Validate and associate API Gateway to AWS Lambda function for above steps

5. Create Snowflake integration object

6. Create Snowflake external function

7. Trigger the external function from within Snowflake

The following steps creates a UDF and performs a project-context switching/reset back to No Current Project on Immuta if the user has previously set a different project. It is triggered within Snowflake by passing two parameters: user_api_key and user_email_uid. The UDF is based on a Snowflake integration object’s AWS API gateway info and triggers an AWS Lambda function to perform the API call to Immuta using the provided user_api_key and user_email_uid. Additionally, this approach allows for other Immuta API endpoints to be hit through the same Lambda function.

The following steps create a UDF and perform a project-context switch or reset back to "No Current Project" on Immuta if the user has previously set a different project. This process is triggered within Snowflake by passing two parameters: user_api_key and user_email_uid. The UDF utilizes the AWS API Gateway information of a Snowflake integration object and triggers an AWS Lambda function to make an API call to Immuta using the provided user_api_key and user_email_uid. This is an example; other Immuta API endpoints can be accessed using a similar approach.

Step 1: Create a lambda function on AWS

• Record the name of the Lambda function.

• Set the function payload from the currentProject parameter to None.

• (NOTE) Replace different endpoint and payload as needed to interact with different Immuta endpoints for various functions.

# AWS Lambda function example to reset Immuta project context to 'No Current Project'
import json
import requests

def lambda_handler(event, context):
   # parse Snowflake values
   api_key = event['data'][0][1]
   user_id = event['data'][0][2]
  
   # set api_key for header; hardcode the api_key for certain api calls if users can't access definition
   headers = {
       'Authorization': 'Bearer ' + api_key,
       'Content-Type': 'application/json',
   }

   # setting payload for the endpoint, update for other endpoints
   data = json.dumps({
       "preferences": {
           "sortProjectState": {
               "column": "name",
               "order": "asc",
               "size": 12
           },
           "tabProjectState": 0,
           "tabDataSourceState": 1,
           "currentProject": None
       }
   })
  
   # immuta endpoint and user uid, replace with your Immuta instance URL
   url = 'https://<your immuta instance>.immutacloud.com/bim/iam/immuta/user/' + user_id + '/profile'
  
   method = 'POST'
   response = requests.request(method, url, headers=headers, json=data)


   json_data = json.loads(response.text)
   currentProject = json_data['preferences']['currentProject']
  
   # set value to return to Snowflake 
   value_to_be_returned={'data':[[0,'success']]}
   json_string_to_return = json.dumps({"data": value_to_be_returned})

   return value_to_be_returned

Step 2: Create REST API endpoint on AWS API Gateway

• Create and set api gateway

  • Create Resources with endpoint

  • Create Method, i.e. POST method to set_project_null

  • Associate the Lambda function to the POST request: i.e: bc_set_project_null

  • Deploy API and note the Invoke URL

Step 3: Create AWS IAM role for Snowflake connectivity

• Create IAM and add APIGateway and Lambda access to it

• Note the ARN info

Step 4: (optional) Validate and associate API Gateway to AWS Lambda function

• Go to the AWS lambda function and add it to the lambda function as fit.

Step 5: Create Snowflake integration object

• Snowflake integration object: provide the AWS IAM role info, and the API-gateway endpoint.

• Describing the Snowflake integration object will provide the API_AWS_IAM_USE_ARN and the API_AWS_EXTERNAL_ID, which need to be provided to the AWS IAM role on AWS.

• On AWS: navigate to the IAM role add the above information to the Trust Relationships for the role:

Step 6: Create Snowflake external function

• Create an external function which takes two parameters: api_key and user_id(email)

• Associate the function with the integration object and the API Gateway endpoint

Step 7: Trigger the external function from within Snowflake

• Test the lambda function has been throughly tested before deployment.

• Execute a SELECT query on the function using the parameter values shown in the image below.

• The project context in the Immuta instance will be reset or set to "No Current Project."

snowflakeudfaws-lambda