[Python SDK][Auth] Refresh token id w/o gcloud cli (#1413)

* Auth: refresh token id w/o gcloud sdk

Signed-off-by: Oleksii Moskalenko <moskalenko.alexey@gmail.com>

* fix unit tests

Signed-off-by: Oleksii Moskalenko <moskalenko.alexey@gmail.com>

sdk/python/feast/grpc/auth.py

@@ -12,10 +12,13 @@
 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 # See the License for the specific language governing permissions and
 # limitations under the License.
+import json
 from http import HTTPStatus
 
 import grpc
+import requests
 from google.auth.exceptions import DefaultCredentialsError
+from google.auth.transport import requests as grequests
 
 from feast.config import Config
 from feast.constants import AuthProvider
@@ -98,10 +101,6 @@ def _refresh_token(self, config: Config):
             self._token = self._static_token
             return
 
-        import json
-
-        import requests
-
         headers_token = {"content-type": "application/json"}
         data_token = {
             "grant_type": config.get(opt.OAUTH_GRANT_TYPE),
@@ -157,7 +156,6 @@ def __init__(self, config: Config):
             config: Feast Configuration object
         """
         super(GoogleOpenIDAuthMetadataPlugin, self).__init__()
-        from google.auth.transport import requests
 
         self._static_token = None
         self._token = None
@@ -166,7 +164,7 @@ def __init__(self, config: Config):
         if config.exists(opt.AUTH_TOKEN):
             self._static_token = config.get(opt.AUTH_TOKEN)
 
-        self._request = requests.Request()
+        self._request = RequestWithTimeout(timeout=5)
         self._refresh_token()
 
     def get_signed_meta(self):
@@ -181,21 +179,13 @@ def _refresh_token(self):
             self._token = self._static_token
             return
 
-        # Try to find ID Token from Gcloud SDK
-        import subprocess
-
-        from google.auth import jwt
+        from google.oauth2.id_token import fetch_id_token
 
         try:
-            cli_output = subprocess.run(
-                ["gcloud", "auth", "print-identity-token"], stdout=subprocess.PIPE
-            )
-            token = cli_output.stdout.decode("utf-8").strip()
-            jwt.decode(token, verify=False)  # Ensure the token is valid
-            self._token = token
+            self._token = fetch_id_token(self._request, audience="feast.dev")
             return
-        except (ValueError, FileNotFoundError):
-            pass  # GCloud command not successful
+        except DefaultCredentialsError:
+            pass
 
         # Try to use Google Auth library to find ID Token
         from google import auth as google_auth
@@ -211,8 +201,7 @@ def _refresh_token(self):
 
         # Raise exception otherwise
         raise RuntimeError(
-            "Could not determine Google ID token. Please ensure that the Google Cloud SDK is installed and run: "
-            '"gcloud auth application-default login" or ensure that a service account can be found by setting'
+            "Could not determine Google ID token. Ensure that a service account can be found by setting"
             " the GOOGLE_APPLICATION_CREDENTIALS environmental variable to its path."
         )
 
@@ -234,3 +223,13 @@ def __call__(self, context, callback):
                 be invoked to pass in the authorization metadata.
         """
         callback(self.get_signed_meta(), None)
+
+
+class RequestWithTimeout(grequests.Request):
+    def __init__(self, *args, timeout=None, **kwargs):
+        self._timeout = timeout
+        super().__init__(*args, **kwargs)
+
+    def __call__(self, *args, **kwargs):
+        timeout = kwargs.pop("timeout", self._timeout)
+        return super().__call__(*args, timeout=timeout, **kwargs)

sdk/python/tests/grpc/test_auth.py

@@ -18,6 +18,7 @@
 from http import HTTPStatus
 from unittest.mock import call, patch
 
+from google.auth.exceptions import DefaultCredentialsError
 from pytest import fixture, raises
 
 from feast.config import Config
@@ -146,17 +147,14 @@ def test_get_auth_metadata_plugin_oauth_should_raise_when_config_is_incorrect(
         get_auth_metadata_plugin(config_with_missing_variable)
 
 
-@patch("google.auth.jwt.decode", return_value=GoogleMockResponse("jwt_token"))
-@patch(
-    "subprocess.run", return_value=GoogleMockResponse("std_output".encode("utf-8")),
-)
+@patch("google.oauth2.id_token.fetch_id_token", return_value="Some Token")
 def test_get_auth_metadata_plugin_google_should_pass_with_token_from_gcloud_sdk(
-    subprocess, jwt, config_google
+    fetch_id_token, config_google
 ):
     auth_metadata_plugin = get_auth_metadata_plugin(config_google)
     assert isinstance(auth_metadata_plugin, GoogleOpenIDAuthMetadataPlugin)
     assert auth_metadata_plugin.get_signed_meta() == (
-        ("authorization", "Bearer std_output"),
+        ("authorization", "Bearer Some Token"),
     )
 
 
@@ -167,11 +165,9 @@ def test_get_auth_metadata_plugin_google_should_pass_with_token_from_gcloud_sdk(
         GoogleDefaultResponse("project_id"),
     ],
 )
-@patch(
-    "subprocess.run", return_value=GoogleMockResponse("std_output".encode("utf-8")),
-)
+@patch("google.oauth2.id_token.fetch_id_token", side_effect=DefaultCredentialsError())
 def test_get_auth_metadata_plugin_google_should_pass_with_token_from_google_auth_lib(
-    subprocess, default, config_google
+    fetch_id_token, default, config_google
 ):
     auth_metadata_plugin = get_auth_metadata_plugin(config_google)
     assert isinstance(auth_metadata_plugin, GoogleOpenIDAuthMetadataPlugin)
@@ -187,11 +183,9 @@ def test_get_auth_metadata_plugin_google_should_pass_with_token_from_google_auth
         GoogleDefaultErrorResponse("project_id"),
     ],
 )
-@patch(
-    "subprocess.run", return_value=GoogleMockResponse("std_output".encode("utf-8")),
-)
+@patch("google.oauth2.id_token.fetch_id_token", side_effect=DefaultCredentialsError())
 def test_get_auth_metadata_plugin_google_should_raise_when_token_validation_fails(
-    subprocess, default, config_google
+    fetch_id_token, default, config_google
 ):
     with raises(RuntimeError):
         get_auth_metadata_plugin(config_google)