databricks-sdk 0.32.2__py3-none-any.whl → 0.32.3__py3-none-any.whl

databricks/sdk/credentials_provider.py

@@ -9,14 +9,15 @@ import pathlib
 import platform
 import subprocess
 import sys
+import time
 from datetime import datetime
-from typing import Callable, Dict, List, Optional, Union
+from typing import Callable, Dict, List, Optional, Tuple, Union
 
-import google.auth
+import google.auth  # type: ignore
 import requests
-from google.auth import impersonated_credentials
-from google.auth.transport.requests import Request
-from google.oauth2 import service_account
+from google.auth import impersonated_credentials  # type: ignore
+from google.auth.transport.requests import Request  # type: ignore
+from google.oauth2 import service_account  # type: ignore
 
 from .azure import add_sp_management_token, add_workspace_id_header
 from .oauth import (ClientCredentials, OAuthClient, Refreshable, Token,
@@ -698,6 +699,90 @@ def metadata_service(cfg: 'Config') -> Optional[CredentialsProvider]:
     return inner
 
 
+# This Code is derived from Mlflow DatabricksModelServingConfigProvider
+# https://github.com/mlflow/mlflow/blob/1219e3ef1aac7d337a618a352cd859b336cf5c81/mlflow/legacy_databricks_cli/configure/provider.py#L332
+class ModelServingAuthProvider():
+    _MODEL_DEPENDENCY_OAUTH_TOKEN_FILE_PATH = "/var/credentials-secret/model-dependencies-oauth-token"
+
+    def __init__(self):
+        self.expiry_time = -1
+        self.current_token = None
+        self.refresh_duration = 300 # 300 Seconds
+
+    def should_fetch_model_serving_environment_oauth(self) -> bool:
+        """
+        Check whether this is the model serving environment
+        Additionally check if the oauth token file path exists
+        """
+
+        is_in_model_serving_env = (os.environ.get("IS_IN_DB_MODEL_SERVING_ENV")
+                                   or os.environ.get("IS_IN_DATABRICKS_MODEL_SERVING_ENV") or "false")
+        return (is_in_model_serving_env == "true"
+                and os.path.isfile(self._MODEL_DEPENDENCY_OAUTH_TOKEN_FILE_PATH))
+
+    def get_model_dependency_oauth_token(self, should_retry=True) -> str:
+        # Use Cached value if it is valid
+        if self.current_token is not None and self.expiry_time > time.time():
+            return self.current_token
+
+        try:
+            with open(self._MODEL_DEPENDENCY_OAUTH_TOKEN_FILE_PATH) as f:
+                oauth_dict = json.load(f)
+                self.current_token = oauth_dict["OAUTH_TOKEN"][0]["oauthTokenValue"]
+                self.expiry_time = time.time() + self.refresh_duration
+        except Exception as e:
+            # sleep and retry in case of any race conditions with OAuth refreshing
+            if should_retry:
+                logger.warning("Unable to read oauth token on first attmept in Model Serving Environment",
+                               exc_info=e)
+                time.sleep(0.5)
+                return self.get_model_dependency_oauth_token(should_retry=False)
+            else:
+                raise RuntimeError(
+                    "Unable to read OAuth credentials from the file mounted in Databricks Model Serving"
+                ) from e
+        return self.current_token
+
+    def get_databricks_host_token(self) -> Optional[Tuple[str, str]]:
+        if not self.should_fetch_model_serving_environment_oauth():
+            return None
+
+        # read from DB_MODEL_SERVING_HOST_ENV_VAR if available otherwise MODEL_SERVING_HOST_ENV_VAR
+        host = os.environ.get("DATABRICKS_MODEL_SERVING_HOST_URL") or os.environ.get(
+            "DB_MODEL_SERVING_HOST_URL")
+        token = self.get_model_dependency_oauth_token()
+
+        return (host, token)
+
+
+@credentials_strategy('model-serving', [])
+def model_serving_auth(cfg: 'Config') -> Optional[CredentialsProvider]:
+    try:
+        model_serving_auth_provider = ModelServingAuthProvider()
+        if not model_serving_auth_provider.should_fetch_model_serving_environment_oauth():
+            logger.debug("model-serving: Not in Databricks Model Serving, skipping")
+            return None
+        host, token = model_serving_auth_provider.get_databricks_host_token()
+        if token is None:
+            raise ValueError(
+                "Got malformed auth (empty token) when fetching auth implicitly available in Model Serving Environment. Please contact Databricks support"
+            )
+        if cfg.host is None:
+            cfg.host = host
+    except Exception as e:
+        logger.warning("Unable to get auth from Databricks Model Serving Environment", exc_info=e)
+        return None
+
+    logger.info("Using Databricks Model Serving Authentication")
+
+    def inner() -> Dict[str, str]:
+        # Call here again to get the refreshed token
+        _, token = model_serving_auth_provider.get_databricks_host_token()
+        return {"Authorization": f"Bearer {token}"}
+
+    return inner
+
+
 class DefaultCredentials:
     """ Select the first applicable credential provider from the chain """
 
@@ -706,7 +791,7 @@ class DefaultCredentials:
         self._auth_providers = [
             pat_auth, basic_auth, metadata_service, oauth_service_principal, azure_service_principal,
             github_oidc_azure, azure_cli, external_browser, databricks_cli, runtime_native_auth,
-            google_credentials, google_id
+            google_credentials, google_id, model_serving_auth
         ]
 
     def auth_type(self) -> str:

databricks/sdk/version.py

@@ -1 +1 @@
-__version__ = '0.32.2'
+__version__ = '0.32.3'

{databricks_sdk-0.32.2.dist-info → databricks_sdk-0.32.3.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: databricks-sdk
-Version: 0.32.2
+Version: 0.32.3
 Summary: Databricks SDK for Python (Beta)
 Home-page: https://databricks-sdk-py.readthedocs.io
 Author: Serge Smertin

{databricks_sdk-0.32.2.dist-info → databricks_sdk-0.32.3.dist-info}/RECORD

@@ -6,7 +6,7 @@ databricks/sdk/casing.py,sha256=NKYPrfPbQjM7lU4hhNQK3z1jb_VEA29BfH4FEdby2tg,1137
 databricks/sdk/clock.py,sha256=Ivlow0r_TkXcTJ8UXkxSA0czKrY0GvwHAeOvjPkJnAQ,1360
 databricks/sdk/config.py,sha256=UaD-UcgvvohbrDmvbQgUt-KFd8FP1w3iWvaocsoIz9k,21169
 databricks/sdk/core.py,sha256=PL3RdTMEOtcy4KPt6TKuL6ylm8RaHBXiqGbYwjRu5Gk,14679
-databricks/sdk/credentials_provider.py,sha256=8AEDLUgQfCLA5ODVGYsS_70N61a9yYv7137G8wDilsY,30781
+databricks/sdk/credentials_provider.py,sha256=H_eHS2uMf8avNHCfES6xxS9ybSY5_6Ipt-MxA8bKOVo,34720
 databricks/sdk/data_plane.py,sha256=Er2z2fT-KVupJKzGozGGZ-jCQ3AmDWq-DZppahIK6tU,2591
 databricks/sdk/dbutils.py,sha256=HFCuB-el6SFKhF8qRfJxYANtyLTm-VG9GtQuQgZXFkM,15741
 databricks/sdk/environments.py,sha256=5KoVuVfF-ZX17rua1sH3EJCCtniVrREXBXsMNDEV-UU,4293
@@ -14,7 +14,7 @@ databricks/sdk/oauth.py,sha256=KzcJPYLL3JL6RDvf_Q8SDAaF9xSaoYNCRD4rYInZDuo,18319
 databricks/sdk/py.typed,sha256=pSvaHpbY1UPNEXyVFUjlgBhjPFZMmVC_UNrPC7eMOHI,74
 databricks/sdk/retries.py,sha256=WgLh12bwdBc6fCQlaig3kKu18cVhPzFDGsspvq629Ew,2454
 databricks/sdk/useragent.py,sha256=I2-VnJSE6cg9QV4GXkoQSkHsEB3bDvRGgkawbBNl4G0,5540
-databricks/sdk/version.py,sha256=mteW834a5f5Gxn1SNzRvKxq9Y4nN1qsJoM5ubib79tI,23
+databricks/sdk/version.py,sha256=lcjcbT3Cv9Zt--ZYcNwIJDgsWDwjGtgW49gCMxzU3XE,23
 databricks/sdk/_widgets/__init__.py,sha256=Qm3JB8LmdPgEn_-VgxKkodTO4gn6OdaDPwsYcDmeIRI,2667
 databricks/sdk/_widgets/default_widgets_utils.py,sha256=Rk59AFzVYVpOektB_yC_7j-vSt5OdtZA85IlG0kw0xA,1202
 databricks/sdk/_widgets/ipywidgets_utils.py,sha256=P-AyGeahPiX3S59mxpAMgffi4gyJ0irEOY7Ekkn9nQ0,2850
@@ -57,9 +57,9 @@ databricks/sdk/service/sharing.py,sha256=kalJYd0v1SwuGhlCaq4l2ZhzNlev9OwNbCXFIOK
 databricks/sdk/service/sql.py,sha256=RaXIYMDtHbhvB7gtSMyvQsqiO_E0cMz5NXeTsrqtPVk,334558
 databricks/sdk/service/vectorsearch.py,sha256=a5Y4vrS_oAJJqa69XwKMANhGuZi5glS0PSXBXz1bKGU,62961
 databricks/sdk/service/workspace.py,sha256=FKLf5esRmfFstIXo7HQg6HQCzQ2svrb6ulr8yzZ7-8U,101182
-databricks_sdk-0.32.2.dist-info/LICENSE,sha256=afBgTZo-JsYqj4VOjnejBetMuHKcFR30YobDdpVFkqY,11411
-databricks_sdk-0.32.2.dist-info/METADATA,sha256=MDIMnh2mylfaS-wyBerMRQvfRuBPl08p7mWbKu8stPk,37967
-databricks_sdk-0.32.2.dist-info/NOTICE,sha256=Qnc0m8JjZNTDV80y0h1aJGvsr4GqM63m1nr2VTypg6E,963
-databricks_sdk-0.32.2.dist-info/WHEEL,sha256=GV9aMThwP_4oNCtvEC2ec3qUYutgWeAzklro_0m4WJQ,91
-databricks_sdk-0.32.2.dist-info/top_level.txt,sha256=7kRdatoSgU0EUurRQJ_3F1Nv4EOSHWAr6ng25tJOJKU,11
-databricks_sdk-0.32.2.dist-info/RECORD,,
+databricks_sdk-0.32.3.dist-info/LICENSE,sha256=afBgTZo-JsYqj4VOjnejBetMuHKcFR30YobDdpVFkqY,11411
+databricks_sdk-0.32.3.dist-info/METADATA,sha256=dJZB_OYTTrlvOV4QB9YzKWQIdsI_SgZRYvLIk0NETTw,37967
+databricks_sdk-0.32.3.dist-info/NOTICE,sha256=Qnc0m8JjZNTDV80y0h1aJGvsr4GqM63m1nr2VTypg6E,963
+databricks_sdk-0.32.3.dist-info/WHEEL,sha256=GV9aMThwP_4oNCtvEC2ec3qUYutgWeAzklro_0m4WJQ,91
+databricks_sdk-0.32.3.dist-info/top_level.txt,sha256=7kRdatoSgU0EUurRQJ_3F1Nv4EOSHWAr6ng25tJOJKU,11
+databricks_sdk-0.32.3.dist-info/RECORD,,