feat: postgresql task store

LaurentMnr95 · LaurentMnr95 · commit 433bb5eadcfa · 2025-05-21T11:41:29.000+02:00
diff --git a/docker/postgres/docker-compose.yml b/docker/postgres/docker-compose.yml
@@ -0,0 +1,28 @@
+version: "3.8"
+
+services:
+  postgres:
+    image: postgres:15-alpine
+    ports:
+      - "5432:5432"
+    environment:
+      - POSTGRES_USER=postgres
+      - POSTGRES_PASSWORD=postgres
+      - POSTGRES_DB=a2a_test
+    volumes:
+      - postgres_data:/var/lib/postgresql/data
+      - ./docker/postgres/init.sql:/docker-entrypoint-initdb.d/init.sql
+    networks:
+      - a2a-network
+    healthcheck:
+      test: ["CMD-SHELL", "pg_isready -U postgres"]
+      interval: 5s
+      timeout: 5s
+      retries: 5
+
+volumes:
+  postgres_data:
+
+networks:
+  a2a-network:
+    driver: bridge
diff --git a/docker/postgres/init.sql b/docker/postgres/init.sql
@@ -0,0 +1,8 @@
+-- Create a dedicated user for the application
+CREATE USER a2a WITH PASSWORD 'a2a_password';
+
+-- Create the tasks database
+CREATE DATABASE a2a_tasks;
+
+GRANT ALL PRIVILEGES ON DATABASE a2a_test TO a2a;
+
diff --git a/src/a2a/server/tasks/postgresql_task_store.py b/src/a2a/server/tasks/postgresql_task_store.py
@@ -0,0 +1,136 @@
+import json
+import logging
+
+import asyncpg
+
+from a2a.server.tasks.task_store import TaskStore
+from a2a.types import Task
+
+
+logger = logging.getLogger(__name__)
+
+
+class PostgreSQLTaskStore(TaskStore):
+    """PostgreSQL implementation of TaskStore.
+
+    Stores task objects in a PostgreSQL database.
+    """
+
+    def __init__(
+        self,
+        url: str,
+        table_name: str = 'tasks',
+        create_table: bool = True,
+    ) -> None:
+        """Initializes the PostgreSQLTaskStore.
+
+        Args:
+            url: PostgreSQL connection string in the format:
+                postgresql://username:password@hostname:port/database
+            table_name: The name of the table to store tasks in
+            create_table: Whether to create the table if it doesn't exist
+        """
+        logger.debug('Initializing PostgreSQLTaskStore')
+        self.url = url
+        self.table_name = table_name
+        self.create_table = create_table
+
+        self.pool: asyncpg.Pool | None = None
+
+    async def initialize(self) -> None:
+        """Initialize the database connection pool and create the table
+        if needed.
+        """
+        if self.pool is not None:
+            return
+
+        logger.debug('Creating connection pool')
+        self.pool = await asyncpg.create_pool(self.url)
+
+        if self.create_table:
+            async with self.pool.acquire() as conn:
+                logger.debug('Creating tasks table if not exists')
+                await conn.execute(
+                    f"""
+                    CREATE TABLE IF NOT EXISTS {self.table_name} (
+                        id TEXT PRIMARY KEY,
+                        data JSONB NOT NULL
+
+                    )
+                    """
+                )
+
+    async def close(self) -> None:
+        """Close the database connection pool."""
+        if self.pool is not None:
+            await self.pool.close()
+            self.pool = None
+
+    async def save(self, task: Task) -> None:
+        """Saves or updates a task in the PostgreSQL store."""
+        await self._ensure_initialized()
+
+        assert self.pool is not None
+        async with self.pool.acquire() as conn, conn.transaction():
+            task_json = task.model_dump()
+
+            await conn.execute(
+                f"""
+                INSERT INTO {self.table_name} (id, data)
+                VALUES ($1, $2)
+                ON CONFLICT (id) DO UPDATE
+                SET data = $2
+                """,
+                task.id,
+                json.dumps(task_json),
+            )
+
+            logger.info('Task %s saved successfully.', task.id)
+
+    async def get(self, task_id: str) -> Task | None:
+        """Retrieves a task from the PostgreSQL store by ID."""
+        await self._ensure_initialized()
+
+        assert self.pool is not None
+        async with self.pool.acquire() as conn, conn.transaction():
+            logger.debug('Attempting to get task with id: %s', task_id)
+
+            row = await conn.fetchrow(
+                f'SELECT data FROM {self.table_name} WHERE id = $1',
+                task_id,
+            )
+
+            if row:
+                task_json = json.loads(row['data'])
+                task = Task.model_validate(task_json)
+                logger.debug('Task %s retrieved successfully.', task_id)
+                return task
+
+            logger.debug('Task %s not found in store.', task_id)
+            return None
+
+    async def delete(self, task_id: str) -> None:
+        """Deletes a task from the PostgreSQL store by ID."""
+        await self._ensure_initialized()
+
+        assert self.pool is not None
+        async with self.pool.acquire() as conn, conn.transaction():
+            logger.debug('Attempting to delete task with id: %s', task_id)
+
+            result = await conn.execute(
+                f'DELETE FROM {self.table_name} WHERE id = $1',
+                task_id,
+            )
+
+            if result.split()[-1] != '0':  # Check if rows were affected
+                logger.info('Task %s deleted successfully.', task_id)
+            else:
+                logger.warning(
+                    'Attempted to delete nonexistent task with id: %s',
+                    task_id,
+                )
+
+    async def _ensure_initialized(self) -> None:
+        """Ensure the database connection is initialized."""
+        if self.pool is None:
+            await self.initialize()
diff --git a/tests/server/tasks/test_postgresql_task_store.py b/tests/server/tasks/test_postgresql_task_store.py
@@ -0,0 +1,207 @@
+import json
+import os
+from typing import AsyncGenerator
+
+import pytest
+import pytest_asyncio
+
+from a2a.server.tasks.postgresql_task_store import PostgreSQLTaskStore
+from a2a.types import Task, TaskState, TaskStatus
+
+# Use a proper Task object instead of a dict for the minimal task
+task_status = TaskStatus(state=TaskState.submitted)
+MINIMAL_TASK_OBJ = Task(
+    id='task-abc',
+    contextId='session-xyz',
+    status=task_status,
+    kind='task',
+)
+
+
+# Get PostgreSQL connection string from environment or use a default for testing
+POSTGRES_TEST_DSN = os.environ.get(
+    'POSTGRES_TEST_DSN',
+    'postgresql://postgres:postgres@localhost:5432/a2a_test',
+)
+
+
+@pytest_asyncio.fixture
+async def postgres_store() -> AsyncGenerator[PostgreSQLTaskStore, None]:
+    """Fixture that provides a PostgreSQLTaskStore connected to a real database.
+
+    This fixture requires a running PostgreSQL instance
+    """
+
+    store = PostgreSQLTaskStore(POSTGRES_TEST_DSN)
+    await store.initialize()
+
+    # Clean up any test data that might be left from previous runs
+    if store.pool is not None:
+        async with store.pool.acquire() as conn:
+            await conn.execute(
+                f"DELETE FROM {store.table_name} WHERE id LIKE 'test-%'"
+            )
+
+    yield store
+    await store.close()
+
+
+@pytest.mark.asyncio
+async def test_initialize_creates_table(
+    postgres_store: PostgreSQLTaskStore,
+) -> None:
+    """Test that initialize creates the table if it doesn't exist."""
+    await postgres_store.initialize()
+
+    # Verify the pool was created
+    assert postgres_store.pool is not None
+
+    # Verify the table creation query was executed
+    async with postgres_store.pool.acquire() as conn:
+        async with conn.transaction():
+            exists = await conn.fetchval(
+                f"SELECT EXISTS (SELECT 1 FROM pg_tables WHERE tablename = '{postgres_store.table_name}')"
+            )
+            assert exists
+
+
+@pytest.mark.asyncio
+async def test_save_task(postgres_store: PostgreSQLTaskStore) -> None:
+    """Test saving a task to the PostgreSQL store."""
+    # Use the pre-created Task object to avoid serialization issues
+    task = MINIMAL_TASK_OBJ
+    await postgres_store.save(task)
+    assert postgres_store.pool is not None
+
+    # Verify the insert query was executed
+    async with postgres_store.pool.acquire() as conn:
+        async with conn.transaction():
+            row = await conn.fetchrow(
+                f'SELECT data FROM {postgres_store.table_name} WHERE id = $1',
+                task.id,
+            )
+            assert row is not None
+            # Convert the task to a dictionary with proper enum handling
+
+            # Parse the JSON string from the database
+            db_dict = (
+                json.loads(row['data'])
+                if isinstance(row['data'], str)
+                else row['data']
+            )
+            assert db_dict == task.model_dump()
+
+
+@pytest.mark.asyncio
+async def test_get_task(postgres_store: PostgreSQLTaskStore) -> None:
+    """Test retrieving a task from the PostgreSQL store."""
+    retrieved_task = await postgres_store.get(MINIMAL_TASK_OBJ.id)
+
+    # Verify the task was correctly reconstructed
+    assert retrieved_task is not None
+    assert retrieved_task.id == MINIMAL_TASK_OBJ.id
+    assert retrieved_task.contextId == MINIMAL_TASK_OBJ.contextId
+
+
+@pytest.mark.asyncio
+async def test_get_nonexistent_task(
+    postgres_store: PostgreSQLTaskStore,
+) -> None:
+    """Test retrieving a nonexistent task."""
+
+    retrieved_task = await postgres_store.get('nonexistent')
+
+    # Verify None was returned
+    assert retrieved_task is None
+
+
+@pytest.mark.asyncio
+async def test_delete_task(
+    postgres_store: PostgreSQLTaskStore,
+) -> None:
+    """Test deleting a task from the PostgreSQL store."""
+    await postgres_store.initialize()
+    await postgres_store.delete(MINIMAL_TASK_OBJ.id)
+
+
+@pytest.mark.asyncio
+async def test_delete_nonexistent_task(
+    postgres_store: PostgreSQLTaskStore,
+) -> None:
+    """Test deleting a nonexistent task."""
+    await postgres_store.initialize()
+    await postgres_store.delete('nonexistent')
+
+
+@pytest.mark.asyncio
+async def test_close_connection_pool(
+    postgres_store: PostgreSQLTaskStore,
+) -> None:
+    """Test closing the database connection pool."""
+    await postgres_store.close()
+    assert postgres_store.pool is None
+
+
+@pytest.mark.asyncio
+async def test_save_and_get_task(
+    postgres_store: PostgreSQLTaskStore,
+) -> None:
+    """Test for saving and retrieving a task from a real PostgreSQL database."""
+    # Create a unique test task
+    test_task = Task(
+        id='test-1',
+        contextId='test-session-1',
+        status=TaskStatus(state=TaskState.submitted),
+        kind='task',
+    )
+
+    # Save the task
+    await postgres_store.save(test_task)
+
+    # Retrieve the task
+    retrieved_task = await postgres_store.get(test_task.id)
+
+    # Verify task was retrieved correctly
+    assert retrieved_task is not None
+    assert retrieved_task.id == test_task.id
+    assert retrieved_task.contextId == test_task.contextId
+    assert retrieved_task.status.state == test_task.status.state
+
+    # Clean up
+    await postgres_store.delete(test_task.id)
+
+    # Verify deletion
+    deleted_task = await postgres_store.get(test_task.id)
+    assert deleted_task is None
+
+
+@pytest.mark.asyncio
+async def test_update_task(
+    postgres_store: PostgreSQLTaskStore,
+) -> None:
+    """Test for updating a task in a real PostgreSQL database."""
+    # Create a test task
+    test_task = Task(
+        id='test-2',
+        contextId='test-session-2',
+        status=TaskStatus(state=TaskState.submitted),
+        kind='task',
+    )
+
+    # Save the task
+    await postgres_store.save(test_task)
+
+    # Update the task
+    updated_task = test_task.model_copy(deep=True)
+    updated_task.status.state = TaskState.completed
+    await postgres_store.save(updated_task)
+
+    # Retrieve the updated task
+    retrieved_task = await postgres_store.get(test_task.id)
+
+    # Verify the update was successful
+    assert retrieved_task is not None
+    assert retrieved_task.status.state == TaskState.completed
+
+    # Clean up
+    await postgres_store.delete(test_task.id)
