feat: use psycopg pool to avoid conflicts in concurrent request (#42)

* feat: use psycopg pool to avoid conflicts in concurrent request

Signed-off-by: Keming <kemingyang@tensorchord.ai>

* only register vector type on conn creation

Signed-off-by: Keming <kemingyang@tensorchord.ai>

---------

Signed-off-by: Keming <kemingyang@tensorchord.ai>

Author: kemingy

pyproject.toml

@@ -20,6 +20,7 @@ dependencies = [
     "msgspec>=0.19.0",
     "numpy>=2.0.2",
     "pgvector>=0.3.6",
+    "psycopg-pool>=3.2.6",
     "psycopg[binary]>=3.2.3",
     "pypdfium2>=4.30.1",
     "pytrec-eval-terrier>=0.5.6",

tests/test_service.py

@@ -1,3 +1,4 @@
+import asyncio
 from collections.abc import AsyncIterator
 from http import HTTPStatus
 from uuid import UUID
@@ -94,3 +95,16 @@ async def test_service_pipeline(client):
     assert len(chunks) == len(text.split())
     for i, chunk in enumerate(chunks):
         assert chunk["text"] == f"num[{i + 1}]"
+
+
+@pytest.mark.parametrize("registry", [(DefaultDocument, DefaultChunk)], indirect=True)
+async def test_concurrent_db_transaction(client):
+    requests = [
+        client.simulate_post(
+            "/api/pipeline", json={"text": " ".join(map(str, range(i, i + 5)))}
+        )
+        for i in range(5)
+    ]
+    responses = await asyncio.gather(*requests)
+    for resp in responses:
+        assert resp.status_code == HTTPStatus.OK, resp.content

tests/test_table.py

@@ -1,3 +1,4 @@
+import asyncio
 from datetime import datetime, timezone
 from functools import partial
 from typing import Annotated, Optional
@@ -322,22 +323,27 @@ async def test_multivec_copy(registry):
 
 @pytest.mark.db
 @pytest.mark.parametrize("registry", [(Document, Chunk)], indirect=True)
-async def test_search_by_return(registry):
+async def test_search_return(registry):
     num = 100
     topk = 5
     await registry.insert(Document(text="hello world"))
+    chunks = []
     for i in range(num):
         text = f"hello {i}"
-        await registry.insert(
+        chunks.append(
             Chunk(doc_id=1, text=text, vector=gen_vector(), keyword=Keyword(text))
         )
+    await asyncio.gather(*[registry.insert(chunk) for chunk in chunks])
 
-    inserted = await registry.select_by(Chunk.partial_init(), fields=["text"])
+    inserted: list[Chunk] = await registry.select_by(
+        Chunk.partial_init(), fields=["text"]
+    )
     assert len(inserted) == num
-    for i, record in enumerate(inserted):
-        assert record.text == f"hello {i}"
-        # vector field is not selected
+    for record in inserted:
+        assert record.text.startswith("hello")
+        # vector field is not selected by default
         assert record.vector is msgspec.UNSET
+        assert record.keyword is msgspec.UNSET
 
     res = await registry.search_by_vector(Chunk, gen_vector(), topk=topk)
     assert len(res) == topk