juniper-data 0.4.2__py3-none-any.whl

juniper_data/storage/postgres_store.py

@@ -0,0 +1,373 @@
+"""PostgreSQL-backed dataset storage for metadata with file system artifacts."""
+
+import io
+import json
+from datetime import datetime  # noqa: F401 - used by DatasetMeta serialization
+from pathlib import Path
+from typing import Any
+
+import numpy as np
+
+from juniper_data.core.models import DatasetMeta
+
+from .base import DatasetStore
+
+try:
+    import psycopg2
+    from psycopg2.extras import RealDictCursor
+
+    POSTGRES_AVAILABLE = True
+except ImportError:
+    POSTGRES_AVAILABLE = False
+    psycopg2 = None  # type: ignore[assignment]
+
+
+class PostgresDatasetStore(DatasetStore):
+    """PostgreSQL-backed dataset storage.
+
+    Stores metadata in PostgreSQL and artifacts on the local filesystem.
+    Suitable for production deployments with database-backed metadata.
+
+    Requires the `psycopg2` package: pip install psycopg2-binary
+    """
+
+    SCHEMA_SQL = """
+    CREATE TABLE IF NOT EXISTS datasets (
+        dataset_id VARCHAR(255) PRIMARY KEY,
+        generator VARCHAR(100) NOT NULL,
+        generator_version VARCHAR(50) NOT NULL,
+        params JSONB NOT NULL,
+        n_samples INTEGER NOT NULL,
+        n_features INTEGER NOT NULL,
+        n_classes INTEGER NOT NULL,
+        n_train INTEGER NOT NULL,
+        n_test INTEGER NOT NULL,
+        class_distribution JSONB NOT NULL,
+        artifact_formats TEXT[] NOT NULL DEFAULT ARRAY['npz'],
+        created_at TIMESTAMP WITH TIME ZONE NOT NULL,
+        checksum VARCHAR(64),
+        tags TEXT[] NOT NULL DEFAULT ARRAY[]::TEXT[],
+        ttl_seconds INTEGER,
+        expires_at TIMESTAMP WITH TIME ZONE,
+        last_accessed_at TIMESTAMP WITH TIME ZONE,
+        access_count INTEGER NOT NULL DEFAULT 0
+    );
+
+    CREATE INDEX IF NOT EXISTS idx_datasets_generator ON datasets(generator);
+    CREATE INDEX IF NOT EXISTS idx_datasets_created_at ON datasets(created_at);
+    CREATE INDEX IF NOT EXISTS idx_datasets_expires_at ON datasets(expires_at);
+    """
+
+    def __init__(
+        self,
+        host: str = "localhost",
+        port: int = 5432,
+        database: str = "juniper_data",
+        user: str = "postgres",
+        password: str | None = None,
+        artifact_path: Path | None = None,
+        connection_string: str | None = None,
+        auto_create_schema: bool = True,
+    ) -> None:
+        """Initialize PostgreSQL connection.
+
+        Args:
+            host: PostgreSQL server hostname.
+            port: PostgreSQL server port.
+            database: Database name.
+            user: Database user.
+            password: Database password.
+            artifact_path: Path for storing NPZ artifacts.
+            connection_string: Optional full connection string (overrides other params).
+            auto_create_schema: Automatically create tables if they don't exist.
+
+        Raises:
+            ImportError: If psycopg2 package is not installed.
+        """
+        if not POSTGRES_AVAILABLE:
+            raise ImportError("psycopg2 package not installed. Install with: pip install psycopg2-binary")
+
+        self._artifact_path = artifact_path or Path("./data/datasets")
+        self._artifact_path.mkdir(parents=True, exist_ok=True)
+
+        if connection_string:
+            self._conn_params: dict[str, Any] = {"dsn": connection_string}
+        else:
+            self._conn_params = {
+                "host": host,
+                "port": str(port),
+                "database": database,
+                "user": user,
+                "password": password or "",
+            }
+
+        if auto_create_schema:
+            self._create_schema()
+
+    def _get_connection(self) -> Any:
+        """Get a new database connection."""
+        return psycopg2.connect(**self._conn_params)
+
+    def _create_schema(self) -> None:
+        """Create database schema if it doesn't exist."""
+        with self._get_connection() as conn:
+            with conn.cursor() as cur:
+                cur.execute(self.SCHEMA_SQL)
+            conn.commit()
+
+    def _artifact_file(self, dataset_id: str) -> Path:
+        """Get the artifact file path for a dataset."""
+        return self._artifact_path / f"{dataset_id}.npz"
+
+    def _meta_to_row(self, meta: DatasetMeta) -> dict:
+        """Convert DatasetMeta to database row dict."""
+        return {
+            "dataset_id": meta.dataset_id,
+            "generator": meta.generator,
+            "generator_version": meta.generator_version,
+            "params": json.dumps(meta.params),
+            "n_samples": meta.n_samples,
+            "n_features": meta.n_features,
+            "n_classes": meta.n_classes,
+            "n_train": meta.n_train,
+            "n_test": meta.n_test,
+            "class_distribution": json.dumps(meta.class_distribution),
+            "artifact_formats": meta.artifact_formats,
+            "created_at": meta.created_at,
+            "checksum": meta.checksum,
+            "tags": meta.tags,
+            "ttl_seconds": meta.ttl_seconds,
+            "expires_at": meta.expires_at,
+            "last_accessed_at": meta.last_accessed_at,
+            "access_count": meta.access_count,
+        }
+
+    def _row_to_meta(self, row: dict) -> DatasetMeta:
+        """Convert database row to DatasetMeta."""
+        return DatasetMeta(
+            dataset_id=row["dataset_id"],
+            generator=row["generator"],
+            generator_version=row["generator_version"],
+            params=row["params"] if isinstance(row["params"], dict) else json.loads(row["params"]),
+            n_samples=row["n_samples"],
+            n_features=row["n_features"],
+            n_classes=row["n_classes"],
+            n_train=row["n_train"],
+            n_test=row["n_test"],
+            class_distribution=row["class_distribution"]
+            if isinstance(row["class_distribution"], dict)
+            else json.loads(row["class_distribution"]),
+            artifact_formats=list(row["artifact_formats"]),
+            created_at=row["created_at"],
+            checksum=row["checksum"],
+            tags=list(row["tags"]) if row["tags"] else [],
+            ttl_seconds=row["ttl_seconds"],
+            expires_at=row["expires_at"],
+            last_accessed_at=row["last_accessed_at"],
+            access_count=row["access_count"],
+        )
+
+    def save(
+        self,
+        dataset_id: str,
+        meta: DatasetMeta,
+        arrays: dict[str, np.ndarray],
+    ) -> None:
+        """Save dataset to PostgreSQL and filesystem.
+
+        Args:
+            dataset_id: Unique identifier for the dataset.
+            meta: Dataset metadata.
+            arrays: Dictionary of numpy arrays.
+        """
+        row = self._meta_to_row(meta)
+
+        insert_sql = """
+        INSERT INTO datasets (
+            dataset_id, generator, generator_version, params, n_samples,
+            n_features, n_classes, n_train, n_test, class_distribution,
+            artifact_formats, created_at, checksum, tags, ttl_seconds,
+            expires_at, last_accessed_at, access_count
+        ) VALUES (
+            %(dataset_id)s, %(generator)s, %(generator_version)s, %(params)s::jsonb,
+            %(n_samples)s, %(n_features)s, %(n_classes)s, %(n_train)s, %(n_test)s,
+            %(class_distribution)s::jsonb, %(artifact_formats)s, %(created_at)s,
+            %(checksum)s, %(tags)s, %(ttl_seconds)s, %(expires_at)s,
+            %(last_accessed_at)s, %(access_count)s
+        ) ON CONFLICT (dataset_id) DO UPDATE SET
+            generator = EXCLUDED.generator,
+            generator_version = EXCLUDED.generator_version,
+            params = EXCLUDED.params,
+            n_samples = EXCLUDED.n_samples,
+            n_features = EXCLUDED.n_features,
+            n_classes = EXCLUDED.n_classes,
+            n_train = EXCLUDED.n_train,
+            n_test = EXCLUDED.n_test,
+            class_distribution = EXCLUDED.class_distribution,
+            artifact_formats = EXCLUDED.artifact_formats,
+            checksum = EXCLUDED.checksum,
+            tags = EXCLUDED.tags,
+            ttl_seconds = EXCLUDED.ttl_seconds,
+            expires_at = EXCLUDED.expires_at,
+            last_accessed_at = EXCLUDED.last_accessed_at,
+            access_count = EXCLUDED.access_count
+        """
+
+        with self._get_connection() as conn:
+            with conn.cursor() as cur:
+                cur.execute(insert_sql, row)
+            conn.commit()
+
+        artifact_path = self._artifact_file(dataset_id)
+        buffer = io.BytesIO()
+        np.savez_compressed(buffer, **arrays)  # type: ignore[arg-type]
+        artifact_path.write_bytes(buffer.getvalue())
+
+    def get_meta(self, dataset_id: str) -> DatasetMeta | None:
+        """Get dataset metadata from PostgreSQL.
+
+        Args:
+            dataset_id: Unique identifier for the dataset.
+
+        Returns:
+            Dataset metadata if found, None otherwise.
+        """
+        with self._get_connection() as conn:
+            with conn.cursor(cursor_factory=RealDictCursor) as cur:
+                cur.execute("SELECT * FROM datasets WHERE dataset_id = %s", (dataset_id,))
+                row = cur.fetchone()
+
+        if row is None:
+            return None
+
+        return self._row_to_meta(dict(row))
+
+    def get_artifact_bytes(self, dataset_id: str) -> bytes | None:
+        """Get dataset artifact bytes from filesystem.
+
+        Args:
+            dataset_id: Unique identifier for the dataset.
+
+        Returns:
+            NPZ bytes if found, None otherwise.
+        """
+        artifact_path = self._artifact_file(dataset_id)
+        if not artifact_path.exists():
+            return None
+        return artifact_path.read_bytes()
+
+    def exists(self, dataset_id: str) -> bool:
+        """Check if dataset exists in PostgreSQL.
+
+        Args:
+            dataset_id: Unique identifier for the dataset.
+
+        Returns:
+            True if the dataset exists, False otherwise.
+        """
+        with self._get_connection() as conn:
+            with conn.cursor() as cur:
+                cur.execute("SELECT 1 FROM datasets WHERE dataset_id = %s", (dataset_id,))
+                return cur.fetchone() is not None
+
+    def delete(self, dataset_id: str) -> bool:
+        """Delete dataset from PostgreSQL and filesystem.
+
+        Args:
+            dataset_id: Unique identifier for the dataset.
+
+        Returns:
+            True if the dataset was deleted, False if it didn't exist.
+        """
+        with self._get_connection() as conn:
+            with conn.cursor() as cur:
+                cur.execute(
+                    "DELETE FROM datasets WHERE dataset_id = %s RETURNING dataset_id",
+                    (dataset_id,),
+                )
+                deleted = cur.fetchone() is not None
+            conn.commit()
+
+        artifact_path = self._artifact_file(dataset_id)
+        if artifact_path.exists():
+            artifact_path.unlink()
+
+        return deleted
+
+    def list_datasets(self, limit: int = 100, offset: int = 0) -> list[str]:
+        """List dataset IDs from PostgreSQL.
+
+        Args:
+            limit: Maximum number of dataset IDs to return.
+            offset: Number of dataset IDs to skip.
+
+        Returns:
+            List of dataset IDs.
+        """
+        with self._get_connection() as conn:
+            with conn.cursor() as cur:
+                cur.execute(
+                    "SELECT dataset_id FROM datasets ORDER BY created_at DESC LIMIT %s OFFSET %s",
+                    (limit, offset),
+                )
+                rows = cur.fetchall()
+
+        return [row[0] for row in rows]
+
+    def update_meta(self, dataset_id: str, meta: DatasetMeta) -> bool:
+        """Update dataset metadata in PostgreSQL.
+
+        Args:
+            dataset_id: Unique identifier for the dataset.
+            meta: Updated dataset metadata.
+
+        Returns:
+            True if the dataset was updated, False if it didn't exist.
+        """
+        row = self._meta_to_row(meta)
+
+        update_sql = """
+        UPDATE datasets SET
+            generator = %(generator)s,
+            generator_version = %(generator_version)s,
+            params = %(params)s::jsonb,
+            n_samples = %(n_samples)s,
+            n_features = %(n_features)s,
+            n_classes = %(n_classes)s,
+            n_train = %(n_train)s,
+            n_test = %(n_test)s,
+            class_distribution = %(class_distribution)s::jsonb,
+            artifact_formats = %(artifact_formats)s,
+            checksum = %(checksum)s,
+            tags = %(tags)s,
+            ttl_seconds = %(ttl_seconds)s,
+            expires_at = %(expires_at)s,
+            last_accessed_at = %(last_accessed_at)s,
+            access_count = %(access_count)s
+        WHERE dataset_id = %(dataset_id)s
+        """
+
+        with self._get_connection() as conn:
+            with conn.cursor() as cur:
+                cur.execute(update_sql, row)
+                updated = cur.rowcount > 0
+            conn.commit()
+
+        return updated
+
+    def list_all_metadata(self) -> list[DatasetMeta]:
+        """List all dataset metadata from PostgreSQL.
+
+        Returns:
+            List of all DatasetMeta objects.
+        """
+        with self._get_connection() as conn:
+            with conn.cursor(cursor_factory=RealDictCursor) as cur:
+                cur.execute("SELECT * FROM datasets ORDER BY created_at DESC")
+                rows = cur.fetchall()
+
+        return [self._row_to_meta(dict(row)) for row in rows]
+
+    def close(self) -> None:
+        """Close database connections (no-op for connection-per-request pattern)."""
+        pass

juniper_data/storage/redis_store.py

@@ -0,0 +1,264 @@
+"""Redis-backed dataset storage for caching and distributed deployments."""
+
+import io
+import json
+from datetime import datetime  # noqa: F401 - used by DatasetMeta serialization
+from typing import Any
+
+import numpy as np
+
+from juniper_data.core.models import DatasetMeta
+
+from .base import DatasetStore
+
+try:
+    import redis
+
+    REDIS_AVAILABLE = True
+except ImportError:
+    REDIS_AVAILABLE = False
+    redis = None  # type: ignore[assignment]
+
+
+class RedisDatasetStore(DatasetStore):
+    """Redis-backed dataset storage.
+
+    Uses Redis for both metadata (as JSON) and artifact storage (as bytes).
+    Suitable for caching layers and distributed deployments.
+
+    Requires the `redis` package: pip install redis
+    """
+
+    def __init__(
+        self,
+        host: str = "localhost",
+        port: int = 6379,
+        db: int = 0,
+        password: str | None = None,
+        key_prefix: str = "juniper:dataset:",
+        default_ttl: int | None = None,
+        connection_pool: Any | None = None,
+    ) -> None:
+        """Initialize Redis connection.
+
+        Args:
+            host: Redis server hostname.
+            port: Redis server port.
+            db: Redis database number.
+            password: Redis password (optional).
+            key_prefix: Prefix for all Redis keys.
+            default_ttl: Default TTL for stored datasets in seconds (optional).
+            connection_pool: Optional existing Redis connection pool.
+
+        Raises:
+            ImportError: If redis package is not installed.
+        """
+        if not REDIS_AVAILABLE:
+            raise ImportError("Redis package not installed. Install with: pip install redis")
+
+        self._key_prefix = key_prefix
+        self._default_ttl = default_ttl
+
+        if connection_pool:
+            self._client: redis.Redis[bytes] = redis.Redis(connection_pool=connection_pool)
+        else:
+            self._client = redis.Redis(
+                host=host,
+                port=port,
+                db=db,
+                password=password,
+                decode_responses=False,
+            )
+
+    def _meta_key(self, dataset_id: str) -> str:
+        """Get Redis key for metadata."""
+        return f"{self._key_prefix}{dataset_id}:meta"
+
+    def _artifact_key(self, dataset_id: str) -> str:
+        """Get Redis key for artifact data."""
+        return f"{self._key_prefix}{dataset_id}:artifact"
+
+    def _encode_meta(self, meta: DatasetMeta) -> bytes:
+        """Encode metadata to JSON bytes."""
+        data = meta.model_dump(mode="json")
+        return json.dumps(data).encode("utf-8")
+
+    def _decode_meta(self, data: bytes) -> DatasetMeta:
+        """Decode metadata from JSON bytes."""
+        parsed = json.loads(data.decode("utf-8"))
+        return DatasetMeta(**parsed)
+
+    def _encode_arrays(self, arrays: dict[str, np.ndarray]) -> bytes:
+        """Encode arrays to NPZ bytes."""
+        buffer = io.BytesIO()
+        np.savez_compressed(buffer, **arrays)  # type: ignore[arg-type]
+        return buffer.getvalue()
+
+    def save(
+        self,
+        dataset_id: str,
+        meta: DatasetMeta,
+        arrays: dict[str, np.ndarray],
+    ) -> None:
+        """Save dataset to Redis.
+
+        Args:
+            dataset_id: Unique identifier for the dataset.
+            meta: Dataset metadata.
+            arrays: Dictionary of numpy arrays.
+        """
+        meta_bytes = self._encode_meta(meta)
+        artifact_bytes = self._encode_arrays(arrays)
+
+        meta_key = self._meta_key(dataset_id)
+        artifact_key = self._artifact_key(dataset_id)
+
+        ttl = meta.ttl_seconds or self._default_ttl
+
+        pipe = self._client.pipeline()
+        if ttl:
+            pipe.setex(meta_key, ttl, meta_bytes)
+            pipe.setex(artifact_key, ttl, artifact_bytes)
+        else:
+            pipe.set(meta_key, meta_bytes)
+            pipe.set(artifact_key, artifact_bytes)
+        pipe.execute()
+
+    def get_meta(self, dataset_id: str) -> DatasetMeta | None:
+        """Get dataset metadata from Redis.
+
+        Args:
+            dataset_id: Unique identifier for the dataset.
+
+        Returns:
+            Dataset metadata if found, None otherwise.
+        """
+        data = self._client.get(self._meta_key(dataset_id))
+        if data is None:
+            return None
+        return self._decode_meta(data)
+
+    def get_artifact_bytes(self, dataset_id: str) -> bytes | None:
+        """Get dataset artifact bytes from Redis.
+
+        Args:
+            dataset_id: Unique identifier for the dataset.
+
+        Returns:
+            NPZ bytes if found, None otherwise.
+        """
+        return self._client.get(self._artifact_key(dataset_id))
+
+    def exists(self, dataset_id: str) -> bool:
+        """Check if dataset exists in Redis.
+
+        Args:
+            dataset_id: Unique identifier for the dataset.
+
+        Returns:
+            True if the dataset exists, False otherwise.
+        """
+        return bool(self._client.exists(self._meta_key(dataset_id)))
+
+    def delete(self, dataset_id: str) -> bool:
+        """Delete dataset from Redis.
+
+        Args:
+            dataset_id: Unique identifier for the dataset.
+
+        Returns:
+            True if the dataset was deleted, False if it didn't exist.
+        """
+        meta_key = self._meta_key(dataset_id)
+        artifact_key = self._artifact_key(dataset_id)
+
+        deleted = self._client.delete(meta_key, artifact_key)
+        return deleted > 0
+
+    def list_datasets(self, limit: int = 100, offset: int = 0) -> list[str]:
+        """List dataset IDs from Redis.
+
+        Args:
+            limit: Maximum number of dataset IDs to return.
+            offset: Number of dataset IDs to skip.
+
+        Returns:
+            List of dataset IDs.
+        """
+        pattern = f"{self._key_prefix}*:meta"
+        keys = list(self._client.scan_iter(match=pattern))
+
+        dataset_ids = []
+        for key in keys:
+            key_str = key.decode("utf-8") if isinstance(key, bytes) else key
+            dataset_id = key_str[len(self._key_prefix) : -5]
+            dataset_ids.append(dataset_id)
+
+        dataset_ids.sort()
+        return dataset_ids[offset : offset + limit]
+
+    def update_meta(self, dataset_id: str, meta: DatasetMeta) -> bool:
+        """Update dataset metadata in Redis.
+
+        Args:
+            dataset_id: Unique identifier for the dataset.
+            meta: Updated dataset metadata.
+
+        Returns:
+            True if the dataset was updated, False if it didn't exist.
+        """
+        meta_key = self._meta_key(dataset_id)
+        if not self._client.exists(meta_key):
+            return False
+
+        ttl = self._client.ttl(meta_key)
+        meta_bytes = self._encode_meta(meta)
+
+        if ttl > 0:
+            self._client.setex(meta_key, ttl, meta_bytes)
+        else:
+            self._client.set(meta_key, meta_bytes)
+
+        return True
+
+    def list_all_metadata(self) -> list[DatasetMeta]:
+        """List all dataset metadata from Redis.
+
+        Returns:
+            List of all DatasetMeta objects.
+        """
+        pattern = f"{self._key_prefix}*:meta"
+        keys = list(self._client.scan_iter(match=pattern))
+
+        metadata: list[DatasetMeta] = []
+        for key in keys:
+            data = self._client.get(key)
+            if data is not None:
+                metadata.append(self._decode_meta(data))
+
+        return metadata
+
+    def ping(self) -> bool:
+        """Check if Redis connection is alive.
+
+        Returns:
+            True if connected, False otherwise.
+        """
+        try:
+            return self._client.ping()
+        except Exception:
+            return False
+
+    def flush_prefix(self) -> int:
+        """Delete all keys with this store's prefix.
+
+        WARNING: This deletes all datasets in this store.
+
+        Returns:
+            Number of keys deleted.
+        """
+        pattern = f"{self._key_prefix}*"
+        keys = list(self._client.scan_iter(match=pattern))
+        if keys:
+            return self._client.delete(*keys)
+        return 0

juniper_data/tests/__init__.py

@@ -0,0 +1 @@
+"""Test suite for Juniper Data."""

juniper_data/tests/conftest.py

@@ -0,0 +1,68 @@
+"""Common pytest fixtures for juniper_data tests."""
+
+import numpy as np
+import pytest
+
+from juniper_data.generators.spiral import SpiralGenerator, SpiralParams
+
+
+@pytest.fixture
+def default_spiral_params() -> SpiralParams:
+    """Default spiral parameters for testing."""
+    return SpiralParams()
+
+
+@pytest.fixture
+def two_spiral_params() -> SpiralParams:
+    """Parameters for a 2-spiral dataset with 100 points per spiral."""
+    return SpiralParams(
+        n_spirals=2,
+        n_points_per_spiral=100,
+        seed=42,
+    )
+
+
+@pytest.fixture
+def three_spiral_params() -> SpiralParams:
+    """Parameters for a 3-spiral dataset with 50 points per spiral."""
+    return SpiralParams(
+        n_spirals=3,
+        n_points_per_spiral=50,
+        seed=42,
+    )
+
+
+@pytest.fixture
+def minimal_spiral_params() -> SpiralParams:
+    """Minimal valid spiral parameters for fast tests."""
+    return SpiralParams(
+        n_spirals=2,
+        n_points_per_spiral=10,
+        seed=42,
+    )
+
+
+@pytest.fixture
+def generated_two_spiral_dataset(two_spiral_params: SpiralParams) -> dict[str, np.ndarray]:
+    """Generate a 2-spiral dataset for testing."""
+    return SpiralGenerator.generate(two_spiral_params)
+
+
+@pytest.fixture
+def generated_three_spiral_dataset(three_spiral_params: SpiralParams) -> dict[str, np.ndarray]:
+    """Generate a 3-spiral dataset for testing."""
+    return SpiralGenerator.generate(three_spiral_params)
+
+
+@pytest.fixture
+def generated_minimal_dataset(minimal_spiral_params: SpiralParams) -> dict[str, np.ndarray]:
+    """Generate a minimal dataset for fast tests."""
+    return SpiralGenerator.generate(minimal_spiral_params)
+
+
+@pytest.fixture
+def sample_arrays() -> dict[str, np.ndarray]:
+    """Simple sample arrays for split/shuffle testing."""
+    X = np.arange(20).reshape(10, 2).astype(np.float32)
+    y = np.eye(2, dtype=np.float32)[np.arange(10) % 2]
+    return {"X": X, "y": y}