chore: import upstream snapshot with attribution

mlflow/utils/workspace_utils.py

@@ -0,0 +1,118 @@
+from __future__ import annotations
+
+from mlflow.environment_variables import MLFLOW_WORKSPACE, MLFLOW_WORKSPACE_STORE_URI
+
+_workspace_store_uri: str | None = None
+
+DEFAULT_WORKSPACE_NAME = "default"
+WORKSPACES_DIR_NAME = "workspaces"
+WORKSPACE_HEADER_NAME = "X-MLFLOW-WORKSPACE"
+
+
+def _normalize_workspace(workspace: str | None) -> str | None:
+    """Normalize a workspace identifier.
+
+    Args:
+        workspace: Raw workspace value, possibly ``None`` or whitespace padded.
+
+    Returns:
+        The trimmed workspace string or ``None`` when the input is empty.
+    """
+    if workspace is None:
+        return None
+    value = workspace.strip()
+    return value or None
+
+
+def resolve_entity_workspace_name(workspace: str | None) -> str:
+    """Determine the workspace to associate with client-facing entities.
+
+    Preference order:
+        1. Explicit ``workspace`` argument provided by the backend store.
+        2. Active workspace bound via ``mlflow.set_workspace``.
+        3. ``MLFLOW_WORKSPACE`` environment variable.
+        4. ``DEFAULT_WORKSPACE_NAME``.
+
+    Args:
+        workspace: Optional workspace name provided by the store layer.
+
+    Returns:
+        A normalized workspace name honoring the preference order.
+    """
+
+    if candidate := _normalize_workspace(workspace):
+        return candidate
+
+    from mlflow.utils.workspace_context import get_request_workspace
+
+    if candidate := _normalize_workspace(get_request_workspace()):
+        return candidate
+
+    if candidate := _normalize_workspace(MLFLOW_WORKSPACE.get()):
+        return candidate
+
+    return DEFAULT_WORKSPACE_NAME
+
+
+def set_workspace_store_uri(uri: str | None) -> None:
+    """Set the global workspace provider URI override.
+
+    Args:
+        uri: URI of the workspace provider or ``None`` to clear the override.
+    """
+
+    global _workspace_store_uri
+    _workspace_store_uri = uri
+    if uri is None:
+        MLFLOW_WORKSPACE_STORE_URI.unset()
+    else:
+        MLFLOW_WORKSPACE_STORE_URI.set(uri)
+
+
+def resolve_workspace_store_uri(
+    workspace_store_uri: str | None = None, tracking_uri: str | None = None
+) -> str | None:
+    """Resolve the workspace provider URI according to precedence rules.
+
+    Args:
+        workspace_store_uri: URI provided explicitly by the caller.
+        tracking_uri: Tracking URI used as the final fallback.
+
+    Returns:
+        The workspace provider URI chosen in this order:
+
+        1. Explicit ``workspace_store_uri`` argument.
+        2. Value configured via :func:`set_workspace_store_uri`
+           or ``MLFLOW_WORKSPACE_STORE_URI``.
+        3. The resolved tracking URI.
+    """
+
+    if workspace_store_uri is not None:
+        return workspace_store_uri
+
+    if configured_uri := get_workspace_store_uri():
+        return configured_uri
+
+    # Lazy import to avoid circular dependency during module import.
+    from mlflow.tracking._tracking_service import utils as tracking_utils
+
+    return tracking_utils._resolve_tracking_uri(tracking_uri)
+
+
+def get_workspace_store_uri() -> str | None:
+    """Get the current workspace provider URI override, if any.
+
+    Returns:
+        The globally configured workspace provider URI, or ``None``.
+    """
+    return _workspace_store_uri or MLFLOW_WORKSPACE_STORE_URI.get()
+
+
+__all__ = [
+    "DEFAULT_WORKSPACE_NAME",
+    "WORKSPACES_DIR_NAME",
+    "WORKSPACE_HEADER_NAME",
+    "resolve_entity_workspace_name",
+    "set_workspace_store_uri",
+    "get_workspace_store_uri",
+]