Merge pull request #327 from DagsHub/create-datasource-url-fix

Fix URL checking of create_datasource
DagsHub · Jul 23, 2023 · 487cbb6 · 487cbb6
2 parents b6fa8f8 + ca89558
commit 487cbb6
Show file tree

Hide file tree

Showing 3 changed files with 7 additions and 20 deletions.
diff --git a/dagshub/data_engine/datasources.py b/dagshub/data_engine/datasources.py
@@ -1,13 +1,11 @@
 import logging
-import urllib.parse
 from typing import Optional, Union, List
 
 from dagshub.common.analytics import send_analytics_event
 from dagshub.common.api.repo import RepoAPI
 from dagshub.data_engine.client.data_client import DataClient
-from dagshub.data_engine.client.models import DatasourceType
 from dagshub.data_engine.model.datasource import Datasource
-from dagshub.data_engine.model.datasource_state import DatasourceState
+from dagshub.data_engine.model.datasource_state import DatasourceState, DatasourceType, path_regexes
 from dagshub.data_engine.model.errors import DatasourceNotFoundError
 
 logger = logging.getLogger(__name__)
@@ -28,24 +26,12 @@ def create_datasource(repo: str, name: str, path: str, revision: Optional[str] =
     :return The created datasource
     """
 
-    parsed = urllib.parse.urlparse(path)
-
-    if parsed.scheme == "" \
-        and parsed.hostname == "" \
-        and parsed.query == "" \
-        and parsed.params == "" \
-        and parsed.fragment == "":
-        return create_from_repo(repo, name, path=parsed.path, revision=revision)
-
-    elif parsed.scheme != "" and parsed.hostname != "":
-        # Bucket URL
-        if parsed.query != "" or parsed.params != "" or parsed.fragment != "":
-            raise ValueError("Invalid bucket URL: ", path)
-        if revision != "":
+    if path_regexes[DatasourceType.BUCKET].fullmatch(path):
+        if revision is not None:
             raise ValueError("revision cannot be used together with bucket URLs")
         return create_from_bucket(repo, name, bucket_url=path)
-
-    raise ValueError("Invalid path used, format could not be determined: ", path)
+    else:
+        return create_from_repo(repo, name, path=path, revision=revision)
 
 
 create = create_datasource

diff --git a/dagshub/data_engine/model/datasource_state.py b/dagshub/data_engine/model/datasource_state.py
@@ -19,7 +19,7 @@
 logger = logging.getLogger(__name__)
 
 path_regexes = {
-    DatasourceType.BUCKET: re.compile(r"(?P<schema>s3|gs)://(?P<bucket>[\w\-]+)(?P<prefix>/.*)?"),
+    DatasourceType.BUCKET: re.compile(r"(?P<schema>s3|gs)://(?P<bucket>[\w\-._]+)(?P<prefix>/.*)?"),
     DatasourceType.REPOSITORY: re.compile(r"repo://(?P<user>[\w\-_.]+)/(?P<repo>[\w\-_.]+)(?P<prefix>/.*)?"),
 }
 

diff --git a/tests/data_engine/test_datasource_state.py b/tests/data_engine/test_datasource_state.py
@@ -64,6 +64,7 @@ def test_repo_regex_incorrect(in_str):
         ("s3://bucket", "s3", "bucket", None),
         ("gs://bucket/prefix", "gs", "bucket", "/prefix"),
         ("s3://bucket/longer/prefix", "s3", "bucket", "/longer/prefix"),
+        ("s3://bucket_with.weird-chars/longer/prefix", "s3", "bucket_with.weird-chars", "/longer/prefix"),
     ]
 )
 def test_bucket_regex(in_str, schema, bucket, prefix):