Fix pgvector sslmode handling for PostgreSQL URIs (#5308)

mem0/vector_stores/pgvector.py

@@ -1,7 +1,9 @@
 import json
 import logging
+import re
 from contextlib import contextmanager
 from typing import Any, List, Optional
+from urllib.parse import parse_qsl, urlencode, urlsplit
 
 from pydantic import BaseModel
 
@@ -113,6 +115,24 @@ def _build_filter_conditions(filters):
     return conditions, params
 
 
+def _with_sslmode(connection_string: str, sslmode: str) -> str:
+    """Add or replace sslmode in URI and keyword conninfo strings.
+
+    Keyword conninfo values are assumed not to contain nested ``sslmode=``
+    substrings, such as inside an ``options`` value.
+    """
+    if "://" in connection_string:
+        parsed = urlsplit(connection_string)
+        query = [(key, value) for key, value in parse_qsl(parsed.query, keep_blank_values=True) if key != "sslmode"]
+        query.append(("sslmode", sslmode))
+        return parsed._replace(query=urlencode(query)).geturl()
+
+    if re.search(r"(^|\s)sslmode=", connection_string):
+        return re.sub(r"(^|\s)sslmode=\S+", lambda match: f"{match.group(1)}sslmode={sslmode}", connection_string)
+
+    return f"{connection_string} sslmode={sslmode}"
+
+
 class OutputData(BaseModel):
     id: Optional[str]
     score: Optional[float]
@@ -168,18 +188,11 @@ class PGVector(VectorStoreBase):
             self.connection_pool = connection_pool
         elif connection_string:
             if sslmode:
-                # Append sslmode to connection string if provided
-                if 'sslmode=' in connection_string:
-                    # Replace existing sslmode
-                    import re
-                    connection_string = re.sub(r'sslmode=[^ ]*', f'sslmode={sslmode}', connection_string)
-                else:
-                    # Add sslmode to connection string
-                    connection_string = f"{connection_string} sslmode={sslmode}"
+                connection_string = _with_sslmode(connection_string, sslmode)
         else:
             connection_string = f"postgresql://{user}:{password}@{host}:{port}/{dbname}"
             if sslmode:
-                connection_string = f"{connection_string} sslmode={sslmode}"
+                connection_string = _with_sslmode(connection_string, sslmode)
         
         if self.connection_pool is None:
             if PSYCOPG_VERSION == 3:

tests/vector_stores/test_pgvector.py

@@ -4,7 +4,7 @@ import unittest
 import uuid
 from unittest.mock import MagicMock, patch
 
-from mem0.vector_stores.pgvector import PGVector, _build_filter_conditions
+from mem0.vector_stores.pgvector import PGVector, _build_filter_conditions, _with_sslmode
 
 
 class TestPGVector(unittest.TestCase):
@@ -2084,8 +2084,8 @@ class TestPGVector(unittest.TestCase):
             connection_string=connection_string
         )
         
-        # Verify ConnectionPool was called with the connection string including sslmode
-        expected_conn_string = f"{connection_string} sslmode=require"
+        # Verify ConnectionPool was called with sslmode as a URI query parameter
+        expected_conn_string = f"{connection_string}?sslmode=require"
         mock_connection_pool.assert_called_with(
             conninfo=expected_conn_string,
             min_size=1,
@@ -2095,6 +2095,36 @@ class TestPGVector(unittest.TestCase):
         self.assertEqual(pgvector.collection_name, "test_collection")
         self.assertEqual(pgvector.embedding_model_dims, 3)
 
+    def test_with_sslmode_appends_uri_query_param(self):
+        """Test sslmode is appended to URI connection strings without corrupting dbname."""
+        connection_string = _with_sslmode(
+            "postgresql://user:pass@localhost:5432/db?connect_timeout=10",
+            "require",
+        )
+
+        self.assertEqual(
+            connection_string,
+            "postgresql://user:pass@localhost:5432/db?connect_timeout=10&sslmode=require",
+        )
+
+    def test_with_sslmode_replaces_existing_uri_sslmode(self):
+        """Test existing URI sslmode values are replaced rather than duplicated."""
+        connection_string = _with_sslmode(
+            "postgresql://user:pass@localhost:5432/db?sslmode=prefer&connect_timeout=10",
+            "require",
+        )
+
+        self.assertEqual(
+            connection_string,
+            "postgresql://user:pass@localhost:5432/db?connect_timeout=10&sslmode=require",
+        )
+
+    def test_with_sslmode_preserves_keyword_conninfo_format(self):
+        """Test non-URI conninfo strings keep PostgreSQL keyword syntax."""
+        connection_string = _with_sslmode("dbname=test user=postgres sslmode=prefer", "require")
+
+        self.assertEqual(connection_string, "dbname=test user=postgres sslmode=require")
+
     # Enhanced Test for Index Creation with DiskANN
     @patch('mem0.vector_stores.pgvector.PSYCOPG_VERSION', 3)
     @patch('mem0.vector_stores.pgvector.ConnectionPool')