let lazy_import_transformers force load the snowflake-connector

Signed-off-by: HH <[email protected]>
flyteorg · Sep 10, 2023 · c16591a · c16591a
1 parent afe39bc
commit c16591a
Show file tree

Hide file tree

Showing 2 changed files with 7 additions and 3 deletions.
diff --git a/flytekit/core/type_engine.py b/flytekit/core/type_engine.py
@@ -854,11 +854,14 @@ def lazy_import_transformers(cls):
             register_arrow_handlers()
         if is_imported("google.cloud.bigquery"):
             register_bigquery_handlers()
-        if is_imported("snowflake.connector"):
-            register_snowflake_handlers()
         if is_imported("numpy"):
             from flytekit.types import numpy  # noqa: F401
 
+        try:
+            register_snowflake_handlers()
+        except ValueError as e:
+            logger.debug(f"Attempted to register the Snowflake handler but failed due to: {str(e)}")
+
     @classmethod
     def to_literal_type(cls, python_type: Type) -> LiteralType:
         """

diff --git a/flytekit/types/structured/__init__.py b/flytekit/types/structured/__init__.py
@@ -81,5 +81,6 @@ def register_snowflake_handlers():
 
     except ImportError:
         logger.info(
-            "We won't register snowflake handler for structured dataset because " "we can't find package snowflake"
+            "We won't register snowflake handler for structured dataset because "
+            "we can't find package snowflakee-connector-python"
         )