From d7041f4cce813d349e9016da55d65a65c1ec2355 Mon Sep 17 00:00:00 2001
From: Mark Snidal <mark.snidal@gmail.com>
Date: Tue, 5 Sep 2023 02:36:03 -0400
Subject: [PATCH] fix: Handle unknown postgres source types gracefully (#3634)

Handle more pg types gracefully

Signed-off-by: Mark Snidal <mark.snidal@gmail.com>
---
 sdk/python/feast/type_map.py | 21 +++++++++++++++++++--
 1 file changed, 19 insertions(+), 2 deletions(-)

diff --git a/sdk/python/feast/type_map.py b/sdk/python/feast/type_map.py
index df85346283..52e2890750 100644
--- a/sdk/python/feast/type_map.py
+++ b/sdk/python/feast/type_map.py
@@ -873,13 +873,26 @@ def feast_value_type_to_pa(
 
 
 def pg_type_code_to_pg_type(code: int) -> str:
-    return {
+    """ Map the postgres type code a Feast type string
+
+    Rather than raise an exception on an unknown type, we return the
+    string representation of the type code. This way rather than raising 
+    an exception on unknown types, Feast will just skip the problem columns.
+
+    Note that json and jsonb are not supported but this shows up in the
+    log as a warning. Since postgres allows custom types we return an unknown for those cases.
+
+    See: https://jdbc.postgresql.org/documentation/publicapi/index.html?constant-values.html
+    """
+    PG_TYPE_MAP = {
         16: "boolean",
         17: "bytea",
         20: "bigint",
         21: "smallint",
         23: "integer",
         25: "text",
+        114: "json",
+        199: "json[]",
         700: "real",
         701: "double precision",
         1000: "boolean[]",
@@ -905,7 +918,11 @@ def pg_type_code_to_pg_type(code: int) -> str:
         1700: "numeric",
         2950: "uuid",
         2951: "uuid[]",
-    }[code]
+        3802: "jsonb",
+        3807: "jsonb[]",
+    }
+
+    return PG_TYPE_MAP.get(code, "unknown")
 
 
 def pg_type_code_to_arrow(code: int) -> str: