Update errors raised from functions instead of a try except approach.

Author: pvk-developer

sdv/datasets/demo.py

@@ -6,7 +6,6 @@
 import os
 import warnings
 from collections import defaultdict
-from functools import wraps
 from pathlib import Path
 from zipfile import ZipFile
 
@@ -58,6 +57,10 @@ def _get_data_from_bucket(object_key, bucket, client):
     return response['Body'].read()
 
 
+def _get_dataset_name_from_prefix(prefix):
+    return prefix.split('/')[1]
+
+
 def _list_objects(prefix, bucket, client):
     """List all objects under a given prefix using pagination.
 
@@ -79,7 +82,21 @@ def _list_objects(prefix, bucket, client):
         contents.extend(resp.get('Contents', []))
 
     if not contents:
-        raise DemoResourceNotFoundError(f"No objects found under '{prefix}' in bucket '{bucket}'.")
+        prefix_parts = prefix.split('/')
+        modality = prefix_parts[0]
+        dataset_name = _get_dataset_name_from_prefix(prefix)
+        if dataset_name:
+            raise DemoResourceNotFoundError(
+                f'Could not download dataset {dataset_name} from bucket {bucket}. '
+                'Make sure the bucket name is correct. If the bucket is private '
+                'make sure to provide your credentials.'
+            )
+        else:
+            raise DemoResourceNotFoundError(
+                f'Could not list datasets in modality {modality} from bucket {bucket}. '
+                'Make sure the bucket name is correct. If the bucket is private '
+                'make sure to provide your credentials.'
+            )
 
     return contents
 
@@ -109,7 +126,7 @@ def _search_contents_keys(contents, match_fn):
     return matches
 
 
-def _find_data_zip_key(contents, dataset_prefix):
+def _find_data_zip_key(contents, dataset_prefix, bucket):
     """Find the 'data.zip' object key under dataset prefix, case-insensitive.
 
     Args:
@@ -131,7 +148,11 @@ def is_data_zip(key):
     if matches:
         return matches[0]
 
-    raise DemoResourceNotFoundError("Could not find 'data.zip' for the requested dataset.")
+    dataset_name = _get_dataset_name_from_prefix(dataset_prefix)
+    raise DemoResourceNotFoundError(
+        f"Could not download dataset '{dataset_name}' from bucket '{bucket}'. "
+        "The dataset is missing 'data.zip' file."
+    )
 
 
 def _get_first_v1_metadata_bytes(contents, dataset_prefix, bucket, client):
@@ -167,61 +188,13 @@ def is_direct_json_under_prefix(key):
         except Exception:
             continue
 
+    dataset_name = _get_dataset_name_from_prefix(dataset_prefix)
     raise DemoResourceNotFoundError(
-        'Could not find a valid metadata JSON with METADATA_SPEC_VERSION "V1".'
+        f"Could not download dataset '{dataset_name}' from bucket '{bucket}'. "
+        'The dataset is missing a valid metadata.'
     )
 
 
-def handle_download_failure():
-    """Decorator to handle download exceptions.
-
-    Returns:
-        func:
-            A wrapped function.
-    """
-
-    def decorator(func):
-        @wraps(func)
-        def wrapper(*args, **kwargs):
-            try:
-                function_result = func(*args, **kwargs)
-            except DemoResourceNotFoundError as error:
-                bucket = kwargs.get('bucket', args[2])
-                dataset_name = kwargs.get('dataset_name', args[1])
-                error_msg = error.message
-                if 'No objects found under' in error_msg:
-                    raise DemoResourceNotFoundError(
-                        f'Could not download dataset {dataset_name} from bucket {bucket}. '
-                        'Make sure the bucket name is correct. If the bucket is private '
-                        'make sure to provide your credentials.'
-                    ) from error
-
-                if 'METADATA_SPEC_VERSION' in error_msg or 'metadata' in error_msg:
-                    raise DemoResourceNotFoundError(
-                        f'Could not download dataset {dataset_name} from bucket {bucket}. '
-                        'The dataset is missing a valid metadata.'
-                    ) from error
-
-                if 'no csv files were found in data.zip' in error_msg:
-                    raise DemoResourceNotFoundError(
-                        f'Could not download dataset {dataset_name} from bucket {bucket}. '
-                        'The dataset is missing `csv` file/s.'
-                    ) from error
-
-                if "Could not find 'data.zip'" in error_msg:
-                    raise DemoResourceNotFoundError(
-                        f'Could not download dataset {dataset_name} from bucket {bucket}. '
-                        "The dataset is missing 'data.zip' file."
-                    ) from error
-
-            return function_result
-
-        return wrapper
-
-    return decorator
-
-
-@handle_download_failure()
 def _download(modality, dataset_name, bucket, credentials=None):
     """Download dataset resources from a bucket.
 
@@ -237,7 +210,7 @@ def _download(modality, dataset_name, bucket, credentials=None):
         f'{bucket_url}/{dataset_prefix}'
     )
     contents = _list_objects(dataset_prefix, bucket=bucket, client=client)
-    zip_key = _find_data_zip_key(contents, dataset_prefix)
+    zip_key = _find_data_zip_key(contents, dataset_prefix, bucket)
     zip_bytes = _get_data_from_bucket(zip_key, bucket=bucket, client=client)
     metadata_bytes = _get_first_v1_metadata_bytes(
         contents, dataset_prefix, bucket=bucket, client=client
@@ -313,7 +286,7 @@ def _get_data_without_output_folder(in_memory_directory):
     return data, skipped_files
 
 
-def _get_data(modality, output_folder_name, in_memory_directory):
+def _get_data(modality, output_folder_name, in_memory_directory, bucket, dataset_name):
     if output_folder_name:
         data, skipped_files = _get_data_with_output_folder(output_folder_name)
     else:
@@ -324,7 +297,8 @@ def _get_data(modality, output_folder_name, in_memory_directory):
 
     if not data:
         raise DemoResourceNotFoundError(
-            'Demo data could not be downloaded because no csv files were found in data.zip'
+            f"Could not download dataset '{dataset_name}' from bucket '{bucket}'. "
+            'The dataset is missing `csv` file/s.'
         )
 
     if modality != 'multi_table':
@@ -352,7 +326,10 @@ def _get_metadata(metadata_bytes, dataset_name, output_folder_name=None):
         metadict = json.loads(metadata_bytes)
         metadata = Metadata().load_from_dict(metadict, dataset_name)
     except Exception as e:
-        raise DemoResourceNotFoundError('Failed to parse metadata JSON for the dataset.') from e
+        raise DemoResourceNotFoundError(
+            f"Could not parse the metadata for dataset '{dataset_name}'. "
+            'The dataset is missing a valid metadata file.'
+        ) from e
 
     if output_folder_name:
         try:
@@ -413,7 +390,13 @@ def download_demo(
 
     data_io, metadata_bytes = _download(modality, dataset_name, s3_bucket_name, credentials)
     in_memory_directory = _extract_data(data_io, output_folder_name)
-    data = _get_data(modality, output_folder_name, in_memory_directory)
+    data = _get_data(
+        modality,
+        output_folder_name,
+        in_memory_directory,
+        s3_bucket_name,
+        dataset_name,
+    )
     metadata = _get_metadata(metadata_bytes, dataset_name, output_folder_name)
 
     return data, metadata

tests/integration/datasets/test_demo.py

@@ -1,6 +1,8 @@
 import pandas as pd
+import pytest
 
-from sdv.datasets.demo import get_available_demos
+from sdv.datasets.demo import download_demo, get_available_demos
+from sdv.metadata import Metadata
 
 
 def test_get_available_demos_single_table():
@@ -79,3 +81,40 @@ def test_get_available_demos_multi_table():
     assert not tables_info.empty
     assert (tables_info['num_tables'] > 1).all()
     assert (tables_info['size_MB'] >= 0).all()
+
+
+@pytest.mark.parametrize('output_path', [None, 'tmp_path'])
+def test_download_demo_single_table(output_path, tmp_path):
+    """Test that the `download_demo` function works as intended for single-table."""
+    # Run
+    output_folder_name = tmp_path / 'sdv' if output_path else None
+    data, metadata = download_demo(
+        modality='single_table',
+        dataset_name='fake_hotel_guests',
+        output_folder_name=output_folder_name,
+    )
+
+    # Assert
+    metadata.validate_data({'fake_hotel_guests': data})
+    assert len(data) > 1
+    assert isinstance(metadata, Metadata)
+
+
+@pytest.mark.parametrize('output_path', [None, 'tmp_path'])
+def test_download_demo_multi_table(output_path, tmp_path):
+    """Test that the `download_demo` function works as intended for multi-table."""
+    # Run
+    output_folder_name = tmp_path / 'sdv' if output_path else None
+    data, metadata = download_demo(
+        modality='multi_table',
+        dataset_name='fake_hotels',
+        output_folder_name=output_folder_name,
+    )
+
+    # Assert
+    metadata.validate_data(data)
+    expected_tables = ['hotels', 'guests']
+    assert set(expected_tables) == set(data)
+    assert isinstance(metadata, Metadata)
+    assert len(data['hotels']) > 1
+    assert len(data['guests']) > 1

tests/unit/datasets/test_demo.py

@@ -355,7 +355,7 @@ def test__find_data_zip_key():
     dataset_prefix = 'single_table/fake_hotel_guests/'
 
     # Run
-    zip_key = _find_data_zip_key(contents, dataset_prefix)
+    zip_key = _find_data_zip_key(contents, dataset_prefix, 'bucket')
 
     # Assert
     assert zip_key == 'single_table/fake_hotel_guests/data.ZIP'
@@ -619,7 +619,7 @@ def test_download_demo_missing_zip_raises(mock_list):
 
     # Run and Assert
     expected_msg = (
-        'Could not download dataset word from bucket sdv-datasets-public. '
+        "Could not download dataset 'word' from bucket 'sdv-datasets-public'. "
         "The dataset is missing 'data.zip' file."
     )
     with pytest.raises(DemoResourceNotFoundError, match=expected_msg):
@@ -640,7 +640,7 @@ def test_download_demo_no_v1_metadata_raises(mock_list, mock_get):
 
     # Run and Assert
     error_msg = (
-        'Could not download dataset word from bucket sdv-datasets-public. '
+        "Could not download dataset 'word' from bucket 'sdv-datasets-public'. "
         'The dataset is missing a valid metadata.'
     )
     with pytest.raises(DemoResourceNotFoundError, match=error_msg):
@@ -677,8 +677,11 @@ def test__get_metadata_warns_on_save_error(_mock_open, tmp_path):
 def test__get_metadata_raises_on_invalid_json():
     """_get_metadata should raise a helpful error when JSON is invalid."""
     # Run / Assert
-    err = 'Failed to parse metadata JSON for the dataset.'
-    with pytest.raises(DemoResourceNotFoundError, match=re.escape(err)):
+    error_msg = (
+        "Could not parse the metadata for dataset 'dataset1'. "
+        'The dataset is missing a valid metadata file.'
+    )
+    with pytest.raises(DemoResourceNotFoundError, match=error_msg):
         _get_metadata(b'not-json', 'dataset1')
 
 
@@ -1105,8 +1108,11 @@ def test_download_demo_raises_when_no_csv_in_zip_single_table(mock_list, mock_ge
     )
 
     # Run and Assert
-    msg = 'Demo data could not be downloaded because no csv files were found in data.zip'
-    with pytest.raises(DemoResourceNotFoundError, match=re.escape(msg)):
+    error_msg = (
+        "Could not download dataset 'word' from bucket 'sdv-datasets-public'. "
+        'The dataset is missing `csv` file/s.'
+    )
+    with pytest.raises(DemoResourceNotFoundError, match=error_msg):
         download_demo('single_table', 'word')