chore: Pgvector - document expected connection string and raise informative error

Author: anakin87

integrations/pgvector/src/haystack_integrations/document_stores/pgvector/document_store.py

@@ -101,9 +101,9 @@ def __init__(
         A specific table to store Haystack documents will be created if it doesn't exist yet.
 
         :param connection_string: The connection string to use to connect to the PostgreSQL database, defined as an
-            environment variable. It can be provided in either URI format
-            e.g.: `PG_CONN_STR="postgresql://USER:PASSWORD@HOST:PORT/DB_NAME"`, or keyword/value format
-            e.g.: `PG_CONN_STR="host=HOST port=PORT dbname=DBNAME user=USER password=PASSWORD"`
+            environment variable. Supported formats:
+            - URI, e.g. `postgresql://USER:PASSWORD@HOST:PORT/DB_NAME` (use percent-encoding for special characters)
+            - keyword/value format, e.g. `PG_CONN_STR="host=HOST port=PORT dbname=DBNAME user=USER password=PASSWORD"`
             See [PostgreSQL Documentation](https://www.postgresql.org/docs/current/libpq-connect.html#LIBPQ-CONNSTRING)
             for more details.
         :param create_extension: Whether to create the pgvector extension if it doesn't exist.
@@ -360,7 +360,15 @@ def _ensure_db_setup(self):
                 logger.debug("Failed to close connection: {e}", e=str(e))
 
         conn_str = self.connection_string.resolve_value() or ""
-        connection = Connection.connect(conn_str)
+        try:
+            connection = Connection.connect(conn_str)
+        except Error as e:
+            msg = (
+                "Failed to connect to PostgreSQL database.  Ensure the connection string follows the "
+                "PostgreSQL connection specification: "
+                "https://www.postgresql.org/docs/current/libpq-connect.html#LIBPQ-CONNSTRING."
+            )
+            raise DocumentStoreError(msg) from e
         connection.autocommit = True
         if self.create_extension:
             connection.execute("CREATE EXTENSION IF NOT EXISTS vector")
@@ -394,7 +402,15 @@ async def _ensure_db_setup_async(self):
             await self._async_connection.close()
 
         conn_str = self.connection_string.resolve_value() or ""
-        async_connection = await AsyncConnection.connect(conn_str)
+        try:
+            async_connection = await AsyncConnection.connect(conn_str)
+        except Error as e:
+            msg = (
+                "Failed to connect to PostgreSQL database.  Ensure the connection string follows the "
+                "PostgreSQL connection specification: "
+                "https://www.postgresql.org/docs/current/libpq-connect.html#LIBPQ-CONNSTRING."
+            )
+            raise DocumentStoreError(msg) from e
         await async_connection.set_autocommit(True)
         if self.create_extension:
             await async_connection.execute("CREATE EXTENSION IF NOT EXISTS vector")

integrations/pgvector/tests/test_document_store.py

@@ -7,7 +7,7 @@
 import psycopg
 import pytest
 from haystack.dataclasses.document import ByteStream, Document
-from haystack.document_stores.errors import DuplicateDocumentError
+from haystack.document_stores.errors import DocumentStoreError, DuplicateDocumentError
 from haystack.document_stores.types import DuplicatePolicy
 from haystack.testing.document_store import CountDocumentsTest, DeleteDocumentsTest, WriteDocumentsTest
 from haystack.utils import Secret
@@ -60,6 +60,13 @@ def test_delete_all_documents(self, document_store: PgvectorDocumentStore) -> No
         document_store.write_documents([Document(id="1")])
         assert document_store.count_documents() == 1
 
+    def test_invalid_connection_string(self, monkeypatch):
+        monkeypatch.setenv("PG_CONN_STR", "invalid_connection_string")
+        document_store = PgvectorDocumentStore()
+        with pytest.raises(DocumentStoreError) as e:
+            document_store._ensure_db_setup()
+        assert "Failed to connect to PostgreSQL database" in str(e)
+
 
 @pytest.mark.usefixtures("patches_for_unit_tests")
 def test_init(monkeypatch):

integrations/pgvector/tests/test_document_store_async.py

@@ -7,7 +7,7 @@
 import psycopg
 import pytest
 from haystack.dataclasses.document import ByteStream, Document
-from haystack.document_stores.errors import DuplicateDocumentError
+from haystack.document_stores.errors import DocumentStoreError, DuplicateDocumentError
 from haystack.document_stores.types import DuplicatePolicy
 from haystack.utils import Secret
 from psycopg.sql import SQL
@@ -101,6 +101,13 @@ async def test_delete_all_documents_async(self, document_store: PgvectorDocument
         document_store.write_documents([Document(id="1")])
         assert document_store.count_documents() == 1
 
+    async def test_invalid_connection_string(self, monkeypatch):
+        monkeypatch.setenv("PG_CONN_STR", "invalid_connection_string")
+        document_store = PgvectorDocumentStore()
+        with pytest.raises(DocumentStoreError) as e:
+            await document_store._ensure_db_setup_async()
+        assert "Failed to connect to PostgreSQL database" in str(e)
+
 
 @pytest.mark.integration
 @pytest.mark.asyncio