1  import urllib.parse
 2  
 3  from mlflow.environment_variables import MLFLOW_UC_OSS_TOKEN
 4  from mlflow.exceptions import MlflowException
 5  from mlflow.utils.databricks_utils import get_databricks_host_creds
 6  from mlflow.utils.rest_utils import MlflowHostCreds
 7  from mlflow.utils.uri import (
 8      _DATABRICKS_UNITY_CATALOG_SCHEME,
 9  )
10  
11  
12  def get_oss_host_creds(server_uri=None):
13      """
14      Retrieve the host credentials for the OSS server.
15  
16      Args:
17          server_uri (str): The URI of the server.
18  
19      Returns:
20          MlflowHostCreds: The host credentials for the OSS server.
21      """
22      parsed_uri = urllib.parse.urlparse(server_uri)
23  
24      if parsed_uri.scheme != "uc":
25          raise MlflowException("The scheme of the server_uri should be 'uc'")
26  
27      if parsed_uri.path == _DATABRICKS_UNITY_CATALOG_SCHEME:
28          return get_databricks_host_creds(parsed_uri.path)
29      return MlflowHostCreds(host=parsed_uri.path, token=MLFLOW_UC_OSS_TOKEN.get())