Fixing and refactoring parse_dataset_uri()

src/datachain/catalog/catalog.py

@@ -681,10 +681,7 @@ def _row_to_node(d: dict[str, Any]) -> Node:
         for src in sources:  # Opt: parallel
             listing: Optional[Listing]
             if src.startswith("ds://"):
-                ds_name, ds_version = parse_dataset_uri(src)
-                ds_namespace, ds_project, ds_name = parse_dataset_name(ds_name)
-                assert ds_namespace
-                assert ds_project
+                (ds_namespace, ds_project, ds_name, ds_version) = parse_dataset_uri(src)
                 dataset = self.get_dataset(
                     ds_name, namespace_name=ds_namespace, project_name=ds_project
                 )
@@ -1515,13 +1512,12 @@ def _instantiate(ds_uri: str) -> None:
         studio_client = StudioClient()
 
         try:
-            remote_ds_name, version = parse_dataset_uri(remote_ds_uri)
+            (remote_namespace, remote_project, remote_ds_name, version) = (
+                parse_dataset_uri(remote_ds_uri)
+            )
         except Exception as e:
             raise DataChainError("Error when parsing dataset uri") from e
 
-        remote_namespace, remote_project, remote_ds_name = parse_dataset_name(
-            remote_ds_name
-        )
         if not remote_namespace or not remote_project:
             raise DataChainError(
                 f"Invalid fully qualified dataset name {remote_ds_name}, namespace"

src/datachain/dataset.py

@@ -1,5 +1,6 @@
 import builtins
 import json
+import re
 from dataclasses import dataclass, fields
 from datetime import datetime
 from functools import cached_property
@@ -10,7 +11,6 @@
     TypeVar,
     Union,
 )
-from urllib.parse import urlparse
 
 from packaging.specifiers import SpecifierSet
 from packaging.version import Version
@@ -43,25 +43,58 @@
 StorageURI = NewType("StorageURI", str)
 
 
-def parse_dataset_uri(uri: str) -> tuple[str, Optional[str]]:
+def parse_dataset_uri(uri: str) -> tuple[str, str, str, str | None]:
     """
-    Parse dataser uri to extract name and version out of it (if version is defined)
-    Example:
-        Input: ds://[email protected]
-        Output: (zalando, 3.0.1)
+    Parse a dataset URI of the form:
+
+        ds://<namespace>.<project>.<name>[@v<semver>]
+
+    Components:
+    - `ds://`        : required prefix identifying dataset URIs.
+    - `namespace`    : required namespace, may start with '@' (e.g., "@user").
+    - `project`      : required project name inside the namespace.
+    - `name`         : required dataset name.
+    - `@v<semver>`   : optional version suffix. Must start with '@v' and
+                       be a semantic version string MAJOR.MINOR.PATCH
+                       (e.g., "1.0.4").
+
+    Returns:
+        tuple[str, str, str, str | None]:
+            (namespace, project, name, version) where version is None
+            if not provided.
+
+    Raises:
+        ValueError: if the URI does not start with 'ds://' or does not
+                    match the expected format.
     """
-    p = urlparse(uri)
-    if p.scheme != "ds":
-        raise Exception("Dataset uri should start with ds://")
-    s = p.netloc.split("@v")
-    name = s[0]
-    if len(s) == 1:
-        return name, None
-    if len(s) != 2:
-        raise Exception(
-            "Wrong dataset uri format, it should be: ds://<name>@v<version>"
-        )
-    return name, s[1]
+
+    if not uri.startswith("ds://"):
+        raise ValueError(f"Invalid dataset URI: {uri}")
+
+    body = uri[len("ds://") :]
+
+    pattern = re.compile(
+        r"""
+        ^(?P<namespace>@?\w+)      # namespace, may start with '@'
+        \. (?P<project>\w+)        # project
+        \. (?P<name>\w+)           # dataset name
+        (?:@v                      # optional version prefix must be '@v'
+            (?P<version>\d+\.\d+\.\d+)
+        )?$                        # end of string
+        """,
+        re.VERBOSE,
+    )
+
+    match = pattern.match(body)
+    if not match:
+        raise ValueError(f"Invalid dataset URI format: {uri}")
+
+    return (
+        match.group("namespace"),
+        match.group("project"),
+        match.group("name"),
+        match.group("version"),
+    )
 
 
 def create_dataset_uri(

tests/func/test_pull.py

@@ -286,9 +286,9 @@ def test_pull_dataset_wrong_version(
 
     with pytest.raises(DataChainError) as exc_info:
         catalog.pull_dataset(
-            f"ds://{REMOTE_NAMESPACE_NAME}.{REMOTE_PROJECT_NAME}.dogs@v5"
+            f"ds://{REMOTE_NAMESPACE_NAME}.{REMOTE_PROJECT_NAME}.dogs@v5.0.0"
         )
-    assert str(exc_info.value) == "Dataset dogs doesn't have version 5 on server"
+    assert str(exc_info.value) == "Dataset dogs doesn't have version 5.0.0 on server"
 
 
 @pytest.mark.parametrize("cloud_type, version_aware", [("s3", False)], indirect=True)

tests/unit/test_dataset.py

@@ -13,6 +13,7 @@
     DatasetRecord,
     DatasetVersion,
     parse_dataset_name,
+    parse_dataset_uri,
 )
 from datachain.error import InvalidDatasetNameError
 from datachain.sql.types import (
@@ -180,3 +181,34 @@ def test_parse_dataset_name(full_name, namespace, project, name):
 def test_parse_dataset_name_empty_name():
     with pytest.raises(InvalidDatasetNameError):
         assert parse_dataset_name(None)
+
+
+@pytest.mark.parametrize(
+    "uri,namespace,project,name,version",
+    [
+        ("ds://global.dev.result", "global", "dev", "result", None),
+        ("ds://[email protected]", "global", "dev", "result", "1.0.5"),
+        ("ds://@ilongin.dev.result", "@ilongin", "dev", "result", None),
+        ("ds://@[email protected]", "@ilongin", "dev", "result", "1.0.4"),
+        ("ds://@vlad.dev.result", "@vlad", "dev", "result", None),
+        ("ds://@[email protected]", "@vlad", "dev", "result", "1.0.5"),
+    ],
+)
+def test_parse_dataset_uri(uri, namespace, project, name, version):
+    assert parse_dataset_uri(uri) == (namespace, project, name, version)
+
+
+@pytest.mark.parametrize(
+    "uri",
+    [
+        "ds://result",
+        "ds://[email protected]",
+        "ds://@[email protected]",
+        "ds://@[email protected]",
+    ],
+)
+def test_parse_dataset_uri_invalid_format(uri):
+    with pytest.raises(ValueError) as excinfo:
+        parse_dataset_uri(uri)
+
+    assert str(excinfo.value) == f"Invalid dataset URI format: {uri}"