Merge pull request #2 from MITLibraries/TIMX-415-load-dataset

ghukill · web-flow · commit 9f9477afb2a0 · 2024-12-04T14:36:02.000-05:00
TIMX 415 - load dataset
diff --git a/Pipfile b/Pipfile
@@ -21,6 +21,8 @@ pytest = "*"
 ruff = "*"
 setuptools = "*"
 pandas-stubs = "*"
+moto = "*"
+pytest-mock = "*"
 
 [requires]
 python_version = "3.12"
diff --git a/Pipfile.lock b/Pipfile.lock
diff --git a/README.md b/README.md
@@ -26,6 +26,15 @@ timdex_dataset_api = {git = "https://github.com/MITLibraries/timdex-dataset-api.
 ... other dependencies...
 ```
 
+## Environment Variables
+
+### Required
+
+### Optional
+```shell
+TDA_LOG_LEVEL=# log level for timdex-dataset-api, accepts [DEBUG, INFO, WARNING, ERROR], default INFO
+```
+
 ## Usage
 
 _TODO..._
diff --git a/pyproject.toml b/pyproject.toml
@@ -94,6 +94,7 @@ ignore = [
     "S320",
     "S321",
     "S608",
+    "TRY003"
 ]
 
 fixable = ["E", "F", "I", "Q"]
diff --git a/tests/conftest.py b/tests/conftest.py
@@ -0,0 +1,24 @@
+"""tests/conftest.py"""
+
+import pytest
+
+from timdex_dataset_api import TIMDEXDataset
+
+
+@pytest.fixture(autouse=True)
+def _test_env(monkeypatch):
+    monkeypatch.setenv("TDA_LOG_LEVEL", "INFO")
+    monkeypatch.setenv("AWS_ACCESS_KEY_ID", "fake_access_key")
+    monkeypatch.setenv("AWS_SECRET_ACCESS_KEY", "fake_secret_key")
+    monkeypatch.setenv("AWS_SESSION_TOKEN", "fake_session_token")
+    monkeypatch.setenv("AWS_DEFAULT_REGION", "us-east-1")
+
+
+@pytest.fixture
+def local_dataset_location():
+    return "tests/fixtures/local_datasets/dataset"
+
+
+@pytest.fixture
+def local_dataset(local_dataset_location):
+    return TIMDEXDataset.load(local_dataset_location)
diff --git a/tests/fixtures/local_datasets/dataset/source=alma/run_date=2023-03-06/run_type=daily/action=index/run_id=74afc7ba-9bbe-4f52-827d-c0595fa82036/records-0.parquet b/tests/fixtures/local_datasets/dataset/source=alma/run_date=2023-03-06/run_type=daily/action=index/run_id=74afc7ba-9bbe-4f52-827d-c0595fa82036/records-0.parquet
diff --git a/tests/fixtures/local_datasets/dataset/source=alma/run_date=2023-03-07/run_type=daily/action=delete/run_id=b831b653-028d-42eb-bf9f-c9fcdd46a982/records-0.parquet b/tests/fixtures/local_datasets/dataset/source=alma/run_date=2023-03-07/run_type=daily/action=delete/run_id=b831b653-028d-42eb-bf9f-c9fcdd46a982/records-0.parquet
diff --git a/tests/fixtures/local_datasets/dataset/source=libguides/run_date=2023-08-09/run_type=full/action=index/run_id=e90832e8-399f-476c-9b33-9ebe4120b5ab/records-0.parquet b/tests/fixtures/local_datasets/dataset/source=libguides/run_date=2023-08-09/run_type=full/action=index/run_id=e90832e8-399f-476c-9b33-9ebe4120b5ab/records-0.parquet
diff --git a/tests/test_config.py b/tests/test_config.py
@@ -0,0 +1,28 @@
+from timdex_dataset_api.config import configure_logger
+
+
+def test_configure_logger_default_info_level(monkeypatch, caplog):
+    caplog.set_level("DEBUG")
+    logger = configure_logger(__name__)
+
+    info_msg = "hello INFO world"
+    logger.info(info_msg)
+    assert info_msg in caplog.text
+
+    debug_msg = "hello DEBUG world"
+    logger.debug(debug_msg)
+    assert debug_msg not in caplog.text  # NOT captured
+
+
+def test_configure_logger_env_var_sets_debug_level(monkeypatch, caplog):
+    caplog.set_level("DEBUG")
+    monkeypatch.setenv("TDA_LOG_LEVEL", "DEBUG")
+    logger = configure_logger(__name__)
+
+    info_msg = "hello INFO world"
+    logger.info(info_msg)
+    assert info_msg in caplog.text
+
+    debug_msg = "hello DEBUG world"
+    logger.debug(debug_msg)
+    assert debug_msg in caplog.text  # IS captured
diff --git a/tests/test_dataset.py b/tests/test_dataset.py
@@ -0,0 +1,122 @@
+# ruff: noqa: S105, S106, SLF001
+
+from unittest.mock import MagicMock, patch
+
+import pyarrow as pa
+import pytest
+from pyarrow import fs
+
+from timdex_dataset_api.dataset import DatasetNotLoadedError, TIMDEXDataset
+
+
+@pytest.mark.parametrize(
+    ("location", "expected_filesystem", "expected_source"),
+    [
+        ("/path/to/dataset", fs.LocalFileSystem, "/path/to/dataset"),
+        (
+            ["/path/to/records1.parquet", "/path/to/records2.parquet"],
+            fs.LocalFileSystem,
+            ["/path/to/records1.parquet", "/path/to/records2.parquet"],
+        ),
+        ("s3://bucket/path/to/dataset", fs.S3FileSystem, "bucket/path/to/dataset"),
+        (
+            [
+                "s3://bucket/path/to/dataset/records1.parquet",
+                "s3://bucket/path/to/dataset/records2.parquet",
+            ],
+            fs.S3FileSystem,
+            [
+                "bucket/path/to/dataset/records1.parquet",
+                "bucket/path/to/dataset/records2.parquet",
+            ],
+        ),
+    ],
+)
+@patch("timdex_dataset_api.dataset.TIMDEXDataset.get_s3_filesystem")
+def test_parse_location_single_local_directory(
+    get_s3_filesystem,
+    location,
+    expected_filesystem,
+    expected_source,
+):
+    get_s3_filesystem.return_value = fs.S3FileSystem()
+    filesystem, source = TIMDEXDataset.parse_location(location)
+    assert isinstance(filesystem, expected_filesystem)
+    assert source == expected_source
+
+
+def test_get_s3_filesystem_success(mocker):
+    mocked_s3_filesystem = mocker.spy(fs, "S3FileSystem")
+    s3_filesystem = TIMDEXDataset.get_s3_filesystem()
+
+    assert mocked_s3_filesystem.call_args[1] == {
+        "secret_key": "fake_secret_key",
+        "access_key": "fake_access_key",
+        "region": "us-east-1",
+        "session_token": "fake_session_token",
+    }
+    assert isinstance(s3_filesystem, pa._s3fs.S3FileSystem)
+
+
+@patch("timdex_dataset_api.dataset.fs.LocalFileSystem")
+@patch("timdex_dataset_api.dataset.ds.dataset")
+def test_load_local_dataset_correct_filesystem_and_source(mock_pyarrow_ds, mock_local_fs):
+    mock_local_fs.return_value = MagicMock()
+    mock_pyarrow_ds.return_value = MagicMock()
+
+    timdex_dataset = TIMDEXDataset(location="local/path/to/dataset")
+    loaded_dataset = timdex_dataset.load_dataset()
+
+    mock_pyarrow_ds.assert_called_once_with(
+        "local/path/to/dataset",
+        schema=timdex_dataset.schema,
+        format="parquet",
+        partitioning="hive",
+        filesystem=mock_local_fs.return_value,
+    )
+    assert loaded_dataset == mock_pyarrow_ds.return_value
+
+
+@patch("timdex_dataset_api.dataset.TIMDEXDataset.get_s3_filesystem")
+@patch("timdex_dataset_api.dataset.ds.dataset")
+def test_load_s3_dataset_correct_filesystem_and_source(mock_pyarrow_ds, mock_get_s3_fs):
+    mock_get_s3_fs.return_value = MagicMock()
+    mock_pyarrow_ds.return_value = MagicMock()
+
+    timdex_dataset = TIMDEXDataset(location="s3://bucket/path/to/dataset")
+    loaded_dataset = timdex_dataset.load_dataset()
+
+    mock_get_s3_fs.assert_called_once()
+    mock_pyarrow_ds.assert_called_once_with(
+        "bucket/path/to/dataset",
+        schema=timdex_dataset.schema,
+        format="parquet",
+        partitioning="hive",
+        filesystem=mock_get_s3_fs.return_value,
+    )
+    assert loaded_dataset == mock_pyarrow_ds.return_value
+
+
+@patch("timdex_dataset_api.dataset.TIMDEXDataset.load_dataset")
+def test_load_method_loads_dataset_and_returns_timdexdataset_instance(mock_load_dataset):
+    mock_load_dataset.return_value = MagicMock()
+
+    timdex_dataset = TIMDEXDataset.load("s3://bucket/path/to/dataset")
+
+    assert isinstance(timdex_dataset, TIMDEXDataset)
+    assert timdex_dataset.location == "s3://bucket/path/to/dataset"
+    mock_load_dataset.assert_called_once()
+
+
+def test_local_dataset_is_valid(local_dataset):
+    assert local_dataset.dataset.to_table().validate() is None  # where None is valid
+
+
+def test_local_dataset_row_count_success(local_dataset):
+    assert local_dataset.dataset.count_rows() == local_dataset.row_count
+
+
+def test_local_dataset_row_count_missing_dataset_exception(local_dataset):
+    td = TIMDEXDataset(location="path/to/nowhere")
+    with pytest.raises(DatasetNotLoadedError):
+        _ = td.row_count
diff --git a/timdex_dataset_api/__init__.py b/timdex_dataset_api/__init__.py
@@ -1,3 +1,9 @@
 """timdex_dataset_api/__init__.py"""
 
+from timdex_dataset_api.dataset import TIMDEXDataset
+
 __version__ = "0.1.0"
+
+__all__ = [
+    "TIMDEXDataset",
+]
diff --git a/timdex_dataset_api/config.py b/timdex_dataset_api/config.py
@@ -0,0 +1,31 @@
+import logging
+import os
+
+
+def configure_logger(name: str) -> logging.Logger:
+    """Prepares and returns a logger instance for a given module name.
+
+    This approach is suitable for an installed and imported library such as this, where
+    any calling application logging levels and handlers should be utilized.
+
+    Args:
+        name (str): The name of the logger, typically __name__ is passed by caller
+    """
+    logger = logging.getLogger(name)
+    logger.addHandler(logging.NullHandler())
+
+    log_level = os.getenv("TDA_LOG_LEVEL", "INFO").strip().upper()
+    if log_level not in ["DEBUG", "INFO", "WARNING", "ERROR", "CRITICAL"]:
+        raise ValueError(f"Invalid log level: '{log_level}'")
+    logger.setLevel(getattr(logging, log_level))
+
+    handler = logging.StreamHandler()
+    handler.setFormatter(
+        logging.Formatter(
+            "%(asctime)s %(levelname)s %(name)s.%(funcName)s() "
+            "line %(lineno)d: %(message)s"
+        )
+    )
+    logger.addHandler(handler)
+
+    return logger
diff --git a/timdex_dataset_api/dataset.py b/timdex_dataset_api/dataset.py
@@ -0,0 +1,133 @@
+"""timdex_dataset_api/dataset.py"""
+
+import time
+
+import boto3
+import pyarrow as pa
+import pyarrow.dataset as ds
+from pyarrow import fs
+
+from timdex_dataset_api.config import configure_logger
+from timdex_dataset_api.exceptions import DatasetNotLoadedError
+
+logger = configure_logger(__name__)
+
+TIMDEX_DATASET_SCHEMA = pa.schema(
+    (
+        pa.field("timdex_record_id", pa.string()),
+        pa.field("source_record", pa.binary()),
+        pa.field("transformed_record", pa.binary()),
+        pa.field("source", pa.string()),
+        pa.field("run_date", pa.date32()),
+        pa.field("run_type", pa.string()),
+        pa.field("run_id", pa.string()),
+        pa.field("action", pa.string()),
+    )
+)
+
+
+class TIMDEXDataset:
+
+    def __init__(self, location: str | list[str]):
+        self.location = location
+        self.dataset: ds.Dataset = None  # type: ignore[assignment]
+        self.schema = TIMDEX_DATASET_SCHEMA
+
+    @classmethod
+    def load(cls, location: str) -> "TIMDEXDataset":
+        """Return an instantiated TIMDEXDataset object given a dataset location.
+
+        Argument 'location' may be a local filesystem path or an S3 URI to a parquet
+        dataset.
+        """
+        timdex_dataset = cls(location=location)
+        timdex_dataset.dataset = timdex_dataset.load_dataset()
+        return timdex_dataset
+
+    @staticmethod
+    def get_s3_filesystem() -> fs.FileSystem:
+        """Instantiate a pyarrow S3 Filesystem for dataset loading."""
+        session = boto3.session.Session()
+        credentials = session.get_credentials()
+        if not credentials:
+            raise RuntimeError("Could not locate AWS credentials")
+        return fs.S3FileSystem(
+            secret_key=credentials.secret_key,
+            access_key=credentials.access_key,
+            region=session.region_name,
+            session_token=credentials.token,
+        )
+
+    @classmethod
+    def parse_location(
+        cls,
+        location: str | list[str],
+    ) -> tuple[fs.FileSystem, str | list[str]]:
+        """Parse and return the filesystem and normalized source location(s).
+
+        Handles both single location strings and lists of Parquet file paths.
+        """
+        match location:
+            case str():
+                return cls._parse_single_location(location)
+            case list():
+                return cls._parse_multiple_locations(location)
+            case _:
+                raise TypeError("Location type must be str or list[str].")
+
+    @classmethod
+    def _parse_single_location(
+        cls, location: str
+    ) -> tuple[fs.FileSystem, str | list[str]]:
+        """Get filesystem and normalized location for single location."""
+        if location.startswith("s3://"):
+            filesystem = TIMDEXDataset.get_s3_filesystem()
+            source = location.removeprefix("s3://")
+        else:
+            filesystem = fs.LocalFileSystem()
+            source = location
+        return filesystem, source
+
+    @classmethod
+    def _parse_multiple_locations(
+        cls, location: list[str]
+    ) -> tuple[fs.FileSystem, str | list[str]]:
+        """Get filesystem and normalized location for multiple locations."""
+        if all(loc.startswith("s3://") for loc in location):
+            filesystem = TIMDEXDataset.get_s3_filesystem()
+            source = [loc.removeprefix("s3://") for loc in location]
+        elif all(not loc.startswith("s3://") for loc in location):
+            filesystem = fs.LocalFileSystem()
+            source = location
+        else:
+            raise ValueError("Mixed S3 and local paths are not supported.")
+        return filesystem, source
+
+    def load_dataset(self) -> ds.Dataset:
+        """Lazy load a pyarrow.Dataset for an already instantiated TIMDEXDataset object.
+
+        The dataset is loaded via the expected schema as defined by module constant
+        TIMDEX_DATASET_SCHEMA.  If the target dataset differs in any way, errors may be
+        raised when reading or writing data.
+        """
+        start_time = time.perf_counter()
+        filesystem, source = self.parse_location(self.location)
+        dataset = ds.dataset(
+            source,
+            schema=self.schema,
+            format="parquet",
+            partitioning="hive",
+            filesystem=filesystem,
+        )
+        logger.info(
+            f"Dataset successfully loaded: '{self.location}', "
+            f"{round(time.perf_counter()-start_time, 2)}s"
+        )
+        return dataset
+
+    @property
+    def row_count(self) -> int:
+        """Get row count from loaded dataset."""
+        if not self.dataset:
+            raise DatasetNotLoadedError
+        return self.dataset.count_rows()
diff --git a/timdex_dataset_api/exceptions.py b/timdex_dataset_api/exceptions.py
@@ -0,0 +1,5 @@
+"""timdex_dataset_api/exceptions.py"""
+
+
+class DatasetNotLoadedError(Exception):
+    """Custom exception for accessing methods requiring a loaded dataset."""

Original file line number	Diff line number	Diff line change
`@@ -94,6 +94,7 @@ ignore = [`
`94`	`94`	`"S320",`
`95`	`95`	`"S321",`
`96`	`96`	`"S608",`
	`97`	`+ "TRY003"`
`97`	`98`	`]`
`98`	`99`
`99`	`100`	`fixable = ["E", "F", "I", "Q"]`