agentmesh-platform 1.0.0a1__py3-none-any.whl → 1.0.0a2__py3-none-any.whl

agentmesh/storage/memory_provider.py

@@ -0,0 +1,232 @@
+"""
+In-Memory Storage Provider.
+
+Simple in-memory implementation for development and testing.
+"""
+
+from typing import Optional
+import asyncio
+from collections import defaultdict
+
+from .provider import AbstractStorageProvider, StorageConfig
+
+
+class MemoryStorageProvider(AbstractStorageProvider):
+    """
+    In-memory storage provider.
+    
+    Uses Python dictionaries for storage. Data is lost on restart.
+    Suitable for development and testing only.
+    """
+    
+    def __init__(self, config: StorageConfig):
+        """Initialize in-memory storage."""
+        super().__init__(config)
+        self._data: dict[str, str] = {}
+        self._hashes: dict[str, dict[str, str]] = defaultdict(dict)
+        self._lists: dict[str, list[str]] = defaultdict(list)
+        self._sorted_sets: dict[str, dict[str, float]] = defaultdict(dict)
+        self._ttls: dict[str, float] = {}
+        self._connected = False
+    
+    async def connect(self) -> None:
+        """Establish connection (no-op for memory)."""
+        self._connected = True
+    
+    async def disconnect(self) -> None:
+        """Close connection (no-op for memory)."""
+        self._connected = False
+    
+    async def health_check(self) -> bool:
+        """Check if storage is healthy."""
+        return self._connected
+    
+    # Key-Value Operations
+    
+    async def get(self, key: str) -> Optional[str]:
+        """Get value by key."""
+        return self._data.get(key)
+    
+    async def set(
+        self,
+        key: str,
+        value: str,
+        ttl_seconds: Optional[int] = None,
+    ) -> bool:
+        """Set value with optional TTL."""
+        self._data[key] = value
+        if ttl_seconds is not None:
+            import time
+            self._ttls[key] = time.time() + ttl_seconds
+        return True
+    
+    async def delete(self, key: str) -> bool:
+        """Delete key."""
+        if key in self._data:
+            del self._data[key]
+            return True
+        return False
+    
+    async def exists(self, key: str) -> bool:
+        """Check if key exists."""
+        return key in self._data
+    
+    # Hash Operations
+    
+    async def hget(self, key: str, field: str) -> Optional[str]:
+        """Get hash field value."""
+        return self._hashes.get(key, {}).get(field)
+    
+    async def hset(self, key: str, field: str, value: str) -> bool:
+        """Set hash field value."""
+        self._hashes[key][field] = value
+        return True
+    
+    async def hgetall(self, key: str) -> dict[str, str]:
+        """Get all hash fields."""
+        return dict(self._hashes.get(key, {}))
+    
+    async def hdel(self, key: str, field: str) -> bool:
+        """Delete hash field."""
+        if key in self._hashes and field in self._hashes[key]:
+            del self._hashes[key][field]
+            return True
+        return False
+    
+    async def hkeys(self, key: str) -> list[str]:
+        """Get all hash field names."""
+        return list(self._hashes.get(key, {}).keys())
+    
+    # List Operations
+    
+    async def lpush(self, key: str, value: str) -> int:
+        """Push value to head of list."""
+        self._lists[key].insert(0, value)
+        return len(self._lists[key])
+    
+    async def rpush(self, key: str, value: str) -> int:
+        """Push value to tail of list."""
+        self._lists[key].append(value)
+        return len(self._lists[key])
+    
+    async def lrange(self, key: str, start: int, stop: int) -> list[str]:
+        """Get list range [start, stop]."""
+        lst = self._lists.get(key, [])
+        if stop == -1:
+            return lst[start:]
+        return lst[start:stop + 1]
+    
+    async def llen(self, key: str) -> int:
+        """Get list length."""
+        return len(self._lists.get(key, []))
+    
+    # Sorted Set Operations
+    
+    async def zadd(self, key: str, score: float, member: str) -> bool:
+        """Add member to sorted set with score."""
+        self._sorted_sets[key][member] = score
+        return True
+    
+    async def zscore(self, key: str, member: str) -> Optional[float]:
+        """Get score of member in sorted set."""
+        return self._sorted_sets.get(key, {}).get(member)
+    
+    async def zrange(
+        self,
+        key: str,
+        start: int,
+        stop: int,
+        with_scores: bool = False,
+    ) -> list[str] | list[tuple[str, float]]:
+        """Get sorted set range."""
+        sorted_set = self._sorted_sets.get(key, {})
+        sorted_items = sorted(sorted_set.items(), key=lambda x: x[1])
+        
+        if stop == -1:
+            items = sorted_items[start:]
+        else:
+            items = sorted_items[start:stop + 1]
+        
+        if with_scores:
+            return items
+        return [member for member, _ in items]
+    
+    async def zrangebyscore(
+        self,
+        key: str,
+        min_score: float,
+        max_score: float,
+        with_scores: bool = False,
+    ) -> list[str] | list[tuple[str, float]]:
+        """Get sorted set range by score."""
+        sorted_set = self._sorted_sets.get(key, {})
+        items = [
+            (member, score)
+            for member, score in sorted(sorted_set.items(), key=lambda x: x[1])
+            if min_score <= score <= max_score
+        ]
+        
+        if with_scores:
+            return items
+        return [member for member, _ in items]
+    
+    # Atomic Operations
+    
+    async def incr(self, key: str) -> int:
+        """Increment value atomically."""
+        current = int(self._data.get(key, "0"))
+        new_value = current + 1
+        self._data[key] = str(new_value)
+        return new_value
+    
+    async def decr(self, key: str) -> int:
+        """Decrement value atomically."""
+        current = int(self._data.get(key, "0"))
+        new_value = current - 1
+        self._data[key] = str(new_value)
+        return new_value
+    
+    async def incrby(self, key: str, amount: int) -> int:
+        """Increment value by amount."""
+        current = int(self._data.get(key, "0"))
+        new_value = current + amount
+        self._data[key] = str(new_value)
+        return new_value
+    
+    # Batch Operations
+    
+    async def mget(self, keys: list[str]) -> list[Optional[str]]:
+        """Get multiple values."""
+        return [self._data.get(key) for key in keys]
+    
+    async def mset(self, mapping: dict[str, str]) -> bool:
+        """Set multiple key-value pairs."""
+        self._data.update(mapping)
+        return True
+    
+    # Pattern Operations
+    
+    async def keys(self, pattern: str) -> list[str]:
+        """Get keys matching pattern."""
+        import fnmatch
+        return [key for key in self._data.keys() if fnmatch.fnmatch(key, pattern)]
+    
+    async def scan(
+        self,
+        cursor: int = 0,
+        match: Optional[str] = None,
+        count: int = 100,
+    ) -> tuple[int, list[str]]:
+        """Scan keys with cursor."""
+        all_keys = list(self._data.keys())
+        
+        if match:
+            import fnmatch
+            all_keys = [key for key in all_keys if fnmatch.fnmatch(key, match)]
+        
+        start = cursor
+        end = cursor + count
+        keys = all_keys[start:end]
+        
+        new_cursor = end if end < len(all_keys) else 0
+        return new_cursor, keys

agentmesh/storage/postgres_provider.py

@@ -0,0 +1,463 @@
+"""
+PostgreSQL Storage Provider.
+
+Enterprise-grade PostgreSQL backend with async SQLAlchemy ORM.
+"""
+
+from typing import Optional
+import json
+
+from .provider import AbstractStorageProvider, StorageConfig
+
+
+class PostgresStorageProvider(AbstractStorageProvider):
+    """
+    PostgreSQL storage provider.
+    
+    Features:
+    - Async SQLAlchemy ORM
+    - Connection pooling
+    - JSONB support for structured data
+    - Full ACID compliance
+    
+    Requires: sqlalchemy[asyncio], asyncpg packages
+    """
+    
+    def __init__(self, config: StorageConfig):
+        """Initialize PostgreSQL storage."""
+        super().__init__(config)
+        self._engine = None
+        self._session_factory = None
+    
+    async def connect(self) -> None:
+        """Establish connection to PostgreSQL."""
+        try:
+            from sqlalchemy.ext.asyncio import (
+                create_async_engine,
+                async_sessionmaker,
+            )
+        except ImportError:
+            raise ImportError(
+                "sqlalchemy[asyncio] and asyncpg packages are required for PostgresStorageProvider. "
+                "Install with: pip install sqlalchemy[asyncio] asyncpg"
+            )
+        
+        # Build connection string
+        if self.config.connection_string:
+            conn_str = self.config.connection_string
+        else:
+            password_part = (
+                f":{self.config.postgres_password}"
+                if self.config.postgres_password
+                else ""
+            )
+            conn_str = (
+                f"postgresql+asyncpg://{self.config.postgres_user}"
+                f"{password_part}@{self.config.postgres_host}"
+                f":{self.config.postgres_port}/{self.config.postgres_database}"
+            )
+            
+            if self.config.postgres_ssl_mode != "disable":
+                conn_str += f"?ssl={self.config.postgres_ssl_mode}"
+        
+        # Create engine
+        self._engine = create_async_engine(
+            conn_str,
+            pool_size=self.config.pool_size,
+            max_overflow=20,
+            pool_pre_ping=True,
+            echo=False,
+        )
+        
+        self._session_factory = async_sessionmaker(
+            self._engine,
+            expire_on_commit=False,
+        )
+        
+        # Initialize schema
+        await self._init_schema()
+    
+    async def _init_schema(self) -> None:
+        """Initialize database schema."""
+        from sqlalchemy import text
+        
+        # Create tables for key-value, hashes, lists, etc.
+        async with self._engine.begin() as conn:
+            await conn.execute(text(
+                """
+                CREATE TABLE IF NOT EXISTS agentmesh_kv (
+                    key VARCHAR(512) PRIMARY KEY,
+                    value TEXT NOT NULL,
+                    expires_at TIMESTAMP
+                );
+                CREATE INDEX IF NOT EXISTS idx_kv_expires ON agentmesh_kv(expires_at);
+                
+                CREATE TABLE IF NOT EXISTS agentmesh_hash (
+                    key VARCHAR(512) NOT NULL,
+                    field VARCHAR(512) NOT NULL,
+                    value TEXT NOT NULL,
+                    PRIMARY KEY (key, field)
+                );
+                
+                CREATE TABLE IF NOT EXISTS agentmesh_list (
+                    key VARCHAR(512) NOT NULL,
+                    idx INTEGER NOT NULL,
+                    value TEXT NOT NULL,
+                    PRIMARY KEY (key, idx)
+                );
+                CREATE INDEX IF NOT EXISTS idx_list_key ON agentmesh_list(key, idx);
+                
+                CREATE TABLE IF NOT EXISTS agentmesh_zset (
+                    key VARCHAR(512) NOT NULL,
+                    member VARCHAR(512) NOT NULL,
+                    score DOUBLE PRECISION NOT NULL,
+                    PRIMARY KEY (key, member)
+                );
+                CREATE INDEX IF NOT EXISTS idx_zset_score ON agentmesh_zset(key, score);
+                """
+            ))
+    
+    async def disconnect(self) -> None:
+        """Close connection to PostgreSQL."""
+        if self._engine:
+            await self._engine.dispose()
+    
+    async def health_check(self) -> bool:
+        """Check if PostgreSQL is healthy."""
+        try:
+            if self._engine:
+                from sqlalchemy import text
+                async with self._engine.begin() as conn:
+                    await conn.execute(text("SELECT 1"))
+                return True
+        except Exception:
+            pass
+        return False
+    
+    # Key-Value Operations
+    
+    async def get(self, key: str) -> Optional[str]:
+        """Get value by key."""
+        async with self._session_factory() as session:
+            result = await session.execute(
+                "SELECT value FROM agentmesh_kv WHERE key = :key "
+                "AND (expires_at IS NULL OR expires_at > NOW())",
+                {"key": key},
+            )
+            row = result.fetchone()
+            return row[0] if row else None
+    
+    async def set(
+        self,
+        key: str,
+        value: str,
+        ttl_seconds: Optional[int] = None,
+    ) -> bool:
+        """Set value with optional TTL."""
+        async with self._session_factory() as session:
+            if ttl_seconds:
+                await session.execute(
+                    "INSERT INTO agentmesh_kv (key, value, expires_at) "
+                    "VALUES (:key, :value, NOW() + INTERVAL '1 second' * :ttl) "
+                    "ON CONFLICT (key) DO UPDATE SET value = :value, expires_at = NOW() + INTERVAL '1 second' * :ttl",
+                    {"key": key, "value": value, "ttl": ttl_seconds},
+                )
+            else:
+                await session.execute(
+                    "INSERT INTO agentmesh_kv (key, value, expires_at) "
+                    "VALUES (:key, :value, NULL) "
+                    "ON CONFLICT (key) DO UPDATE SET value = :value, expires_at = NULL",
+                    {"key": key, "value": value},
+                )
+            await session.commit()
+        return True
+    
+    async def delete(self, key: str) -> bool:
+        """Delete key."""
+        async with self._session_factory() as session:
+            result = await session.execute(
+                "DELETE FROM agentmesh_kv WHERE key = :key",
+                {"key": key},
+            )
+            await session.commit()
+            return result.rowcount > 0
+    
+    async def exists(self, key: str) -> bool:
+        """Check if key exists."""
+        result = await self.get(key)
+        return result is not None
+    
+    # Hash Operations
+    
+    async def hget(self, key: str, field: str) -> Optional[str]:
+        """Get hash field value."""
+        async with self._session_factory() as session:
+            result = await session.execute(
+                "SELECT value FROM agentmesh_hash WHERE key = :key AND field = :field",
+                {"key": key, "field": field},
+            )
+            row = result.fetchone()
+            return row[0] if row else None
+    
+    async def hset(self, key: str, field: str, value: str) -> bool:
+        """Set hash field value."""
+        async with self._session_factory() as session:
+            await session.execute(
+                "INSERT INTO agentmesh_hash (key, field, value) "
+                "VALUES (:key, :field, :value) "
+                "ON CONFLICT (key, field) DO UPDATE SET value = :value",
+                {"key": key, "field": field, "value": value},
+            )
+            await session.commit()
+        return True
+    
+    async def hgetall(self, key: str) -> dict[str, str]:
+        """Get all hash fields."""
+        async with self._session_factory() as session:
+            result = await session.execute(
+                "SELECT field, value FROM agentmesh_hash WHERE key = :key",
+                {"key": key},
+            )
+            return {row[0]: row[1] for row in result.fetchall()}
+    
+    async def hdel(self, key: str, field: str) -> bool:
+        """Delete hash field."""
+        async with self._session_factory() as session:
+            result = await session.execute(
+                "DELETE FROM agentmesh_hash WHERE key = :key AND field = :field",
+                {"key": key, "field": field},
+            )
+            await session.commit()
+            return result.rowcount > 0
+    
+    async def hkeys(self, key: str) -> list[str]:
+        """Get all hash field names."""
+        async with self._session_factory() as session:
+            result = await session.execute(
+                "SELECT field FROM agentmesh_hash WHERE key = :key",
+                {"key": key},
+            )
+            return [row[0] for row in result.fetchall()]
+    
+    # List Operations
+    
+    async def lpush(self, key: str, value: str) -> int:
+        """Push value to head of list."""
+        async with self._session_factory() as session:
+            # Shift all indices up
+            await session.execute(
+                "UPDATE agentmesh_list SET idx = idx + 1 WHERE key = :key",
+                {"key": key},
+            )
+            # Insert at position 0
+            await session.execute(
+                "INSERT INTO agentmesh_list (key, idx, value) VALUES (:key, 0, :value)",
+                {"key": key, "value": value},
+            )
+            await session.commit()
+            # Get new length
+            result = await session.execute(
+                "SELECT COUNT(*) FROM agentmesh_list WHERE key = :key",
+                {"key": key},
+            )
+            return result.scalar()
+    
+    async def rpush(self, key: str, value: str) -> int:
+        """Push value to tail of list."""
+        async with self._session_factory() as session:
+            # Get max index
+            result = await session.execute(
+                "SELECT COALESCE(MAX(idx), -1) FROM agentmesh_list WHERE key = :key",
+                {"key": key},
+            )
+            max_idx = result.scalar()
+            # Insert at end
+            await session.execute(
+                "INSERT INTO agentmesh_list (key, idx, value) VALUES (:key, :idx, :value)",
+                {"key": key, "idx": max_idx + 1, "value": value},
+            )
+            await session.commit()
+            return max_idx + 2
+    
+    async def lrange(self, key: str, start: int, stop: int) -> list[str]:
+        """Get list range [start, stop]."""
+        async with self._session_factory() as session:
+            if stop == -1:
+                result = await session.execute(
+                    "SELECT value FROM agentmesh_list WHERE key = :key AND idx >= :start ORDER BY idx",
+                    {"key": key, "start": start},
+                )
+            else:
+                result = await session.execute(
+                    "SELECT value FROM agentmesh_list WHERE key = :key AND idx >= :start AND idx <= :stop ORDER BY idx",
+                    {"key": key, "start": start, "stop": stop},
+                )
+            return [row[0] for row in result.fetchall()]
+    
+    async def llen(self, key: str) -> int:
+        """Get list length."""
+        async with self._session_factory() as session:
+            result = await session.execute(
+                "SELECT COUNT(*) FROM agentmesh_list WHERE key = :key",
+                {"key": key},
+            )
+            return result.scalar()
+    
+    # Sorted Set Operations
+    
+    async def zadd(self, key: str, score: float, member: str) -> bool:
+        """Add member to sorted set with score."""
+        async with self._session_factory() as session:
+            await session.execute(
+                "INSERT INTO agentmesh_zset (key, member, score) "
+                "VALUES (:key, :member, :score) "
+                "ON CONFLICT (key, member) DO UPDATE SET score = :score",
+                {"key": key, "member": member, "score": score},
+            )
+            await session.commit()
+        return True
+    
+    async def zscore(self, key: str, member: str) -> Optional[float]:
+        """Get score of member in sorted set."""
+        async with self._session_factory() as session:
+            result = await session.execute(
+                "SELECT score FROM agentmesh_zset WHERE key = :key AND member = :member",
+                {"key": key, "member": member},
+            )
+            row = result.fetchone()
+            return row[0] if row else None
+    
+    async def zrange(
+        self,
+        key: str,
+        start: int,
+        stop: int,
+        with_scores: bool = False,
+    ) -> list[str] | list[tuple[str, float]]:
+        """Get sorted set range."""
+        async with self._session_factory() as session:
+            if stop == -1:
+                result = await session.execute(
+                    "SELECT member, score FROM agentmesh_zset WHERE key = :key "
+                    "ORDER BY score OFFSET :start",
+                    {"key": key, "start": start},
+                )
+            else:
+                result = await session.execute(
+                    "SELECT member, score FROM agentmesh_zset WHERE key = :key "
+                    "ORDER BY score LIMIT :limit OFFSET :start",
+                    {"key": key, "start": start, "limit": stop - start + 1},
+                )
+            rows = result.fetchall()
+            if with_scores:
+                return [(row[0], row[1]) for row in rows]
+            return [row[0] for row in rows]
+    
+    async def zrangebyscore(
+        self,
+        key: str,
+        min_score: float,
+        max_score: float,
+        with_scores: bool = False,
+    ) -> list[str] | list[tuple[str, float]]:
+        """Get sorted set range by score."""
+        async with self._session_factory() as session:
+            result = await session.execute(
+                "SELECT member, score FROM agentmesh_zset "
+                "WHERE key = :key AND score >= :min AND score <= :max ORDER BY score",
+                {"key": key, "min": min_score, "max": max_score},
+            )
+            rows = result.fetchall()
+            if with_scores:
+                return [(row[0], row[1]) for row in rows]
+            return [row[0] for row in rows]
+    
+    # Atomic Operations
+    
+    async def incr(self, key: str) -> int:
+        """Increment value atomically."""
+        return await self.incrby(key, 1)
+    
+    async def decr(self, key: str) -> int:
+        """Decrement value atomically."""
+        return await self.incrby(key, -1)
+    
+    async def incrby(self, key: str, amount: int) -> int:
+        """Increment value by amount."""
+        async with self._session_factory() as session:
+            # Use PostgreSQL's atomic UPDATE ... RETURNING
+            result = await session.execute(
+                "INSERT INTO agentmesh_kv (key, value) VALUES (:key, :amount) "
+                "ON CONFLICT (key) DO UPDATE SET value = "
+                "(CAST(agentmesh_kv.value AS INTEGER) + :amount)::TEXT "
+                "RETURNING CAST(value AS INTEGER)",
+                {"key": key, "amount": str(amount)},
+            )
+            await session.commit()
+            return result.scalar()
+    
+    # Batch Operations
+    
+    async def mget(self, keys: list[str]) -> list[Optional[str]]:
+        """Get multiple values."""
+        async with self._session_factory() as session:
+            result = await session.execute(
+                "SELECT key, value FROM agentmesh_kv WHERE key = ANY(:keys) "
+                "AND (expires_at IS NULL OR expires_at > NOW())",
+                {"keys": keys},
+            )
+            values_dict = {row[0]: row[1] for row in result.fetchall()}
+            return [values_dict.get(key) for key in keys]
+    
+    async def mset(self, mapping: dict[str, str]) -> bool:
+        """Set multiple key-value pairs."""
+        async with self._session_factory() as session:
+            for key, value in mapping.items():
+                await session.execute(
+                    "INSERT INTO agentmesh_kv (key, value) VALUES (:key, :value) "
+                    "ON CONFLICT (key) DO UPDATE SET value = :value",
+                    {"key": key, "value": value},
+                )
+            await session.commit()
+        return True
+    
+    # Pattern Operations
+    
+    async def keys(self, pattern: str) -> list[str]:
+        """Get keys matching pattern."""
+        # Convert glob pattern to SQL LIKE pattern
+        sql_pattern = pattern.replace("*", "%").replace("?", "_")
+        async with self._session_factory() as session:
+            result = await session.execute(
+                "SELECT key FROM agentmesh_kv WHERE key LIKE :pattern "
+                "AND (expires_at IS NULL OR expires_at > NOW())",
+                {"pattern": sql_pattern},
+            )
+            return [row[0] for row in result.fetchall()]
+    
+    async def scan(
+        self,
+        cursor: int = 0,
+        match: Optional[str] = None,
+        count: int = 100,
+    ) -> tuple[int, list[str]]:
+        """Scan keys with cursor."""
+        async with self._session_factory() as session:
+            if match:
+                sql_pattern = match.replace("*", "%").replace("?", "_")
+                result = await session.execute(
+                    "SELECT key FROM agentmesh_kv WHERE key LIKE :pattern "
+                    "AND (expires_at IS NULL OR expires_at > NOW()) "
+                    "ORDER BY key LIMIT :count OFFSET :cursor",
+                    {"pattern": sql_pattern, "count": count, "cursor": cursor},
+                )
+            else:
+                result = await session.execute(
+                    "SELECT key FROM agentmesh_kv "
+                    "WHERE expires_at IS NULL OR expires_at > NOW() "
+                    "ORDER BY key LIMIT :count OFFSET :cursor",
+                    {"count": count, "cursor": cursor},
+                )
+            keys = [row[0] for row in result.fetchall()]
+            new_cursor = cursor + count if len(keys) == count else 0
+            return new_cursor, keys