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 <[email protected]>
ExpediaGroup · Oct 2, 2024 · c2ed01c · c2ed01c
1 parent 49a9c5f
commit c2ed01c
Show file tree

Hide file tree

Showing 2 changed files with 20 additions and 13 deletions.
diff --git 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
@@ -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,