Repository: mlflow/mlflow

Issue:
[BUG] uc_volume_dataset_source only validates file paths, not folder paths
### Issues Policy acknowledgement

- [X] I have read and agree to submit bug reports in accordance with the [issues policy](https://www.github.com/mlflow/mlflow/blob/master/ISSUE_POLICY.md)

### Where did you encounter this bug?

Local machine

### Willingness to contribute

Yes. I would be willing to contribute a fix for this bug with guidance from the MLflow community.

### MLflow version

mlflow-2.12.2

### System information

- **OS Platform and Distribution (e.g., Linux Ubuntu 16.04)**:
- **Python version**:
- **yarn version, if running the dev UI**:


### Describe the problem

https://github.com/mlflow/mlflow/blob/72df4a2a0f44c52179dfbdc7d47ad10f58ceec39/mlflow/data/uc_volume_dataset_source.py#L28 doesn't verify folder paths, only file paths

### Tracking information

<!-- PLEASE KEEP BACKTICKS AND CHECK PREVIEW -->
```shell
REPLACE_ME
```


### Code to reproduce issue

<!-- PLEASE KEEP BACKTICKS AND CHECK PREVIEW -->
```
REPLACE_ME
```


### Stack trace

<!-- PLEASE KEEP BACKTICKS AND CHECK PREVIEW -->
```
REPLACE_ME
```


### Other info / logs

<!-- PLEASE KEEP BACKTICKS AND CHECK PREVIEW -->
```
REPLACE_ME
```


### What component(s) does this bug affect?

- [ ] `area/artifacts`: Artifact stores and artifact logging
- [ ] `area/build`: Build and test infrastructure for MLflow
- [ ] `area/deployments`: MLflow Deployments client APIs, server, and third-party Deployments integrations
- [ ] `area/docs`: MLflow documentation pages
- [ ] `area/examples`: Example code
- [ ] `area/model-registry`: Model Registry service, APIs, and the fluent client calls for Model Registry
- [ ] `area/models`: MLmodel format, model serialization/deserialization, flavors
- [ ] `area/recipes`: Recipes, Recipe APIs, Recipe configs, Recipe Templates
- [ ] `area/projects`: MLproject format, project running backends
- [ ] `area/scoring`: MLflow Model server, model deployment tools, Spark UDFs
- [ ] `area/server-infra`: MLflow Tracking server backend
- [ ] `area/tracking`: Tracking Service, tracking client APIs, autologging

### What interface(s) does this bug affect?

- [ ] `area/uiux`: Front-end, user experience, plotting, JavaScript, JavaScript dev server
- [ ] `area/docker`: Docker use across MLflow's components, such as MLflow Projects and MLflow Models
- [ ] `area/sqlalchemy`: Use of SQLAlchemy in the Tracking Service or Model Registry
- [ ] `area/windows`: Windows support

### What language(s) does this bug affect?

- [ ] `language/r`: R APIs and clients
- [ ] `language/java`: Java APIs and clients
- [ ] `language/new`: Proposals for new client languages

### What integration(s) does this bug affect?

- [ ] `integrations/azure`: Azure and Azure ML integrations
- [ ] `integrations/sagemaker`: SageMaker integrations
- [ ] `integrations/databricks`: Databricks integrations


Discussion:

Resolving patch:
diff --git a/mlflow/data/uc_volume_dataset_source.py b/mlflow/data/uc_volume_dataset_source.py
--- a/mlflow/data/uc_volume_dataset_source.py
+++ b/mlflow/data/uc_volume_dataset_source.py
@@ -22,10 +22,10 @@ class UCVolumeDatasetSource(DatasetSource):
     """
 
     def __init__(self, path: str):
-        self._verify_uc_path_is_valid(path)
         self.path = path
+        self._verify_uc_path_is_valid()
 
-    def _verify_uc_path_is_valid(self, path):
+    def _verify_uc_path_is_valid(self):
         """Verify if the path exists in Databricks Unified Catalog."""
         try:
             from databricks.sdk import WorkspaceClient
@@ -49,9 +49,17 @@ def _verify_uc_path_is_valid(self, path):
             return
 
         try:
-            w.files.get_metadata(path)
+            # Check if `self.path` points to a valid UC file.
+            w.files.get_metadata(self.path)
         except Exception:
-            raise MlflowException(f"{path} does not exist in Databricks Unified Catalog.")
+            try:
+                # Check if `self.path` points to a valid UC directory.
+                w.files.get_directory_metadata(self.path)
+                # Append a slash to `self.path` to indicate it's a directory.
+                self.path += "/" if not self.path.endswith("/") else ""
+            except Exception:
+                # Neither file nor directory exists, we throw an exception.
+                raise MlflowException(f"{self.path} does not exist in Databricks Unified Catalog.")
 
     @staticmethod
     def _get_source_type() -> str:


Test patch: