Establish TIMDEXDataset class

Why these changes are being introduced:

All operations for a dataset can be wrapped in a TIMDEXDataset class.
This class would be responsible for loading a dataset that read and
write operations would be performed against.

How this addresses that need:
* Establishes TIMDEXDataset class
* Adds methods for loading a database
* Support for local or S3 location for dataset

Side effects of this change:
* None

Relevant ticket(s):
* https://mitlibraries.atlassian.net/browse/TIMX-415

Author: ghukill

tests/conftest.py

@@ -1,6 +1,18 @@
 import pytest
 
+from timdex_dataset_api import TIMDEXDataset
+
 
 @pytest.fixture(autouse=True)
 def _test_env(monkeypatch):
     monkeypatch.setenv("TDA_LOG_LEVEL", "INFO")
+
+
+@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)

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

@@ -0,0 +1,129 @@
+# ruff: noqa: S105, S106, SLF001
+
+from unittest.mock import MagicMock, patch
+
+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
+
+
+@patch("timdex_dataset_api.dataset.fs.S3FileSystem")
+@patch("timdex_dataset_api.dataset.boto3.session.Session")
+def test_get_s3_filesystem_success(mock_session, mock_s3_filesystem):
+    mock_credentials = MagicMock()
+    mock_credentials.secret_key = "fake_secret_key"
+    mock_credentials.access_key = "fake_access_key"
+    mock_credentials.token = "fake_session_token"
+    mock_session.return_value.get_credentials.return_value = mock_credentials
+    mock_session.return_value.region_name = "us-east-1"
+
+    s3_filesystem = TIMDEXDataset.get_s3_filesystem()
+
+    mock_s3_filesystem.assert_called_once_with(
+        secret_key="fake_secret_key",
+        access_key="fake_access_key",
+        region="us-east-1",
+        session_token="fake_session_token",
+    )
+    assert isinstance(s3_filesystem, MagicMock)
+
+
+@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

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

@@ -1,3 +1,9 @@
 """timdex_dataset_api/__init__.py"""
 
+from timdex_dataset_api.dataset import TIMDEXDataset
+
 __version__ = "0.1.0"
+
+__all__ = [
+    "TIMDEXDataset",
+]

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

@@ -0,0 +1,118 @@
+"""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,
+        )
+
+    @staticmethod
+    def parse_location(
+        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.
+        """
+        source: str | list[str]
+        if isinstance(location, str):
+            if location.startswith("s3://"):
+                filesystem = TIMDEXDataset.get_s3_filesystem()
+                source = location.removeprefix("s3://")
+            else:
+                filesystem = fs.LocalFileSystem()
+                source = location
+        elif isinstance(location, list):
+            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.")
+        else:
+            raise TypeError("Location type must be str or list[str].")
+
+        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()

tests/test_dataset.py

@@ -0,0 +1,5 @@
+"""timdex_dataset_api/exceptions.py"""
+
+
+class DatasetNotLoadedError(Exception):
+    """Custom exception for accessing methods requiring a loaded dataset."""