From 651ce341687034ce07ca959f805f3c90dccfd4cc Mon Sep 17 00:00:00 2001
From: sfc-gh-madkins <82121043+sfc-gh-madkins@users.noreply.github.com>
Date: Mon, 1 Aug 2022 12:24:21 -0400
Subject: [PATCH] fix: Snowflake_online_read fix (#2988)

Signed-off-by: Miles Adkins <miles.adkins@snowflake.com>
---
 .../feast/infra/online_stores/snowflake.py    | 19 ++++++++++---------
 .../templates/snowflake/feature_store.yaml    |  1 +
 2 files changed, 11 insertions(+), 9 deletions(-)

diff --git a/sdk/python/feast/infra/online_stores/snowflake.py b/sdk/python/feast/infra/online_stores/snowflake.py
index e5e7b680be..73c68e4bc0 100644
--- a/sdk/python/feast/infra/online_stores/snowflake.py
+++ b/sdk/python/feast/infra/online_stores/snowflake.py
@@ -96,11 +96,11 @@ def online_write_batch(
             for j, (feature_name, val) in enumerate(values.items()):
                 df.loc[j, "entity_feature_key"] = serialize_entity_key(
                     entity_key,
-                    entity_key_serialization_version=entity_key_serialization_version,
+                    entity_key_serialization_version,
                 ) + bytes(feature_name, encoding="utf-8")
                 df.loc[j, "entity_key"] = serialize_entity_key(
                     entity_key,
-                    entity_key_serialization_version=entity_key_serialization_version,
+                    entity_key_serialization_version,
                 )
                 df.loc[j, "feature_name"] = feature_name
                 df.loc[j, "value"] = val.SerializeToString()
@@ -156,12 +156,18 @@ def online_read(
 
         result: List[Tuple[Optional[datetime], Optional[Dict[str, ValueProto]]]] = []
 
+        entity_key_serialization_version = (
+            config.entity_key_serialization_version
+            if config.entity_key_serialization_version
+            else 2
+        )
+
         entity_fetch_str = ",".join(
             [
                 (
                     "TO_BINARY("
                     + hexlify(
-                        serialize_entity_key(combo[0])
+                        serialize_entity_key(combo[0], entity_key_serialization_version)
                         + bytes(combo[1], encoding="utf-8")
                     ).__str__()[1:]
                     + ")"
@@ -187,15 +193,10 @@ def online_read(
                 .fetch_pandas_all()
             )
 
-        entity_key_serialization_version = (
-            config.entity_key_serialization_version
-            if config.entity_key_serialization_version
-            else 2
-        )
         for entity_key in entity_keys:
             entity_key_bin = serialize_entity_key(
                 entity_key,
-                entity_key_serialization_version=entity_key_serialization_version,
+                entity_key_serialization_version,
             )
             res = {}
             res_ts = None
diff --git a/sdk/python/feast/templates/snowflake/feature_store.yaml b/sdk/python/feast/templates/snowflake/feature_store.yaml
index 3e2e3c3cea..39f266f89f 100644
--- a/sdk/python/feast/templates/snowflake/feature_store.yaml
+++ b/sdk/python/feast/templates/snowflake/feature_store.yaml
@@ -9,6 +9,7 @@ offline_store:
     role: SNOWFLAKE_ROLE
     warehouse: SNOWFLAKE_WAREHOUSE
     database: SNOWFLAKE_DATABASE
+entity_key_serialization_version: 2
 online_store:
     type: snowflake.online
     account: SNOWFLAKE_DEPLOYMENT_URL