atdata 0.2.2b1__py3-none-any.whl → 0.3.0b1__py3-none-any.whl

atdata/promote.py

@@ -4,27 +4,25 @@ This module provides functionality to promote locally-indexed datasets to the
 ATProto atmosphere network. This enables sharing datasets with the broader
 federation while maintaining schema consistency.
 
-Example:
-    ::
-
-        >>> from atdata.local import LocalIndex, Repo
-        >>> from atdata.atmosphere import AtmosphereClient, AtmosphereIndex
-        >>> from atdata.promote import promote_to_atmosphere
-        >>>
-        >>> # Setup
-        >>> local_index = LocalIndex()
-        >>> client = AtmosphereClient()
-        >>> client.login("handle.bsky.social", "app-password")
-        >>>
-        >>> # Promote a dataset
-        >>> entry = local_index.get_dataset("my-dataset")
-        >>> at_uri = promote_to_atmosphere(entry, local_index, client)
+Examples:
+    >>> from atdata.local import Index, Repo
+    >>> from atdata.atmosphere import AtmosphereClient, AtmosphereIndex
+    >>> from atdata.promote import promote_to_atmosphere
+    >>>
+    >>> # Setup
+    >>> local_index = Index()
+    >>> client = AtmosphereClient()
+    >>> client.login("handle.bsky.social", "app-password")
+    >>>
+    >>> # Promote a dataset
+    >>> entry = local_index.get_dataset("my-dataset")
+    >>> at_uri = promote_to_atmosphere(entry, local_index, client)
 """
 
 from typing import TYPE_CHECKING, Type
 
 if TYPE_CHECKING:
-    from .local import LocalDatasetEntry, Index as LocalIndex
+    from .local import LocalDatasetEntry, Index
     from .atmosphere import AtmosphereClient
     from ._protocols import AbstractDataStore, Packable
 
@@ -96,7 +94,7 @@ def _find_or_publish_schema(
 
 def promote_to_atmosphere(
     local_entry: "LocalDatasetEntry",
-    local_index: "LocalIndex",
+    local_index: "Index",
     atmosphere_client: "AtmosphereClient",
     *,
     data_store: "AbstractDataStore | None" = None,
@@ -128,13 +126,11 @@ def promote_to_atmosphere(
         KeyError: If schema not found in local index.
         ValueError: If local entry has no data URLs.
 
-    Example:
-        ::
-
-            >>> entry = local_index.get_dataset("mnist-train")
-            >>> uri = promote_to_atmosphere(entry, local_index, client)
-            >>> print(uri)
-            at://did:plc:abc123/ac.foundation.dataset.datasetIndex/...
+    Examples:
+        >>> entry = local_index.get_dataset("mnist-train")
+        >>> uri = promote_to_atmosphere(entry, local_index, client)
+        >>> print(uri)
+        at://did:plc:abc123/ac.foundation.dataset.datasetIndex/...
     """
     from .atmosphere import DatasetPublisher
     from ._schema_codec import schema_to_type

atdata/providers/__init__.py

@@ -0,0 +1,25 @@
+"""Storage provider backends for the local Index.
+
+This package defines the ``IndexProvider`` abstract base class and concrete
+implementations for Redis, SQLite, and PostgreSQL. The ``Index`` class in
+``atdata.local`` delegates all persistence to an ``IndexProvider``.
+
+Providers:
+    RedisProvider: Redis-backed storage (existing default).
+    SqliteProvider: SQLite file-based storage (zero external dependencies).
+    PostgresProvider: PostgreSQL storage (requires ``psycopg``).
+
+Examples:
+    >>> from atdata.providers import IndexProvider, create_provider
+    >>> provider = create_provider("sqlite", path="~/.atdata/index.db")
+    >>> from atdata.local import Index
+    >>> index = Index(provider=provider)
+"""
+
+from ._base import IndexProvider
+from ._factory import create_provider
+
+__all__ = [
+    "IndexProvider",
+    "create_provider",
+]

atdata/providers/_base.py

@@ -0,0 +1,140 @@
+"""Abstract base class for index storage providers.
+
+The ``IndexProvider`` ABC defines the persistence contract that the ``Index``
+class delegates to.  Each provider handles storage and retrieval of two entity
+types — dataset entries and schema records — using whatever backend it wraps.
+
+Concrete implementations live in sibling modules:
+    ``_redis.py``, ``_sqlite.py``, ``_postgres.py``
+"""
+
+from __future__ import annotations
+
+from abc import ABC, abstractmethod
+from typing import TYPE_CHECKING, Iterator
+
+if TYPE_CHECKING:
+    from ..local import LocalDatasetEntry
+
+
+class IndexProvider(ABC):
+    """Storage backend for the ``Index`` class.
+
+    Implementations persist ``LocalDatasetEntry`` objects and schema JSON
+    records.  The ``Index`` class owns all business logic (CID generation,
+    version bumping, schema building); the provider is a pure persistence
+    layer.
+
+    Examples:
+        >>> from atdata.providers import create_provider
+        >>> provider = create_provider("sqlite", path="/tmp/index.db")
+        >>> provider.store_schema("MySample", "1.0.0", '{"name": "MySample"}')
+        >>> provider.get_schema_json("MySample", "1.0.0")
+        '{"name": "MySample"}'
+    """
+
+    # ------------------------------------------------------------------
+    # Dataset entry operations
+    # ------------------------------------------------------------------
+
+    @abstractmethod
+    def store_entry(self, entry: LocalDatasetEntry) -> None:
+        """Persist a dataset entry (upsert by CID).
+
+        Args:
+            entry: The dataset entry to store.  The entry's ``cid`` property
+                is used as the primary key.
+        """
+
+    @abstractmethod
+    def get_entry_by_cid(self, cid: str) -> LocalDatasetEntry:
+        """Load a dataset entry by its content identifier.
+
+        Args:
+            cid: Content-addressable identifier.
+
+        Returns:
+            The matching ``LocalDatasetEntry``.
+
+        Raises:
+            KeyError: If no entry exists for *cid*.
+        """
+
+    @abstractmethod
+    def get_entry_by_name(self, name: str) -> LocalDatasetEntry:
+        """Load a dataset entry by its human-readable name.
+
+        Args:
+            name: Dataset name.
+
+        Returns:
+            The first matching ``LocalDatasetEntry``.
+
+        Raises:
+            KeyError: If no entry exists with *name*.
+        """
+
+    @abstractmethod
+    def iter_entries(self) -> Iterator[LocalDatasetEntry]:
+        """Iterate over all stored dataset entries.
+
+        Yields:
+            ``LocalDatasetEntry`` objects in unspecified order.
+        """
+
+    # ------------------------------------------------------------------
+    # Schema operations
+    # ------------------------------------------------------------------
+
+    @abstractmethod
+    def store_schema(self, name: str, version: str, schema_json: str) -> None:
+        """Persist a schema record (upsert by name + version).
+
+        Args:
+            name: Schema name (e.g. ``"MySample"``).
+            version: Semantic version string (e.g. ``"1.0.0"``).
+            schema_json: JSON-serialized schema record.
+        """
+
+    @abstractmethod
+    def get_schema_json(self, name: str, version: str) -> str | None:
+        """Load a schema's JSON by name and version.
+
+        Args:
+            name: Schema name.
+            version: Semantic version string.
+
+        Returns:
+            The JSON string, or ``None`` if not found.
+        """
+
+    @abstractmethod
+    def iter_schemas(self) -> Iterator[tuple[str, str, str]]:
+        """Iterate over all stored schemas.
+
+        Yields:
+            Tuples of ``(name, version, schema_json)``.
+        """
+
+    @abstractmethod
+    def find_latest_version(self, name: str) -> str | None:
+        """Find the latest semantic version for a schema name.
+
+        Args:
+            name: Schema name to search for.
+
+        Returns:
+            The latest version string (e.g. ``"1.2.3"``), or ``None``
+            if no schema with *name* exists.
+        """
+
+    # ------------------------------------------------------------------
+    # Lifecycle
+    # ------------------------------------------------------------------
+
+    def close(self) -> None:
+        """Release any resources held by the provider.
+
+        The default implementation is a no-op.  Providers that hold
+        connections (SQLite, PostgreSQL) should override this.
+        """

atdata/providers/_factory.py

@@ -0,0 +1,69 @@
+"""Factory for creating index providers by name.
+
+Examples:
+    >>> from atdata.providers._factory import create_provider
+    >>> provider = create_provider("sqlite", path="/tmp/index.db")
+    >>> provider = create_provider("redis")
+    >>> provider = create_provider("postgres", dsn="postgresql://localhost/mydb")
+"""
+
+from __future__ import annotations
+
+from pathlib import Path
+from typing import Any
+
+from ._base import IndexProvider
+
+
+def create_provider(
+    name: str,
+    *,
+    path: str | Path | None = None,
+    dsn: str | None = None,
+    redis: Any = None,
+    **kwargs: Any,
+) -> IndexProvider:
+    """Instantiate an ``IndexProvider`` by backend name.
+
+    Args:
+        name: One of ``"redis"``, ``"sqlite"``, or ``"postgres"``.
+        path: Database file path (SQLite).  Defaults to
+            ``~/.atdata/index.db`` when *name* is ``"sqlite"``.
+        dsn: Connection string (PostgreSQL).
+        redis: An existing ``redis.Redis`` connection (Redis).  When
+            ``None`` and *name* is ``"redis"``, a new connection is
+            created from *kwargs*.
+        **kwargs: Extra arguments forwarded to the provider constructor
+            (e.g. Redis host/port).
+
+    Returns:
+        A ready-to-use ``IndexProvider``.
+
+    Raises:
+        ValueError: If *name* is not a recognised backend.
+    """
+    name = name.lower().strip()
+
+    if name == "redis":
+        from ._redis import RedisProvider
+        from redis import Redis as _Redis
+
+        if redis is not None:
+            return RedisProvider(redis)
+        return RedisProvider(_Redis(**kwargs))
+
+    if name == "sqlite":
+        from ._sqlite import SqliteProvider
+
+        return SqliteProvider(path=path)
+
+    if name in ("postgres", "postgresql"):
+        from ._postgres import PostgresProvider
+
+        if dsn is None:
+            raise ValueError("dsn is required for the postgres provider")
+        return PostgresProvider(dsn=dsn)
+
+    raise ValueError(
+        f"Unknown provider {name!r}. Choose from: 'redis', 'sqlite', 'postgres'."
+    )

atdata/providers/_postgres.py

@@ -0,0 +1,214 @@
+"""PostgreSQL-backed index provider.
+
+Stores dataset entries and schema records in PostgreSQL tables.
+Requires the ``psycopg`` (v3) package, which is an optional dependency::
+
+    pip install "atdata[postgres]"
+
+The provider lazily imports ``psycopg`` so that ``import atdata`` never
+fails when the package is absent.
+"""
+
+from __future__ import annotations
+
+from typing import Iterator
+
+import msgpack
+
+from ._base import IndexProvider
+from .._type_utils import parse_semver
+
+_CREATE_TABLES = """\
+CREATE TABLE IF NOT EXISTS dataset_entries (
+    cid         TEXT PRIMARY KEY,
+    name        TEXT NOT NULL,
+    schema_ref  TEXT NOT NULL,
+    data_urls   BYTEA NOT NULL,
+    metadata    BYTEA,
+    legacy_uuid TEXT,
+    created_at  TIMESTAMPTZ DEFAULT now()
+);
+
+CREATE INDEX IF NOT EXISTS idx_entries_name
+    ON dataset_entries(name);
+
+CREATE TABLE IF NOT EXISTS schemas (
+    name        TEXT NOT NULL,
+    version     TEXT NOT NULL,
+    schema_json TEXT NOT NULL,
+    created_at  TIMESTAMPTZ DEFAULT now(),
+    PRIMARY KEY (name, version)
+);
+"""
+
+
+class PostgresProvider(IndexProvider):
+    """Index provider backed by PostgreSQL.
+
+    Args:
+        dsn: PostgreSQL connection string, e.g.
+            ``"postgresql://user:pass@host:5432/dbname"``.
+
+    Raises:
+        ImportError: If ``psycopg`` is not installed.
+
+    Examples:
+        >>> provider = PostgresProvider(dsn="postgresql://localhost/atdata")
+        >>> provider.store_schema("MySample", "1.0.0", '{"name":"MySample"}')
+    """
+
+    def __init__(self, dsn: str) -> None:
+        try:
+            import psycopg
+        except ImportError as exc:
+            raise ImportError(
+                "The postgres provider requires the 'psycopg' package. "
+                "Install it with: pip install 'atdata[postgres]'"
+            ) from exc
+
+        self._conn = psycopg.connect(dsn, autocommit=False)
+        with self._conn.cursor() as cur:
+            cur.execute(_CREATE_TABLES)
+        self._conn.commit()
+
+    # ------------------------------------------------------------------
+    # Dataset entry operations
+    # ------------------------------------------------------------------
+
+    def store_entry(self, entry: "LocalDatasetEntry") -> None:  # noqa: F821
+        with self._conn.cursor() as cur:
+            cur.execute(
+                """INSERT INTO dataset_entries
+                   (cid, name, schema_ref, data_urls, metadata, legacy_uuid)
+                   VALUES (%s, %s, %s, %s, %s, %s)
+                   ON CONFLICT (cid) DO UPDATE SET
+                       name = EXCLUDED.name,
+                       schema_ref = EXCLUDED.schema_ref,
+                       data_urls = EXCLUDED.data_urls,
+                       metadata = EXCLUDED.metadata,
+                       legacy_uuid = EXCLUDED.legacy_uuid""",
+                (
+                    entry.cid,
+                    entry.name,
+                    entry.schema_ref,
+                    msgpack.packb(entry.data_urls),
+                    msgpack.packb(entry.metadata)
+                    if entry.metadata is not None
+                    else None,
+                    entry._legacy_uuid,
+                ),
+            )
+        self._conn.commit()
+
+    def get_entry_by_cid(self, cid: str) -> "LocalDatasetEntry":  # noqa: F821
+        with self._conn.cursor() as cur:
+            cur.execute(
+                "SELECT cid, name, schema_ref, data_urls, metadata, legacy_uuid "
+                "FROM dataset_entries WHERE cid = %s",
+                (cid,),
+            )
+            row = cur.fetchone()
+        if row is None:
+            raise KeyError(f"LocalDatasetEntry not found: {cid}")
+        return _row_to_entry(row)
+
+    def get_entry_by_name(self, name: str) -> "LocalDatasetEntry":  # noqa: F821
+        with self._conn.cursor() as cur:
+            cur.execute(
+                "SELECT cid, name, schema_ref, data_urls, metadata, legacy_uuid "
+                "FROM dataset_entries WHERE name = %s LIMIT 1",
+                (name,),
+            )
+            row = cur.fetchone()
+        if row is None:
+            raise KeyError(f"No entry with name: {name}")
+        return _row_to_entry(row)
+
+    def iter_entries(self) -> Iterator["LocalDatasetEntry"]:  # noqa: F821
+        with self._conn.cursor() as cur:
+            cur.execute(
+                "SELECT cid, name, schema_ref, data_urls, metadata, legacy_uuid "
+                "FROM dataset_entries"
+            )
+            for row in cur:
+                yield _row_to_entry(row)
+
+    # ------------------------------------------------------------------
+    # Schema operations
+    # ------------------------------------------------------------------
+
+    def store_schema(self, name: str, version: str, schema_json: str) -> None:
+        with self._conn.cursor() as cur:
+            cur.execute(
+                """INSERT INTO schemas (name, version, schema_json)
+                   VALUES (%s, %s, %s)
+                   ON CONFLICT (name, version) DO UPDATE SET
+                       schema_json = EXCLUDED.schema_json""",
+                (name, version, schema_json),
+            )
+        self._conn.commit()
+
+    def get_schema_json(self, name: str, version: str) -> str | None:
+        with self._conn.cursor() as cur:
+            cur.execute(
+                "SELECT schema_json FROM schemas WHERE name = %s AND version = %s",
+                (name, version),
+            )
+            row = cur.fetchone()
+        if row is None:
+            return None
+        return row[0]
+
+    def iter_schemas(self) -> Iterator[tuple[str, str, str]]:
+        with self._conn.cursor() as cur:
+            cur.execute("SELECT name, version, schema_json FROM schemas")
+            for row in cur:
+                yield row[0], row[1], row[2]
+
+    def find_latest_version(self, name: str) -> str | None:
+        with self._conn.cursor() as cur:
+            cur.execute(
+                "SELECT version FROM schemas WHERE name = %s",
+                (name,),
+            )
+            latest: tuple[int, int, int] | None = None
+            latest_str: str | None = None
+            for (version_str,) in cur:
+                try:
+                    v = parse_semver(version_str)
+                    if latest is None or v > latest:
+                        latest = v
+                        latest_str = version_str
+                except ValueError:
+                    continue
+        return latest_str
+
+    # ------------------------------------------------------------------
+    # Lifecycle
+    # ------------------------------------------------------------------
+
+    def close(self) -> None:
+        """Close the PostgreSQL connection."""
+        self._conn.close()
+
+
+# ------------------------------------------------------------------
+# Helpers
+# ------------------------------------------------------------------
+
+
+def _row_to_entry(row: tuple) -> "LocalDatasetEntry":  # noqa: F821
+    """Convert a database row to a ``LocalDatasetEntry``."""
+    from ..local import LocalDatasetEntry
+
+    cid, name, schema_ref, data_urls_blob, metadata_blob, legacy_uuid = row
+    return LocalDatasetEntry(
+        name=name,
+        schema_ref=schema_ref,
+        data_urls=msgpack.unpackb(bytes(data_urls_blob)),
+        metadata=msgpack.unpackb(bytes(metadata_blob))
+        if metadata_blob is not None
+        else None,
+        _cid=cid,
+        _legacy_uuid=legacy_uuid,
+    )

atdata/providers/_redis.py

@@ -0,0 +1,171 @@
+"""Redis-backed index provider.
+
+This module extracts the Redis persistence logic that was previously
+inlined in ``atdata.local.Index`` and ``LocalDatasetEntry`` into a
+standalone ``IndexProvider`` implementation.
+"""
+
+from __future__ import annotations
+
+from typing import Iterator
+
+import msgpack
+from redis import Redis
+
+from ._base import IndexProvider
+from .._type_utils import parse_semver
+
+# Redis key prefixes — kept in sync with local.py constants
+_KEY_DATASET_ENTRY = "LocalDatasetEntry"
+_KEY_SCHEMA = "LocalSchema"
+
+
+class RedisProvider(IndexProvider):
+    """Index provider backed by a Redis connection.
+
+    This reproduces the exact storage layout used by the original
+    ``Index`` class so that existing Redis data is fully compatible.
+
+    Args:
+        redis: An active ``redis.Redis`` connection.
+    """
+
+    def __init__(self, redis: Redis) -> None:
+        self._redis = redis
+
+    @property
+    def redis(self) -> Redis:
+        """The underlying Redis connection (for advanced use / migration)."""
+        return self._redis
+
+    # ------------------------------------------------------------------
+    # Dataset entry operations
+    # ------------------------------------------------------------------
+
+    def store_entry(self, entry: "LocalDatasetEntry") -> None:  # noqa: F821
+        save_key = f"{_KEY_DATASET_ENTRY}:{entry.cid}"
+        data: dict[str, str | bytes] = {
+            "name": entry.name,
+            "schema_ref": entry.schema_ref,
+            "data_urls": msgpack.packb(entry.data_urls),
+            "cid": entry.cid,
+        }
+        if entry.metadata is not None:
+            data["metadata"] = msgpack.packb(entry.metadata)
+        if entry._legacy_uuid is not None:
+            data["legacy_uuid"] = entry._legacy_uuid
+
+        self._redis.hset(save_key, mapping=data)  # type: ignore[arg-type]
+
+    def get_entry_by_cid(self, cid: str) -> "LocalDatasetEntry":  # noqa: F821
+        save_key = f"{_KEY_DATASET_ENTRY}:{cid}"
+        raw_data = self._redis.hgetall(save_key)
+        if not raw_data:
+            raise KeyError(f"{_KEY_DATASET_ENTRY} not found: {cid}")
+
+        return _entry_from_redis_hash(raw_data)
+
+    def get_entry_by_name(self, name: str) -> "LocalDatasetEntry":  # noqa: F821
+        for entry in self.iter_entries():
+            if entry.name == name:
+                return entry
+        raise KeyError(f"No entry with name: {name}")
+
+    def iter_entries(self) -> Iterator["LocalDatasetEntry"]:  # noqa: F821
+        prefix = f"{_KEY_DATASET_ENTRY}:"
+        for key in self._redis.scan_iter(match=f"{prefix}*"):
+            key_str = key.decode("utf-8") if isinstance(key, bytes) else key
+            cid = key_str[len(prefix) :]
+            yield self.get_entry_by_cid(cid)
+
+    # ------------------------------------------------------------------
+    # Schema operations
+    # ------------------------------------------------------------------
+
+    def store_schema(self, name: str, version: str, schema_json: str) -> None:
+        redis_key = f"{_KEY_SCHEMA}:{name}@{version}"
+        self._redis.set(redis_key, schema_json)
+
+    def get_schema_json(self, name: str, version: str) -> str | None:
+        redis_key = f"{_KEY_SCHEMA}:{name}@{version}"
+        value = self._redis.get(redis_key)
+        if value is None:
+            return None
+        if isinstance(value, bytes):
+            return value.decode("utf-8")
+        return value  # type: ignore[return-value]
+
+    def iter_schemas(self) -> Iterator[tuple[str, str, str]]:
+        prefix = f"{_KEY_SCHEMA}:"
+        for key in self._redis.scan_iter(match=f"{prefix}*"):
+            key_str = key.decode("utf-8") if isinstance(key, bytes) else key
+            schema_id = key_str[len(prefix) :]
+
+            if "@" not in schema_id:
+                continue
+
+            raw_name, version = schema_id.rsplit("@", 1)
+            # Handle legacy format: module.Class -> Class
+            if "." in raw_name:
+                raw_name = raw_name.rsplit(".", 1)[1]
+
+            value = self._redis.get(key)
+            if value is None:
+                continue
+            schema_json = value.decode("utf-8") if isinstance(value, bytes) else value
+            yield raw_name, version, schema_json  # type: ignore[misc]
+
+    def find_latest_version(self, name: str) -> str | None:
+        latest: tuple[int, int, int] | None = None
+        latest_str: str | None = None
+
+        for schema_name, version, _ in self.iter_schemas():
+            if schema_name != name:
+                continue
+            try:
+                v = parse_semver(version)
+                if latest is None or v > latest:
+                    latest = v
+                    latest_str = version
+            except ValueError:
+                continue
+
+        return latest_str
+
+    # ------------------------------------------------------------------
+    # Lifecycle
+    # ------------------------------------------------------------------
+
+    def close(self) -> None:
+        """Close the Redis connection."""
+        self._redis.close()
+
+
+# ------------------------------------------------------------------
+# Helpers
+# ------------------------------------------------------------------
+
+
+def _entry_from_redis_hash(raw_data: dict) -> "LocalDatasetEntry":  # noqa: F821
+    """Reconstruct a ``LocalDatasetEntry`` from a Redis hash mapping."""
+    from ..local import LocalDatasetEntry
+    from typing import cast
+
+    raw = cast(dict[bytes, bytes], raw_data)
+    name = raw[b"name"].decode("utf-8")
+    schema_ref = raw[b"schema_ref"].decode("utf-8")
+    cid_value = raw.get(b"cid", b"").decode("utf-8") or None
+    legacy_uuid = raw.get(b"legacy_uuid", b"").decode("utf-8") or None
+    data_urls = msgpack.unpackb(raw[b"data_urls"])
+    metadata = None
+    if b"metadata" in raw:
+        metadata = msgpack.unpackb(raw[b"metadata"])
+
+    return LocalDatasetEntry(
+        name=name,
+        schema_ref=schema_ref,
+        data_urls=data_urls,
+        metadata=metadata,
+        _cid=cid_value,
+        _legacy_uuid=legacy_uuid,
+    )