Add ability to close connection pool in mem store

This adds an __enter__ and __exit__ function to the PostgresMemoryStore
so that it can be used as a context manager.
Also modify the integration tests to utilize the context manager,
so that connections are freed up after the tests are run.
Without this change, the tests will hold onto the connections
and cause the database to run out of connections.

Author: lossyrob

python/semantic_kernel/connectors/memory/postgres/postgres_memory_store.py

@@ -500,3 +500,13 @@ def __serialize_metadata(self, record: MemoryRecord) -> str:
             "description": record._description,
             "additional_metadata": record._additional_metadata,
         })
+
+    # Enable the connection pool to be closed when using as a context manager
+    def __enter__(self) -> "PostgresMemoryStore":
+        """Enter the runtime context."""
+        return self
+
+    def __exit__(self, exc_type, exc_value, traceback) -> bool:
+        """Exit the runtime context and dispose of the connection pool."""
+        self._connection_pool.close()
+        return False

python/tests/integration/connectors/memory/test_postgres_memory_store.py

@@ -46,166 +46,165 @@ def connection_string():
 
 
 def test_constructor(connection_string):
-    memory = PostgresMemoryStore(connection_string, 2, 1, 5)
-    assert memory._connection_pool is not None
+    with PostgresMemoryStore(connection_string, 2, 1, 5) as memory:
+        assert memory._connection_pool is not None
 
 
 @pytest.mark.asyncio
 async def test_create_and_does_collection_exist(connection_string):
-    memory = PostgresMemoryStore(connection_string, 2, 1, 5)
-    await memory.create_collection("test_collection")
-    result = await memory.does_collection_exist("test_collection")
-    assert result is not None
+    with PostgresMemoryStore(connection_string, 2, 1, 5) as memory:
+        await memory.create_collection("test_collection")
+        result = await memory.does_collection_exist("test_collection")
+        assert result is not None
 
 
 @pytest.mark.asyncio
 async def test_get_collections(connection_string):
-    memory = PostgresMemoryStore(connection_string, 2, 1, 5)
-
-    try:
-        await memory.create_collection("test_collection")
-        result = await memory.get_collections()
-        assert "test_collection" in result
-    except PoolTimeout:
-        pytest.skip("PoolTimeout exception raised, skipping test.")
+    with PostgresMemoryStore(connection_string, 2, 1, 5) as memory:
+        try:
+            await memory.create_collection("test_collection")
+            result = await memory.get_collections()
+            assert "test_collection" in result
+        except PoolTimeout:
+            pytest.skip("PoolTimeout exception raised, skipping test.")
 
 
 @pytest.mark.asyncio
 async def test_delete_collection(connection_string):
-    memory = PostgresMemoryStore(connection_string, 2, 1, 5)
-    try:
-        await memory.create_collection("test_collection")
+    with PostgresMemoryStore(connection_string, 2, 1, 5) as memory:
+        try:
+            await memory.create_collection("test_collection")
 
-        result = await memory.get_collections()
-        assert "test_collection" in result
+            result = await memory.get_collections()
+            assert "test_collection" in result
 
-        await memory.delete_collection("test_collection")
-        result = await memory.get_collections()
-        assert "test_collection" not in result
-    except PoolTimeout:
-        pytest.skip("PoolTimeout exception raised, skipping test.")
+            await memory.delete_collection("test_collection")
+            result = await memory.get_collections()
+            assert "test_collection" not in result
+        except PoolTimeout:
+            pytest.skip("PoolTimeout exception raised, skipping test.")
 
 
 @pytest.mark.asyncio
 async def test_does_collection_exist(connection_string):
-    memory = PostgresMemoryStore(connection_string, 2, 1, 5)
-    try:
-        await memory.create_collection("test_collection")
-        result = await memory.does_collection_exist("test_collection")
-        assert result is True
-    except PoolTimeout:
-        pytest.skip("PoolTimeout exception raised, skipping test.")
+    with PostgresMemoryStore(connection_string, 2, 1, 5) as memory:
+        try:
+            await memory.create_collection("test_collection")
+            result = await memory.does_collection_exist("test_collection")
+            assert result is True
+        except PoolTimeout:
+            pytest.skip("PoolTimeout exception raised, skipping test.")
 
 
 @pytest.mark.asyncio
 async def test_upsert_and_get(connection_string, memory_record1):
-    memory = PostgresMemoryStore(connection_string, 2, 1, 5)
-    try:
-        await memory.create_collection("test_collection")
-        await memory.upsert("test_collection", memory_record1)
-        result = await memory.get("test_collection", memory_record1._id, with_embedding=True)
-        assert result is not None
-        assert result._id == memory_record1._id
-        assert result._text == memory_record1._text
-        assert result._timestamp == memory_record1._timestamp
-        for i in range(len(result._embedding)):
-            assert result._embedding[i] == memory_record1._embedding[i]
-    except PoolTimeout:
-        pytest.skip("PoolTimeout exception raised, skipping test.")
+    with PostgresMemoryStore(connection_string, 2, 1, 5) as memory:
+        try:
+            await memory.create_collection("test_collection")
+            await memory.upsert("test_collection", memory_record1)
+            result = await memory.get("test_collection", memory_record1._id, with_embedding=True)
+            assert result is not None
+            assert result._id == memory_record1._id
+            assert result._text == memory_record1._text
+            assert result._timestamp == memory_record1._timestamp
+            for i in range(len(result._embedding)):
+                assert result._embedding[i] == memory_record1._embedding[i]
+        except PoolTimeout:
+            pytest.skip("PoolTimeout exception raised, skipping test.")
 
 
 @pytest.mark.asyncio
 async def test_upsert_batch_and_get_batch(connection_string, memory_record1, memory_record2):
-    memory = PostgresMemoryStore(connection_string, 2, 1, 5)
-    try:
-        await memory.create_collection("test_collection")
-        await memory.upsert_batch("test_collection", [memory_record1, memory_record2])
-
-        results = await memory.get_batch(
-            "test_collection",
-            [memory_record1._id, memory_record2._id],
-            with_embeddings=True,
-        )
-        assert len(results) == 2
-        assert results[0]._id in [memory_record1._id, memory_record2._id]
-        assert results[1]._id in [memory_record1._id, memory_record2._id]
-    except PoolTimeout:
-        pytest.skip("PoolTimeout exception raised, skipping test.")
+    with PostgresMemoryStore(connection_string, 2, 1, 5) as memory:
+        try:
+            await memory.create_collection("test_collection")
+            await memory.upsert_batch("test_collection", [memory_record1, memory_record2])
+
+            results = await memory.get_batch(
+                "test_collection",
+                [memory_record1._id, memory_record2._id],
+                with_embeddings=True,
+            )
+            assert len(results) == 2
+            assert results[0]._id in [memory_record1._id, memory_record2._id]
+            assert results[1]._id in [memory_record1._id, memory_record2._id]
+        except PoolTimeout:
+            pytest.skip("PoolTimeout exception raised, skipping test.")
 
 
 @pytest.mark.asyncio
 async def test_remove(connection_string, memory_record1):
-    memory = PostgresMemoryStore(connection_string, 2, 1, 5)
-    try:
-        await memory.create_collection("test_collection")
-        await memory.upsert("test_collection", memory_record1)
+    with PostgresMemoryStore(connection_string, 2, 1, 5) as memory:
+        try:
+            await memory.create_collection("test_collection")
+            await memory.upsert("test_collection", memory_record1)
 
-        result = await memory.get("test_collection", memory_record1._id, with_embedding=True)
-        assert result is not None
+            result = await memory.get("test_collection", memory_record1._id, with_embedding=True)
+            assert result is not None
 
-        await memory.remove("test_collection", memory_record1._id)
-        with pytest.raises(ServiceResourceNotFoundError):
-            await memory.get("test_collection", memory_record1._id, with_embedding=True)
-    except PoolTimeout:
-        pytest.skip("PoolTimeout exception raised, skipping test.")
+            await memory.remove("test_collection", memory_record1._id)
+            with pytest.raises(ServiceResourceNotFoundError):
+                await memory.get("test_collection", memory_record1._id, with_embedding=True)
+        except PoolTimeout:
+            pytest.skip("PoolTimeout exception raised, skipping test.")
 
 
 @pytest.mark.asyncio
 async def test_remove_batch(connection_string, memory_record1, memory_record2):
-    memory = PostgresMemoryStore(connection_string, 2, 1, 5)
-    try:
-        await memory.create_collection("test_collection")
-        await memory.upsert_batch("test_collection", [memory_record1, memory_record2])
-        await memory.remove_batch("test_collection", [memory_record1._id, memory_record2._id])
-        with pytest.raises(ServiceResourceNotFoundError):
-            _ = await memory.get("test_collection", memory_record1._id, with_embedding=True)
+    with PostgresMemoryStore(connection_string, 2, 1, 5) as memory:
+        try:
+            await memory.create_collection("test_collection")
+            await memory.upsert_batch("test_collection", [memory_record1, memory_record2])
+            await memory.remove_batch("test_collection", [memory_record1._id, memory_record2._id])
+            with pytest.raises(ServiceResourceNotFoundError):
+                _ = await memory.get("test_collection", memory_record1._id, with_embedding=True)
 
-        with pytest.raises(ServiceResourceNotFoundError):
-            _ = await memory.get("test_collection", memory_record2._id, with_embedding=True)
-    except PoolTimeout:
-        pytest.skip("PoolTimeout exception raised, skipping test.")
+            with pytest.raises(ServiceResourceNotFoundError):
+                _ = await memory.get("test_collection", memory_record2._id, with_embedding=True)
+        except PoolTimeout:
+            pytest.skip("PoolTimeout exception raised, skipping test.")
 
 
 @pytest.mark.asyncio
 async def test_get_nearest_match(connection_string, memory_record1, memory_record2):
-    memory = PostgresMemoryStore(connection_string, 2, 1, 5)
-    try:
-        await memory.create_collection("test_collection")
-        await memory.upsert_batch("test_collection", [memory_record1, memory_record2])
-        test_embedding = memory_record1.embedding.copy()
-        test_embedding[0] = test_embedding[0] + 0.01
-
-        result = await memory.get_nearest_match(
-            "test_collection", test_embedding, min_relevance_score=0.0, with_embedding=True
-        )
-        assert result is not None
-        assert result[0]._id == memory_record1._id
-        assert result[0]._text == memory_record1._text
-        assert result[0]._timestamp == memory_record1._timestamp
-        for i in range(len(result[0]._embedding)):
-            assert result[0]._embedding[i] == memory_record1._embedding[i]
-    except PoolTimeout:
-        pytest.skip("PoolTimeout exception raised, skipping test.")
+    with PostgresMemoryStore(connection_string, 2, 1, 5) as memory:
+        try:
+            await memory.create_collection("test_collection")
+            await memory.upsert_batch("test_collection", [memory_record1, memory_record2])
+            test_embedding = memory_record1.embedding.copy()
+            test_embedding[0] = test_embedding[0] + 0.01
+
+            result = await memory.get_nearest_match(
+                "test_collection", test_embedding, min_relevance_score=0.0, with_embedding=True
+            )
+            assert result is not None
+            assert result[0]._id == memory_record1._id
+            assert result[0]._text == memory_record1._text
+            assert result[0]._timestamp == memory_record1._timestamp
+            for i in range(len(result[0]._embedding)):
+                assert result[0]._embedding[i] == memory_record1._embedding[i]
+        except PoolTimeout:
+            pytest.skip("PoolTimeout exception raised, skipping test.")
 
 
 @pytest.mark.asyncio
 async def test_get_nearest_matches(connection_string, memory_record1, memory_record2, memory_record3):
-    memory = PostgresMemoryStore(connection_string, 2, 1, 5)
-    try:
-        await memory.create_collection("test_collection")
-        await memory.upsert_batch("test_collection", [memory_record1, memory_record2, memory_record3])
-        test_embedding = memory_record2.embedding
-        test_embedding[0] = test_embedding[0] + 0.025
-
-        result = await memory.get_nearest_matches(
-            "test_collection",
-            test_embedding,
-            limit=2,
-            min_relevance_score=0.0,
-            with_embeddings=True,
-        )
-        assert len(result) == 2
-        assert result[0][0]._id in [memory_record3._id, memory_record2._id]
-        assert result[1][0]._id in [memory_record3._id, memory_record2._id]
-    except PoolTimeout:
-        pytest.skip("PoolTimeout exception raised, skipping test.")
+    with PostgresMemoryStore(connection_string, 2, 1, 5) as memory:
+        try:
+            await memory.create_collection("test_collection")
+            await memory.upsert_batch("test_collection", [memory_record1, memory_record2, memory_record3])
+            test_embedding = memory_record2.embedding
+            test_embedding[0] = test_embedding[0] + 0.025
+
+            result = await memory.get_nearest_matches(
+                "test_collection",
+                test_embedding,
+                limit=2,
+                min_relevance_score=0.0,
+                with_embeddings=True,
+            )
+            assert len(result) == 2
+            assert result[0][0]._id in [memory_record3._id, memory_record2._id]
+            assert result[1][0]._id in [memory_record3._id, memory_record2._id]
+        except PoolTimeout:
+            pytest.skip("PoolTimeout exception raised, skipping test.")