From c2ed01cbdf8340806f500b89b1079d8b8497d3cf Mon Sep 17 00:00:00 2001
From: Bhargav Dodla <13788369+EXPEbdodla@users.noreply.github.com>
Date: Wed, 2 Oct 2024 11:50:47 -0700
Subject: [PATCH] fix: Import grpc only for type checking (#141)

* fix: Import grpc only for type checking in errors.py

* chore: Disable warning for experimental SparkSource API

The commit disables a warning message for the experimental SparkSource API in the `spark_source.py` file. The warning was causing confusion for users, so it was decided to disable it. This change improves the user experience by removing unnecessary confusion.

* chore: Remove unused import warnings in spark_source.py

---------

Co-authored-by: Bhargav Dodla <bdodla@expediagroup.com>
---
 sdk/python/feast/errors.py                     | 18 +++++++++++++-----
 .../spark_offline_store/spark_source.py        | 15 +++++++--------
 2 files changed, 20 insertions(+), 13 deletions(-)

diff --git a/sdk/python/feast/errors.py b/sdk/python/feast/errors.py
index 06335a86a7..c8d9e5b5e0 100644
--- a/sdk/python/feast/errors.py
+++ b/sdk/python/feast/errors.py
@@ -1,11 +1,13 @@
 import importlib
 import json
 import logging
-from typing import Any, List, Optional, Set
+from typing import TYPE_CHECKING, Any, List, Optional, Set
 
 from colorama import Fore, Style
 from fastapi import status as HttpStatusCode
-from grpc import StatusCode as GrpcStatusCode
+
+if TYPE_CHECKING:
+    from grpc import StatusCode as GrpcStatusCode
 
 from feast.field import Field
 
@@ -15,7 +17,9 @@
 class FeastError(Exception):
     pass
 
-    def grpc_status_code(self) -> GrpcStatusCode:
+    def grpc_status_code(self) -> "GrpcStatusCode":
+        from grpc import StatusCode as GrpcStatusCode
+
         return GrpcStatusCode.INTERNAL
 
     def http_status_code(self) -> int:
@@ -89,7 +93,9 @@ def __init__(self, ds_name: str):
 class FeastObjectNotFoundException(FeastError):
     pass
 
-    def grpc_status_code(self) -> GrpcStatusCode:
+    def grpc_status_code(self) -> "GrpcStatusCode":
+        from grpc import StatusCode as GrpcStatusCode
+
         return GrpcStatusCode.NOT_FOUND
 
     def http_status_code(self) -> int:
@@ -509,7 +515,9 @@ class FeastPermissionError(FeastError, PermissionError):
     def __init__(self, details: str):
         super().__init__(f"Permission error:\n{details}")
 
-    def grpc_status_code(self) -> GrpcStatusCode:
+    def grpc_status_code(self) -> "GrpcStatusCode":
+        from grpc import StatusCode as GrpcStatusCode
+
         return GrpcStatusCode.PERMISSION_DENIED
 
     def http_status_code(self) -> int:
diff --git a/sdk/python/feast/infra/offline_stores/contrib/spark_offline_store/spark_source.py b/sdk/python/feast/infra/offline_stores/contrib/spark_offline_store/spark_source.py
index 2a2a7f2df0..8825d36383 100644
--- a/sdk/python/feast/infra/offline_stores/contrib/spark_offline_store/spark_source.py
+++ b/sdk/python/feast/infra/offline_stores/contrib/spark_offline_store/spark_source.py
@@ -1,10 +1,8 @@
 import logging
 import traceback
-import warnings
 from enum import Enum
 from typing import Any, Callable, Dict, Iterable, Optional, Tuple
 
-from feast import flags_helper
 from feast.data_source import DataSource
 from feast.errors import DataSourceNoNameException, DataSourceNotFoundException
 from feast.infra.offline_stores.offline_utils import get_temp_entity_table_name
@@ -61,12 +59,13 @@ def __init__(
             owner=owner,
         )
 
-        if not flags_helper.is_test():
-            warnings.warn(
-                "The spark data source API is an experimental feature in alpha development. "
-                "This API is unstable and it could and most probably will be changed in the future.",
-                RuntimeWarning,
-            )
+        # We recommend using SparkSource for our users. So disabling this warning as this can cause confusion.
+        # if not flags_helper.is_test():
+        #     warnings.warn(
+        #         "The spark data source API is an experimental feature in alpha development. "
+        #         "This API is unstable and it could and most probably will be changed in the future.",
+        #         RuntimeWarning,
+        #     )
 
         self.spark_options = SparkOptions(
             table=table,