Merge branch 'databricks:main' into main

Author: pivoshenko

.codegen.json

@@ -1,6 +1,6 @@
 {
   "mode": "py_v0",
-  "changelog_config": ".codegen/changelog_config.yml",
+  "api_changelog": true,  
   "version": {
     "databricks/sdk/version.py": "__version__ = '$VERSION'"
   },

.codegen/changelog_config.yml

@@ -7,8 +7,5 @@ change_types:
     tag: "[Doc]"
   - message: Internal Changes
     tag: "[Internal]"
-  # Does not appear in the Changelog. Only for PR validation.
-  - message: Release
-    tag: "[Release]"
   # Default for messages without a tag
   - message: Other Changes

.github/workflows/next-changelog.yml

@@ -0,0 +1,105 @@
+# Generated file. DO NOT EDIT.
+name: Check for NEXT_CHANGELOG.md Changes
+
+on:
+  # Use pull_request_target to have access to GitHub API
+  pull_request_target:
+
+jobs:
+  check-next-changelog:
+    runs-on:
+      group: databricks-deco-testing-runner-group
+      labels: ubuntu-latest-deco
+
+    steps:
+      - name: Checkout code
+        uses: actions/checkout@v3
+
+      - name: Fetch list of changed files
+        id: changed-files
+        env:
+          GH_TOKEN: ${{ secrets.GITHUB_TOKEN }}
+        run: |
+          # Use the GitHub API to fetch changed files
+          files=$(gh pr view ${{ github.event.pull_request.number }} --json files -q '.files[].path')
+          
+          # Sanitize to avoid code injection
+          sanitized_files=$(echo "$files" | sed 's/[^a-zA-Z0-9._/-]/_/g')
+          
+          # Store the sanitized list of files in a temporary file to avoid env variable issues
+          echo "$sanitized_files" > modified_files.txt
+
+      - name: Fetch PR message
+        id: pr-message
+        env:
+          GH_TOKEN: ${{ secrets.GITHUB_TOKEN }}
+        run: |
+          # Use the GitHub API to fetch the PR message
+          pr_message=$(gh pr view ${{ github.event.pull_request.number }} --json body -q '.body')
+          
+          # Sanitize the PR message to avoid code injection, keeping the equal sign
+          sanitized_pr_message=$(echo "$pr_message" | sed 's/[^a-zA-Z0-9._/-=]/_/g')
+          
+          # Store the sanitized PR message
+          echo "$sanitized_pr_message" > pr_message.txt
+
+      - name: Verify NEXT_CHANGELOG.md was modified or PR message contains NO_CHANGELOG=true
+        run: |
+          # Read the sanitized files and PR message from the temporary files
+          modified_files=$(cat modified_files.txt)
+          pr_message=$(cat pr_message.txt)
+          
+          # Check if NEXT_CHANGELOG.md exists in the list of changed files
+          echo "Changed files: $modified_files"
+          if ! echo "$modified_files" | grep -q "NEXT_CHANGELOG.md"; then
+            echo "NEXT_CHANGELOG.md not modified."
+          
+            # Check if PR message contains NO_CHANGELOG=true
+            if echo "$pr_message" | grep -q "NO_CHANGELOG=true"; then
+              echo "NO_CHANGELOG=true found in PR message. Skipping changelog check."
+              exit 0
+            else
+              echo "WARNING: file NEXT_CHANGELOG.md not changed. If this is expected, add NO_CHANGELOG=true to the PR message."
+              exit 1
+            fi
+          fi
+
+      - name: Comment on PR with instructions if needed
+        if: failure() # This step will only run if the previous step fails (i.e., if NEXT_CHANGELOG.md was not modified and NO_CHANGELOG=true was not in the PR message)
+        env:
+          GH_TOKEN: ${{ secrets.GITHUB_TOKEN }}
+        run: |
+          # Check if a comment exists with the instructions
+          previous_comment_ids=$(gh api "repos/${{ github.repository }}/issues/${{ github.event.pull_request.number }}/comments" \
+            --jq '.[] | select(.body | startswith("<!-- NEXT_CHANGELOG_INSTRUCTIONS -->")) | .id')
+          echo "Previous comment IDs: $previous_comment_ids"
+
+          # If no previous comment exists, add one with instructions
+          if [ -z "$previous_comment_ids" ]; then
+            echo "Adding instructions comment."
+            gh pr comment ${{ github.event.pull_request.number }} --body \
+            "<!-- NEXT_CHANGELOG_INSTRUCTIONS -->
+            Please ensure that the NEXT_CHANGELOG.md file is updated with any relevant changes. 
+            If this is not necessary for your PR, please include the following in your PR description:
+            NO_CHANGELOG=true
+            and rerun the job."
+          fi
+
+      - name: Delete instructions comment on success
+        if: success() # This step will only run if the previous check passed (i.e., if NEXT_CHANGELOG.md was modified or NO_CHANGELOG=true is in the PR message)
+        env:
+          GH_TOKEN: ${{ secrets.GITHUB_TOKEN }}
+        run: |
+          # Check if there is a previous instructions comment
+          previous_comment_ids=$(gh api "repos/${{ github.repository }}/issues/${{ github.event.pull_request.number }}/comments" \
+            --jq '.[] | select(.body | startswith("<!-- NEXT_CHANGELOG_INSTRUCTIONS -->")) | .id')
+          
+            # If a comment exists, delete it
+            if [ -n "$previous_comment_ids" ]; then
+              echo "Deleting previous instructions comment."
+              for comment_id in $previous_comment_ids; do
+                gh api "repos/${{ github.repository }}/issues/comments/$comment_id" --method DELETE
+              done
+            else
+              echo "No instructions comment found to delete."
+            fi

.github/workflows/tagging.yml

@@ -0,0 +1,52 @@
+# Generated file. DO NOT EDIT.
+name: tagging
+
+on:
+  workflow_dispatch:
+  # Enable for automatic tagging
+  #schedule:
+  #  - cron: '0 0 * * TUE'
+
+# Ensure that only a single instance of the workflow is running at a time.
+concurrency:
+  group: "tagging"
+
+
+jobs:
+  tag:
+    environment: "release-is"
+    runs-on:
+      group: databricks-deco-testing-runner-group
+      labels: ubuntu-latest-deco
+    steps:
+      - name: Generate GitHub App Token
+        id: generate-token
+        uses: actions/create-github-app-token@v1
+        with:
+          app-id: ${{ secrets.DECO_SDK_TAGGING_APP_ID }}
+          private-key: ${{ secrets.DECO_SDK_TAGGING_PRIVATE_KEY }}
+
+      - name: Checkout repository
+        uses: actions/checkout@v4
+        with:
+          fetch-depth: 0
+          token: ${{ steps.generate-token.outputs.token }}
+
+      #NOTE: email must be the GitHub App email or the commit will not be verified.
+      - name: Set up Git configuration
+        run: |
+          git config user.name "Databricks SDK Release Bot"
+          git config user.email "DECO-SDK-Tagging[bot]@users.noreply.github.com"
+
+      - name: Install dependencies
+        run: |
+          python -m pip install --upgrade pip
+          pip install PyGithub
+
+      - name: Run script
+        env:
+          GITHUB_TOKEN: ${{ steps.generate-token.outputs.token }}
+          GITHUB_REPOSITORY: ${{ github.repository }}
+        run: |
+          python tagging.py
+

.package.json

@@ -0,0 +1 @@
+{}

CHANGELOG.md

@@ -1,5 +1,18 @@
 # Version changelog
 
+## [Release] Release v0.44.1
+
+### New Features and Improvements
+
+ * Introduce new Credential Strategies for Agents ([#882](https://github.com/databricks/databricks-sdk-py/pull/882)).
+
+
+### Internal Changes
+
+ * GetRun logic paginates more arrays ([#867](https://github.com/databricks/databricks-sdk-py/pull/867)).
+
+
+
 ## [Release] Release v0.44.0
 
 ### Internal Changes

NEXT_CHANGELOG.md

@@ -0,0 +1,15 @@
+# NEXT CHANGELOG
+
+## Release v0.45.0
+
+### New Features and Improvements
+
+### Bug Fixes
+
+### Documentation
+
+### Internal Changes
+* Introduce automated tagging ([#888](https://github.com/databricks/databricks-sdk-py/pull/888))
+* Update Jobs GetJob API to support paginated responses ([#869](https://github.com/databricks/databricks-sdk-py/pull/869)).
+
+### API Changes

databricks/sdk/credentials_provider.py

@@ -9,6 +9,7 @@
 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 inner() -> Dict[str, str]:
 # 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 @@ def should_fetch_model_serving_environment_oauth(self) -> bool:
         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 @@ def get_model_dependency_oauth_token(self, should_retry=True) -> str:
                 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 inner() -> Dict[str, str]:
     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 @@ def __call__(self, cfg: 'Config') -> CredentialsProvider:
         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)