From 177dfc660d99e8261febd7ef86c2878105ec85a6 Mon Sep 17 00:00:00 2001
From: Miles Adkins <miles.adkins@snowflake.com>
Date: Sat, 27 Aug 2022 22:30:02 -0500
Subject: [PATCH] chore: Add tagging for snowflake engine

Signed-off-by: Miles Adkins <miles.adkins@snowflake.com>
---
 .../utils/snowflake/snowpark/snowflake_udfs.py | 18 ++++++++++++++++++
 1 file changed, 18 insertions(+)

diff --git a/sdk/python/feast/infra/utils/snowflake/snowpark/snowflake_udfs.py b/sdk/python/feast/infra/utils/snowflake/snowpark/snowflake_udfs.py
index 7fde4dd3a1c..02311ca55d6 100644
--- a/sdk/python/feast/infra/utils/snowflake/snowpark/snowflake_udfs.py
+++ b/sdk/python/feast/infra/utils/snowflake/snowpark/snowflake_udfs.py
@@ -1,3 +1,4 @@
+import sys
 from binascii import unhexlify
 
 import pandas
@@ -24,6 +25,8 @@
 # ValueType.BYTES = 1
 @vectorized(input=pandas.DataFrame)
 def feast_snowflake_binary_to_bytes_proto(df):
+    sys._xoptions["snowflake_partner_attribution"].append("feast")
+
     df = list(
         map(
             ValueProto.SerializeToString,
@@ -45,6 +48,8 @@ def feast_snowflake_binary_to_bytes_proto(df):
 # ValueType.STRING = 2
 @vectorized(input=pandas.DataFrame)
 def feast_snowflake_varchar_to_string_proto(df):
+    sys._xoptions["snowflake_partner_attribution"].append("feast")
+
     df = list(
         map(
             ValueProto.SerializeToString,
@@ -66,6 +71,8 @@ def feast_snowflake_varchar_to_string_proto(df):
 # ValueType.INT32 = 3
 @vectorized(input=pandas.DataFrame)
 def feast_snowflake_number_to_int32_proto(df):
+    sys._xoptions["snowflake_partner_attribution"].append("feast")
+
     df = list(
         map(
             ValueProto.SerializeToString,
@@ -87,6 +94,8 @@ def feast_snowflake_number_to_int32_proto(df):
 # ValueType.INT64 = 4
 @vectorized(input=pandas.DataFrame)
 def feast_snowflake_number_to_int64_proto(df):
+    sys._xoptions["snowflake_partner_attribution"].append("feast")
+
     df = list(
         map(
             ValueProto.SerializeToString,
@@ -110,6 +119,8 @@ def feast_snowflake_number_to_int64_proto(df):
 # ValueType.FLOAT = 5 & ValueType.DOUBLE = 6
 @vectorized(input=pandas.DataFrame)
 def feast_snowflake_float_to_double_proto(df):
+    sys._xoptions["snowflake_partner_attribution"].append("feast")
+
     df = list(
         map(
             ValueProto.SerializeToString,
@@ -131,6 +142,8 @@ def feast_snowflake_float_to_double_proto(df):
 # ValueType.BOOL = 7
 @vectorized(input=pandas.DataFrame)
 def feast_snowflake_boolean_to_bool_boolean_proto(df):
+    sys._xoptions["snowflake_partner_attribution"].append("feast")
+
     df = list(
         map(
             ValueProto.SerializeToString,
@@ -152,6 +165,7 @@ def feast_snowflake_boolean_to_bool_boolean_proto(df):
 # ValueType.UNIX_TIMESTAMP = 8
 @vectorized(input=pandas.DataFrame)
 def feast_snowflake_timestamp_to_unix_timestamp_proto(df):
+    sys._xoptions["snowflake_partner_attribution"].append("feast")
 
     df = list(
         map(
@@ -177,6 +191,8 @@ def feast_snowflake_timestamp_to_unix_timestamp_proto(df):
 # converts 1 to n many entity keys to a single binary for lookups
 @vectorized(input=pandas.DataFrame)
 def feast_serialize_entity_keys(df):
+    sys._xoptions["snowflake_partner_attribution"].append("feast")
+
     join_keys = create_entity_dict(df[0].values[0], df[2].values[0])
 
     df = pandas.DataFrame.from_dict(
@@ -222,6 +238,8 @@ def feast_serialize_entity_keys(df):
 # converts 1 to n many entity keys to a single binary for lookups
 @vectorized(input=pandas.DataFrame)
 def feast_entity_key_proto_to_string(df):
+    sys._xoptions["snowflake_partner_attribution"].append("feast")
+
     join_keys = create_entity_dict(df[0].values[0], df[2].values[0])
 
     df = pandas.DataFrame.from_dict(