atdata 0.2.3b1__py3-none-any.whl → 0.3.1b1__py3-none-any.whl

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,
+    )

atdata/providers/_sqlite.py

@@ -0,0 +1,191 @@
+"""SQLite-backed index provider.
+
+Stores dataset entries and schema records in a local SQLite database file.
+Uses WAL journal mode for concurrent read access and ``INSERT OR REPLACE``
+for upsert semantics.
+
+No external dependencies — uses Python's built-in ``sqlite3`` module.
+"""
+
+from __future__ import annotations
+
+import sqlite3
+from pathlib import Path
+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   BLOB NOT NULL,
+    metadata    BLOB,
+    legacy_uuid TEXT,
+    created_at  TEXT DEFAULT (datetime('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  TEXT DEFAULT (datetime('now')),
+    PRIMARY KEY (name, version)
+);
+"""
+
+
+class SqliteProvider(IndexProvider):
+    """Index provider backed by a local SQLite database.
+
+    Args:
+        path: Path to the database file.  The parent directory is created
+            automatically.  Defaults to ``~/.atdata/index.db``.
+
+    Examples:
+        >>> provider = SqliteProvider(path="/tmp/test-index.db")
+        >>> provider.store_schema("MySample", "1.0.0", '{"name":"MySample"}')
+        >>> provider.get_schema_json("MySample", "1.0.0")
+        '{"name":"MySample"}'
+    """
+
+    def __init__(self, path: str | Path | None = None) -> None:
+        if path is None:
+            path = Path.home() / ".atdata" / "index.db"
+        self._path = Path(path).expanduser()
+        self._path.parent.mkdir(parents=True, exist_ok=True)
+
+        self._conn = sqlite3.connect(str(self._path))
+        self._conn.execute("PRAGMA journal_mode=WAL")
+        self._conn.executescript(_CREATE_TABLES)
+        self._conn.commit()
+
+    @property
+    def path(self) -> Path:
+        """Path to the SQLite database file."""
+        return self._path
+
+    # ------------------------------------------------------------------
+    # Dataset entry operations
+    # ------------------------------------------------------------------
+
+    def store_entry(self, entry: "LocalDatasetEntry") -> None:  # noqa: F821
+        self._conn.execute(
+            """INSERT OR REPLACE INTO dataset_entries
+               (cid, name, schema_ref, data_urls, metadata, legacy_uuid)
+               VALUES (?, ?, ?, ?, ?, ?)""",
+            (
+                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
+        row = self._conn.execute(
+            "SELECT cid, name, schema_ref, data_urls, metadata, legacy_uuid "
+            "FROM dataset_entries WHERE cid = ?",
+            (cid,),
+        ).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
+        row = self._conn.execute(
+            "SELECT cid, name, schema_ref, data_urls, metadata, legacy_uuid "
+            "FROM dataset_entries WHERE name = ? LIMIT 1",
+            (name,),
+        ).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
+        cursor = self._conn.execute(
+            "SELECT cid, name, schema_ref, data_urls, metadata, legacy_uuid "
+            "FROM dataset_entries"
+        )
+        for row in cursor:
+            yield _row_to_entry(row)
+
+    # ------------------------------------------------------------------
+    # Schema operations
+    # ------------------------------------------------------------------
+
+    def store_schema(self, name: str, version: str, schema_json: str) -> None:
+        self._conn.execute(
+            """INSERT OR REPLACE INTO schemas (name, version, schema_json)
+               VALUES (?, ?, ?)""",
+            (name, version, schema_json),
+        )
+        self._conn.commit()
+
+    def get_schema_json(self, name: str, version: str) -> str | None:
+        row = self._conn.execute(
+            "SELECT schema_json FROM schemas WHERE name = ? AND version = ?",
+            (name, version),
+        ).fetchone()
+        if row is None:
+            return None
+        return row[0]
+
+    def iter_schemas(self) -> Iterator[tuple[str, str, str]]:
+        cursor = self._conn.execute("SELECT name, version, schema_json FROM schemas")
+        yield from cursor
+
+    def find_latest_version(self, name: str) -> str | None:
+        cursor = self._conn.execute(
+            "SELECT version FROM schemas WHERE name = ?",
+            (name,),
+        )
+        latest: tuple[int, int, int] | None = None
+        latest_str: str | None = None
+        for (version_str,) in cursor:
+            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 SQLite 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(data_urls_blob),
+        metadata=msgpack.unpackb(metadata_blob) if metadata_blob is not None else None,
+        _cid=cid,
+        _legacy_uuid=legacy_uuid,
+    )