databricks-sdk 0.44.0__py3-none-any.whl → 0.44.1__py3-none-any.whl

databricks/sdk/credentials_provider.py

@@ -9,6 +9,7 @@ import pathlib
 import platform
 import subprocess
 import sys
+import threading
 import time
 from datetime import datetime
 from typing import Callable, Dict, List, Optional, Tuple, Union
@@ -723,14 +724,17 @@ def metadata_service(cfg: 'Config') -> Optional[CredentialsProvider]:
 # This Code is derived from Mlflow DatabricksModelServingConfigProvider
 # https://github.com/mlflow/mlflow/blob/1219e3ef1aac7d337a618a352cd859b336cf5c81/mlflow/legacy_databricks_cli/configure/provider.py#L332
 class ModelServingAuthProvider():
+    USER_CREDENTIALS = "user_credentials"
+
     _MODEL_DEPENDENCY_OAUTH_TOKEN_FILE_PATH = "/var/credentials-secret/model-dependencies-oauth-token"
 
-    def __init__(self):
+    def __init__(self, credential_type: Optional[str]):
         self.expiry_time = -1
         self.current_token = None
         self.refresh_duration = 300 # 300 Seconds
+        self.credential_type = credential_type
 
-    def should_fetch_model_serving_environment_oauth(self) -> bool:
+    def should_fetch_model_serving_environment_oauth() -> bool:
         """
         Check whether this is the model serving environment
         Additionally check if the oauth token file path exists
@@ -739,15 +743,15 @@ class ModelServingAuthProvider():
         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))
+                and os.path.isfile(ModelServingAuthProvider._MODEL_DEPENDENCY_OAUTH_TOKEN_FILE_PATH))
 
-    def get_model_dependency_oauth_token(self, should_retry=True) -> str:
+    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:
+            with open(ModelServingAuthProvider._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
@@ -757,32 +761,43 @@ class ModelServingAuthProvider():
                 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)
+                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_invokers_token(self):
+        current_thread = threading.current_thread()
+        thread_data = current_thread.__dict__
+        invokers_token = None
+        if "invokers_token" in thread_data:
+            invokers_token = thread_data["invokers_token"]
+
+        if invokers_token is None:
+            raise RuntimeError("Unable to read Invokers Token in Databricks Model Serving")
+
+        return invokers_token
+
     def get_databricks_host_token(self) -> Optional[Tuple[str, str]]:
-        if not self.should_fetch_model_serving_environment_oauth():
+        if not ModelServingAuthProvider.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)
+        if self.credential_type == ModelServingAuthProvider.USER_CREDENTIALS:
+            return (host, self._get_invokers_token())
+        else:
+            return (host, self._get_model_dependency_oauth_token())
 
 
-@credentials_strategy('model-serving', [])
-def model_serving_auth(cfg: 'Config') -> Optional[CredentialsProvider]:
+def model_serving_auth_visitor(cfg: 'Config',
+                               credential_type: Optional[str] = None) -> 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
+        model_serving_auth_provider = ModelServingAuthProvider(credential_type)
         host, token = model_serving_auth_provider.get_databricks_host_token()
         if token is None:
             raise ValueError(
@@ -793,7 +808,6 @@ def model_serving_auth(cfg: 'Config') -> Optional[CredentialsProvider]:
     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]:
@@ -804,6 +818,15 @@ def model_serving_auth(cfg: 'Config') -> Optional[CredentialsProvider]:
     return inner
 
 
+@credentials_strategy('model-serving', [])
+def model_serving_auth(cfg: 'Config') -> Optional[CredentialsProvider]:
+    if not ModelServingAuthProvider.should_fetch_model_serving_environment_oauth():
+        logger.debug("model-serving: Not in Databricks Model Serving, skipping")
+        return None
+
+    return model_serving_auth_visitor(cfg)
+
+
 class DefaultCredentials:
     """ Select the first applicable credential provider from the chain """
 
@@ -846,3 +869,35 @@ class DefaultCredentials:
         raise ValueError(
             f'cannot configure default credentials, please check {auth_flow_url} to configure credentials for your preferred authentication method.'
         )
+
+
+class ModelServingUserCredentials(CredentialsStrategy):
+    """
+    This credential strategy is designed for authenticating the Databricks SDK in the model serving environment using user-specific rights. 
+    In the model serving environment, the strategy retrieves a downscoped user token from the thread-local variable. 
+    In any other environments, the class defaults to the DefaultCredentialStrategy. 
+    To use this credential strategy, instantiate the WorkspaceClient with the ModelServingUserCredentials strategy as follows:
+
+    invokers_client = WorkspaceClient(credential_strategy = ModelServingUserCredentials())
+    """
+
+    def __init__(self):
+        self.credential_type = ModelServingAuthProvider.USER_CREDENTIALS
+        self.default_credentials = DefaultCredentials()
+
+    def auth_type(self):
+        if ModelServingAuthProvider.should_fetch_model_serving_environment_oauth():
+            return "model_serving_" + self.credential_type
+        else:
+            return self.default_credentials.auth_type()
+
+    def __call__(self, cfg: 'Config') -> CredentialsProvider:
+        if ModelServingAuthProvider.should_fetch_model_serving_environment_oauth():
+            header_factory = model_serving_auth_visitor(cfg, self.credential_type)
+            if not header_factory:
+                raise ValueError(
+                    f"Unable to authenticate using {self.credential_type} in Databricks Model Serving Environment"
+                )
+            return header_factory
+        else:
+            return self.default_credentials(cfg)

databricks/sdk/mixins/jobs.py

@@ -11,9 +11,10 @@ class JobsExt(jobs.JobsAPI):
                 include_history: Optional[bool] = None,
                 include_resolved_values: Optional[bool] = None,
                 page_token: Optional[str] = None) -> jobs.Run:
-        """
-        This method fetches the details of a run identified by `run_id`. If the run has multiple pages of tasks or iterations,
-        it will paginate through all pages and aggregate the results.
+        """Get a single job run.
+
+        Retrieve the metadata of a run. If a run has multiple pages of tasks, it will paginate through all pages of tasks, iterations, job_clusters, job_parameters, and repair history.
+
         :param run_id: int
           The canonical identifier of the run for which to retrieve the metadata. This field is required.
         :param include_history: bool (optional)
@@ -21,8 +22,9 @@ class JobsExt(jobs.JobsAPI):
         :param include_resolved_values: bool (optional)
           Whether to include resolved parameter values in the response.
         :param page_token: str (optional)
-          To list the next page or the previous page of job tasks, set this field to the value of the
-          `next_page_token` or `prev_page_token` returned in the GetJob response.
+          To list the next page of job tasks, set this field to the value of the `next_page_token` returned in
+          the GetJob response.
+
         :returns: :class:`Run`
         """
         run = super().get_run(run_id,
@@ -34,6 +36,7 @@ class JobsExt(jobs.JobsAPI):
         # When querying a ForEach task run, a page token is returned when there are more than 100 iterations. Only a single task is returned, corresponding to the ForEach task itself. Therefore, the client only reads the iterations from the next page and not the tasks.
         is_paginating_iterations = run.iterations is not None and len(run.iterations) > 0
 
+        # runs/get response includes next_page_token as long as there are more pages to fetch.
         while run.next_page_token is not None:
             next_run = super().get_run(run_id,
                                        include_history=include_history,
@@ -43,7 +46,10 @@ class JobsExt(jobs.JobsAPI):
                 run.iterations.extend(next_run.iterations)
             else:
                 run.tasks.extend(next_run.tasks)
+            # Each new page of runs/get response includes the next page of the job_clusters, job_parameters, and repair history.
+            run.job_clusters.extend(next_run.job_clusters)
+            run.job_parameters.extend(next_run.job_parameters)
+            run.repair_history.extend(next_run.repair_history)
             run.next_page_token = next_run.next_page_token
 
-        run.prev_page_token = None
         return run

databricks/sdk/version.py

@@ -1 +1 @@
-__version__ = '0.44.0'
+__version__ = '0.44.1'

{databricks_sdk-0.44.0.dist-info → databricks_sdk-0.44.1.dist-info}/METADATA

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

{databricks_sdk-0.44.0.dist-info → databricks_sdk-0.44.1.dist-info}/RECORD

@@ -7,7 +7,7 @@ databricks/sdk/casing.py,sha256=NKYPrfPbQjM7lU4hhNQK3z1jb_VEA29BfH4FEdby2tg,1137
 databricks/sdk/clock.py,sha256=Ivlow0r_TkXcTJ8UXkxSA0czKrY0GvwHAeOvjPkJnAQ,1360
 databricks/sdk/config.py,sha256=fUCRZX7tpY4LNvh3ilPtbaoC6wlFauWj9jyONsMDLWM,20314
 databricks/sdk/core.py,sha256=s8W2UlRg8y9vCencdiFocYs49hLS3bltswLi00_cn38,4086
-databricks/sdk/credentials_provider.py,sha256=IZyjoxdKGwDqhRJ6jTFT-mS1lf8l1LDpxNBbPSdDHlY,35311
+databricks/sdk/credentials_provider.py,sha256=ioH9ZYNc7e9L2KDCQRL2oOUiFifzHqqn11a2wG6kxj0,37817
 databricks/sdk/data_plane.py,sha256=Kr81rnBlgqnlwNlDbV4JWzzTA5NSVlf0J-DQOylp0w4,2582
 databricks/sdk/dbutils.py,sha256=HFCuB-el6SFKhF8qRfJxYANtyLTm-VG9GtQuQgZXFkM,15741
 databricks/sdk/environments.py,sha256=5KoVuVfF-ZX17rua1sH3EJCCtniVrREXBXsMNDEV-UU,4293
@@ -15,7 +15,7 @@ databricks/sdk/oauth.py,sha256=ZlIzEGlKTUgGGgLfv5NQJr3Y_mWpKgTr8-hUEwwqfEE,23861
 databricks/sdk/py.typed,sha256=pSvaHpbY1UPNEXyVFUjlgBhjPFZMmVC_UNrPC7eMOHI,74
 databricks/sdk/retries.py,sha256=kdCKGIJjSkGLZYmQ0oI_hiGj7FP7MIqK9-nIr7WbykU,2574
 databricks/sdk/useragent.py,sha256=o9cojoaVwI7C6tbIZy6jcQ8QiYuUmdL5_zATu6IZSaw,7373
-databricks/sdk/version.py,sha256=MMPeQ2-kYJhgLEqI-j8q_6GpHFGW12eIl3b2Nt7TutU,23
+databricks/sdk/version.py,sha256=-1tkgtTZjJHtveKwFUf8Bh2OZhXwVDjVJDxyd7uV6Ic,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
@@ -34,7 +34,7 @@ databricks/sdk/logger/round_trip_logger.py,sha256=SMtHDfdqy5Noge2iZO-LpuEm92rz3A
 databricks/sdk/mixins/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 databricks/sdk/mixins/compute.py,sha256=khb00BzBckc4RLUF4-GnNMCSO5lXKt_XYMM3IhiUxlA,11237
 databricks/sdk/mixins/files.py,sha256=G5Tv0lbQGgOECwvFxmACX5X24ZMP6ZQD7ZudtkORbDw,28283
-databricks/sdk/mixins/jobs.py,sha256=PqiYWi0O-drOjAD6fsptHuAj6IofCdR5e_F3I2BzHqQ,2622
+databricks/sdk/mixins/jobs.py,sha256=xDXHKG5dSTF-K9tlyfwSu47get_7PF8koBnhHXnROLw,2966
 databricks/sdk/mixins/open_ai_client.py,sha256=gvh8xKm1O8B23gNmIv6vxSxk8tNmphUqvS-XkkS260U,4776
 databricks/sdk/mixins/workspace.py,sha256=dWMNvuEi8jJ5wMhrDt1LiqxNdWSsmEuDTzrcZR-eJzY,4896
 databricks/sdk/runtime/__init__.py,sha256=9NnZkBzeZXZRQxcE1qKzAszQEzcpIgpL7lQzW3_kxEU,7266
@@ -61,9 +61,9 @@ databricks/sdk/service/sharing.py,sha256=vs69Y2FE28A4XnLxDGKL2HvV-LKgB-DV7nSnbYc
 databricks/sdk/service/sql.py,sha256=GU_8ALx1r3lgu_FCmEJs8zTcWFKtJhKQ2CJcIcJgxzo,399590
 databricks/sdk/service/vectorsearch.py,sha256=5p5pW94Bv_Q2tw4j8kFb35nAoFa9GUG5FIHTdfAHWps,77997
 databricks/sdk/service/workspace.py,sha256=BCoi43R1L2eJI9DYq9vwCVdjbMsdLuzDebN6AZvT4kg,128751
-databricks_sdk-0.44.0.dist-info/LICENSE,sha256=afBgTZo-JsYqj4VOjnejBetMuHKcFR30YobDdpVFkqY,11411
-databricks_sdk-0.44.0.dist-info/METADATA,sha256=2Y3c4mHOJ-s54GaO-Up6IP8WToQdcTnC59B2v3F6Ah4,38301
-databricks_sdk-0.44.0.dist-info/NOTICE,sha256=tkRcQYA1k68wDLcnOWbg2xJDsUOJw8G8DGBhb8dnI3w,1588
-databricks_sdk-0.44.0.dist-info/WHEEL,sha256=oiQVh_5PnQM0E3gPdiz09WCNmwiHDMaGer_elqB3coM,92
-databricks_sdk-0.44.0.dist-info/top_level.txt,sha256=7kRdatoSgU0EUurRQJ_3F1Nv4EOSHWAr6ng25tJOJKU,11
-databricks_sdk-0.44.0.dist-info/RECORD,,
+databricks_sdk-0.44.1.dist-info/LICENSE,sha256=afBgTZo-JsYqj4VOjnejBetMuHKcFR30YobDdpVFkqY,11411
+databricks_sdk-0.44.1.dist-info/METADATA,sha256=iCnThrfuodH8SezEZUmnw5lHzZcHnUjK-nanS8oNJh0,38301
+databricks_sdk-0.44.1.dist-info/NOTICE,sha256=tkRcQYA1k68wDLcnOWbg2xJDsUOJw8G8DGBhb8dnI3w,1588
+databricks_sdk-0.44.1.dist-info/WHEEL,sha256=oiQVh_5PnQM0E3gPdiz09WCNmwiHDMaGer_elqB3coM,92
+databricks_sdk-0.44.1.dist-info/top_level.txt,sha256=7kRdatoSgU0EUurRQJ_3F1Nv4EOSHWAr6ng25tJOJKU,11
+databricks_sdk-0.44.1.dist-info/RECORD,,