Update Credential Strategy

Author: aravind-segu

databricks/sdk/credentials_provider.py

@@ -314,12 +314,11 @@ def github_oidc_azure(cfg: 'Config') -> Optional[CredentialsProvider]:
         # detect Azure AD Tenant ID if it's not specified directly
         token_endpoint = cfg.oidc_endpoints.token_endpoint
         cfg.azure_tenant_id = token_endpoint.replace(aad_endpoint, '').split('/')[0]
-    inner = ClientCredentials(
-        client_id=cfg.azure_client_id,
-        client_secret="", # we have no (rotatable) secrets in OIDC flow
-        token_url=f"{aad_endpoint}{cfg.azure_tenant_id}/oauth2/token",
-        endpoint_params=params,
-        use_params=True)
+    inner = ClientCredentials(client_id=cfg.azure_client_id,
+                              client_secret="", # we have no (rotatable) secrets in OIDC flow
+                              token_url=f"{aad_endpoint}{cfg.azure_tenant_id}/oauth2/token",
+                              endpoint_params=params,
+                              use_params=True)
 
     def refreshed_headers() -> Dict[str, str]:
         token = inner.token()
@@ -725,11 +724,10 @@ def inner() -> Dict[str, str]:
 # https://github.com/mlflow/mlflow/blob/1219e3ef1aac7d337a618a352cd859b336cf5c81/mlflow/legacy_databricks_cli/configure/provider.py#L332
 class ModelServingAuthProvider():
     USER_CREDENTIALS = "user_credentials"
-    EMBEDDED_CREDENTIALS = "embedded_credentials"
 
     _MODEL_DEPENDENCY_OAUTH_TOKEN_FILE_PATH = "/var/credentials-secret/model-dependencies-oauth-token"
 
-    def __init__(self, credential_type):
+    def __init__(self, credential_type: Optional[str]):
         self.expiry_time = -1
         self.current_token = None
         self.refresh_duration = 300 # 300 Seconds
@@ -746,7 +744,7 @@ def should_fetch_model_serving_environment_oauth() -> bool:
         return (is_in_model_serving_env == "true"
                 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
@@ -762,14 +760,14 @@ 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):
+    def _get_invokers_token(self):
         current_thread = threading.current_thread()
         thread_data = current_thread.__dict__
         invokers_token = None
@@ -788,18 +786,16 @@ def get_databricks_host_token(self) -> Optional[Tuple[str, str]]:
         # 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(
-        ) if self.credential_type == ModelServingAuthProvider.EMBEDDED_CREDENTIALS else self.get_invokers_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())
 
 
-def model_serving_auth_func(cfg: 'Config', credential_type) -> Optional[CredentialsProvider]:
+def model_serving_auth_visitor(cfg: 'Config',
+                               credential_type: Optional[str] = None) -> Optional[CredentialsProvider]:
     try:
-        if not ModelServingAuthProvider.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:
@@ -823,7 +819,11 @@ def inner() -> Dict[str, str]:
 
 @credentials_strategy('model-serving', [])
 def model_serving_auth(cfg: 'Config') -> Optional[CredentialsProvider]:
-    return model_serving_auth_func(cfg, ModelServingAuthProvider.EMBEDDED_CREDENTIALS)
+    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:
@@ -870,37 +870,25 @@ def __call__(self, cfg: 'Config') -> CredentialsProvider:
         )
 
 
-class AgentCredentials(CredentialsStrategy):
+class ModelServingUserCredentials(CredentialsStrategy):
 
-    def __init__(self, credential_type):
-        self.credential_type = credential_type
+    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 "agent_" + self.credential_type
+            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_func(cfg, self.credential_type)
+            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)
-
-
-class AgentUserCredentials(AgentCredentials):
-
-    def __init__(self):
-        super().__init__(ModelServingAuthProvider.USER_CREDENTIALS)
-
-
-class AgentEmbeddedCredentials(AgentCredentials):
-
-    def __init__(self):
-        super().__init__(ModelServingAuthProvider.EMBEDDED_CREDENTIALS)

tests/test_model_serving_auth.py

@@ -4,8 +4,7 @@
 import pytest
 
 from databricks.sdk.core import Config
-from databricks.sdk.credentials_provider import (AgentEmbeddedCredentials,
-                                                 AgentUserCredentials)
+from databricks.sdk.credentials_provider import ModelServingUserCredentials
 
 from .conftest import raises
 
@@ -27,9 +26,7 @@
                           ([('IS_IN_DATABRICKS_MODEL_SERVING_ENV', 'true'),
                             ('DATABRICKS_MODEL_SERVING_HOST_URL', 'x')
                             ], ['DB_MODEL_SERVING_HOST_URL'], "tests/testdata/model-serving-test-token"), ])
-@pytest.mark.parametrize("use_credential_strategy", [True, False])
-def test_model_serving_auth(env_values, del_env_values, oauth_file_name, use_credential_strategy, monkeypatch,
-                            mocker):
+def test_model_serving_auth(env_values, del_env_values, oauth_file_name, monkeypatch, mocker):
     ## In mlflow we check for these two environment variables to return the correct config
     for (env_name, env_value) in env_values:
         monkeypatch.setenv(env_name, env_value)
@@ -42,12 +39,9 @@ def test_model_serving_auth(env_values, del_env_values, oauth_file_name, use_cre
         "databricks.sdk.credentials_provider.ModelServingAuthProvider._MODEL_DEPENDENCY_OAUTH_TOKEN_FILE_PATH",
         oauth_file_name)
     mocker.patch('databricks.sdk.config.Config._known_file_config_loader')
-    if use_credential_strategy:
-        cfg = Config(credentials_strategy=AgentEmbeddedCredentials())
-        assert cfg.auth_type == 'agent_embedded_credentials'
-    else:
-        cfg = Config()
-        assert cfg.auth_type == 'model-serving'
+
+    cfg = Config()
+    assert cfg.auth_type == 'model-serving'
     headers = cfg.authenticate()
     assert (cfg.host == 'x')
     # Token defined in the test file
@@ -78,8 +72,7 @@ def test_model_serving_auth_errors(env_values, oauth_file_name, monkeypatch):
     Config()
 
 
-@pytest.mark.parametrize("use_credential_strategy", [True, False])
-def test_model_serving_auth_refresh(use_credential_strategy, monkeypatch, mocker):
+def test_model_serving_auth_refresh(monkeypatch, mocker):
     ## In mlflow we check for these two environment variables to return the correct config
     monkeypatch.setenv('IS_IN_DB_MODEL_SERVING_ENV', 'true')
     monkeypatch.setenv('DB_MODEL_SERVING_HOST_URL', 'x')
@@ -90,12 +83,8 @@ def test_model_serving_auth_refresh(use_credential_strategy, monkeypatch, mocker
         "tests/testdata/model-serving-test-token")
     mocker.patch('databricks.sdk.config.Config._known_file_config_loader')
 
-    if use_credential_strategy:
-        cfg = Config(credentials_strategy=AgentEmbeddedCredentials())
-        assert cfg.auth_type == 'agent_embedded_credentials'
-    else:
-        cfg = Config()
-        assert cfg.auth_type == 'model-serving'
+    cfg = Config()
+    assert cfg.auth_type == 'model-serving'
 
     current_time = time.time()
     headers = cfg.authenticate()
@@ -135,8 +124,8 @@ def test_agent_user_credentials(monkeypatch, mocker):
     thread_data = current_thread.__dict__
     thread_data["invokers_token"] = invokers_token_val
 
-    cfg = Config(credentials_strategy=AgentUserCredentials())
-    assert cfg.auth_type == 'agent_user_credentials'
+    cfg = Config(credentials_strategy=ModelServingUserCredentials())
+    assert cfg.auth_type == 'model_serving_user_credentials'
 
     headers = cfg.authenticate()
 
@@ -160,22 +149,7 @@ def test_agent_user_credentials_in_non_model_serving_environments(monkeypatch):
     monkeypatch.setenv('DATABRICKS_HOST', 'x')
     monkeypatch.setenv('DATABRICKS_TOKEN', 'token')
 
-    cfg = Config(credentials_strategy=AgentUserCredentials())
-    assert cfg.auth_type == 'pat' # Auth type is PAT as it is no longer in a model serving environment
-
-    headers = cfg.authenticate()
-
-    assert (cfg.host == 'https://x')
-    assert headers.get("Authorization") == f'Bearer token'
-
-
-# If this credential strategy is being used in a non model serving environments then use default credential strategy instead
-def test_agent_embedded_credentials_in_non_model_serving_environments(monkeypatch):
-
-    monkeypatch.setenv('DATABRICKS_HOST', 'x')
-    monkeypatch.setenv('DATABRICKS_TOKEN', 'token')
-
-    cfg = Config(credentials_strategy=AgentEmbeddedCredentials())
+    cfg = Config(credentials_strategy=ModelServingUserCredentials())
     assert cfg.auth_type == 'pat' # Auth type is PAT as it is no longer in a model serving environment
 
     headers = cfg.authenticate()