src/agent_server/pg_storage.py - master - man4j/agent-server

src/agent_server/pg_storage.py
98 строк3 KB
Vladimir
initial
18 апр 2026, 19:36
18 апр 2026, 19:36a95f341
from __future__ import annotations

import asyncpg
from typing import Any, Dict, Union
from urllib.parse import quote

from chainlit.data.storage_clients.base import BaseStorageClient


class PostgresStorageClient(BaseStorageClient):
    def __init__(self, dsn: str, public_base_url: str, route_prefix: str = "/storage"):
        self._dsn = dsn
        self._public_base_url = public_base_url.rstrip("/")
        self._route_prefix = route_prefix.rstrip("/")
        self._pool: asyncpg.Pool | None = None

    async def _get_pool(self) -> asyncpg.Pool:
        if self._pool is None:
            self._pool = await asyncpg.create_pool(self._dsn, min_size=1, max_size=5)
        return self._pool

    async def upload_file(
        self,
        object_key: str,
        data: Union[bytes, str],
        mime: str = "application/octet-stream",
        overwrite: bool = True,
        content_disposition: str | None = None,
    ) -> Dict[str, Any]:
        payload = data.encode("utf-8") if isinstance(data, str) else data
        pool = await self._get_pool()

        async with pool.acquire() as conn:
            if overwrite:
                await conn.execute(
                    """
                    INSERT INTO chainlit_storage_objects (object_key, mime, content, updated_at)
                    VALUES ($1, $2, $3, now())
                    ON CONFLICT (object_key) DO UPDATE
                    SET mime = EXCLUDED.mime,
                        content = EXCLUDED.content,
                        updated_at = now()
                    """,
                    object_key,
                    mime,
                    payload,
                )
            else:
                await conn.execute(
                    """
                    INSERT INTO chainlit_storage_objects (object_key, mime, content)
                    VALUES ($1, $2, $3)
                    ON CONFLICT DO NOTHING
                    """,
                    object_key,
                    mime,
                    payload,
                )

        return {
            "object_key": object_key,
            "url": f"{self._public_base_url}{self._route_prefix}/{quote(object_key, safe='')}",
        }

    async def delete_file(self, object_key: str) -> bool:
        pool = await self._get_pool()
        async with pool.acquire() as conn:
            result = await conn.execute(
                "DELETE FROM chainlit_storage_objects WHERE object_key = $1",
                object_key,
            )
        # asyncpg returns strings like "DELETE 1"
        return result.endswith("1")

    async def get_read_url(self, object_key: str) -> str:
        return f"{self._public_base_url}{self._route_prefix}/{quote(object_key, safe='')}"

    async def fetch_file(self, object_key: str) -> tuple[bytes, str] | None:
        pool = await self._get_pool()
        async with pool.acquire() as conn:
            row = await conn.fetchrow(
                """
                SELECT content, mime
                FROM chainlit_storage_objects
                WHERE object_key = $1
                """,
                object_key,
            )

        if not row:
            return None

        return bytes(row["content"]), row["mime"]

    async def close(self) -> None:
        if self._pool is not None:
            await self._pool.close()
            self._pool = None