chore: import upstream snapshot with attribution

python/tests/integration/memory/postgres/test_postgres_int.py

@@ -0,0 +1,259 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+import uuid
+from collections.abc import AsyncGenerator, Sequence
+from contextlib import asynccontextmanager
+from typing import Annotated, Any
+
+import pandas as pd
+import pytest
+import pytest_asyncio
+from pydantic import BaseModel
+
+from semantic_kernel.connectors.postgres import PostgresCollection, PostgresSettings, PostgresStore
+from semantic_kernel.data.vector import (
+    DistanceFunction,
+    IndexKind,
+    VectorStoreCollectionDefinition,
+    VectorStoreField,
+    vectorstoremodel,
+)
+from semantic_kernel.exceptions.memory_connector_exceptions import (
+    MemoryConnectorConnectionException,
+    MemoryConnectorInitializationError,
+)
+
+try:
+    import psycopg  # noqa: F401
+    import psycopg_pool  # noqa: F401
+
+    psycopg_pool_installed = True
+except ImportError:
+    psycopg_pool_installed = False
+
+pg_settings: PostgresSettings = PostgresSettings()
+try:
+    connection_params_present = any(pg_settings.get_connection_args().values())
+except MemoryConnectorInitializationError:
+    connection_params_present = False
+
+pytestmark = pytest.mark.skipif(
+    not (psycopg_pool_installed or connection_params_present),
+    reason="psycopg_pool is not installed" if not psycopg_pool_installed else "No connection parameters provided",
+)
+
+
+@vectorstoremodel
+class SimpleDataModel(BaseModel):
+    id: Annotated[int, VectorStoreField("key")]
+    embedding: Annotated[
+        list[float] | str | None,
+        VectorStoreField(
+            "vector",
+            index_kind=IndexKind.HNSW,
+            dimensions=3,
+            distance_function=DistanceFunction.COSINE_SIMILARITY,
+        ),
+    ] = None
+    data: Annotated[
+        dict[str, Any],
+        VectorStoreField("data", type="JSONB"),
+    ]
+
+    def model_post_init(self, context: Any) -> None:
+        if self.embedding is None:
+            self.embedding = self.data
+
+
+def DataModelPandas(record) -> tuple:
+    definition = VectorStoreCollectionDefinition(
+        fields=[
+            VectorStoreField(
+                "vector",
+                name="embedding",
+                index_kind="hnsw",
+                dimensions=3,
+                distance_function="cosine_similarity",
+                type="float",
+            ),
+            VectorStoreField("key", name="id", type="int"),
+            VectorStoreField("data", name="data", type="dict"),
+        ],
+        container_mode=True,
+        to_dict=lambda x: x.to_dict(orient="records"),
+        from_dict=lambda x, **_: pd.DataFrame(x),
+    )
+    df = pd.DataFrame([record])
+    return definition, df
+
+
+@pytest_asyncio.fixture
+async def vector_store() -> AsyncGenerator[PostgresStore, None]:
+    try:
+        async with await pg_settings.create_connection_pool() as pool:
+            yield PostgresStore(connection_pool=pool)
+    except MemoryConnectorConnectionException:
+        pytest.skip("Postgres connection not available")
+        yield None
+        return
+
+
+@asynccontextmanager
+async def create_simple_collection(
+    vector_store: PostgresStore,
+) -> AsyncGenerator[PostgresCollection[int, SimpleDataModel], None]:
+    """Returns a collection with a unique name that is deleted after the context.
+
+    This can be moved to use a fixture with scope=function and loop_scope=session
+    after upgrade to pytest-asyncio 0.24. With the current version, the fixture
+    would both cache and use the event loop of the declared scope.
+    """
+    suffix = str(uuid.uuid4()).replace("-", "")[:8]
+    collection_id = f"test_collection_{suffix}"
+    collection = vector_store.get_collection(collection_name=collection_id, record_type=SimpleDataModel)
+    assert isinstance(collection, PostgresCollection)
+    await collection.ensure_collection_exists()
+    try:
+        yield collection
+    finally:
+        await collection.ensure_collection_deleted()
+
+
+def test_create_store(vector_store):
+    assert vector_store is not None
+    assert vector_store.connection_pool is not None
+
+
+async def test_ensure_collection_exists_exists_and_delete(vector_store: PostgresStore):
+    suffix = str(uuid.uuid4()).replace("-", "")[:8]
+
+    collection = vector_store.get_collection(collection_name=f"test_collection_{suffix}", record_type=SimpleDataModel)
+
+    does_exist_1 = await collection.collection_exists()
+    assert does_exist_1 is False
+
+    await collection.ensure_collection_exists()
+    does_exist_2 = await collection.collection_exists()
+    assert does_exist_2 is True
+
+    await collection.ensure_collection_deleted()
+    does_exist_3 = await collection.collection_exists()
+    assert does_exist_3 is False
+
+
+async def test_list_collection_names(vector_store):
+    async with create_simple_collection(vector_store) as simple_collection:
+        simple_collection_id = simple_collection.collection_name
+        result = await vector_store.list_collection_names()
+        assert simple_collection_id in result
+
+
+async def test_upsert_get_and_delete(vector_store: PostgresStore):
+    record = SimpleDataModel(id=1, embedding=[1.1, 2.2, 3.3], data={"key": "value"})
+    async with create_simple_collection(vector_store) as simple_collection:
+        result_before_upsert = await simple_collection.get(1)
+        assert result_before_upsert is None
+
+        await simple_collection.upsert(record)
+        result = await simple_collection.get(1)
+        assert result is not None
+        assert result.id == record.id
+        assert result.embedding == record.embedding
+        assert result.data == record.data
+
+        # Check that the table has an index
+        connection_pool = simple_collection.connection_pool
+        async with connection_pool.connection() as conn, conn.cursor() as cur:
+            await cur.execute(
+                "SELECT indexname FROM pg_indexes WHERE tablename = %s", (simple_collection.collection_name,)
+            )
+            rows = await cur.fetchall()
+            index_names = [index[0] for index in rows]
+            assert any("embedding_idx" in index_name for index_name in index_names)
+
+        await simple_collection.delete(1)
+        result_after_delete = await simple_collection.get(1)
+        assert result_after_delete is None
+
+
+async def test_upsert_get_and_delete_pandas(vector_store):
+    record = SimpleDataModel(id=1, embedding=[1.1, 2.2, 3.3], data={"key": "value"})
+    definition, df = DataModelPandas(record.model_dump())
+
+    suffix = str(uuid.uuid4()).replace("-", "")[:8]
+    collection = vector_store.get_collection(
+        collection_name=f"test_collection_{suffix}",
+        record_type=pd.DataFrame,
+        definition=definition,
+    )
+    await collection.ensure_collection_exists()
+
+    try:
+        result_before_upsert = await collection.get(1)
+        assert result_before_upsert is None
+
+        await collection.upsert(df)
+        result: pd.DataFrame = await collection.get(1)
+        assert result is not None
+        row = result.iloc[0]
+        assert row.id == record.id
+        assert row.embedding == record.embedding
+        assert row.data == record.data
+
+        await collection.delete(1)
+        result_after_delete = await collection.get(1)
+        assert result_after_delete is None
+    finally:
+        await collection.ensure_collection_deleted()
+
+
+async def test_upsert_get_and_delete_multiple(vector_store: PostgresStore):
+    async with create_simple_collection(vector_store) as simple_collection:
+        record1 = SimpleDataModel(id=1, embedding=[1.1, 2.2, 3.3], data={"key": "value"})
+        record2 = SimpleDataModel(id=2, embedding=[4.4, 5.5, 6.6], data={"key": "value"})
+
+        result_before_upsert = await simple_collection.get([1, 2])
+        assert result_before_upsert is None
+
+        await simple_collection.upsert([record1, record2])
+        # Test get for the two existing keys and one non-existing key;
+        # this should return only the two existing records.
+        result = await simple_collection.get([1, 2, 3])
+        assert result is not None
+        assert isinstance(result, Sequence)
+        assert len(result) == 2
+        assert result[0] is not None
+        assert result[0].id == record1.id
+        assert result[0].embedding == record1.embedding
+        assert result[0].data == record1.data
+        assert result[1] is not None
+        assert result[1].id == record2.id
+        assert result[1].embedding == record2.embedding
+        assert result[1].data == record2.data
+
+        await simple_collection.delete([1, 2])
+        result_after_delete = await simple_collection.get([1, 2])
+        assert result_after_delete is None
+
+
+async def test_search(vector_store: PostgresStore):
+    async with create_simple_collection(vector_store) as simple_collection:
+        records = [
+            SimpleDataModel(id=1, embedding=[1.0, 0.0, 0.0], data={"key": "value1"}),
+            SimpleDataModel(id=2, embedding=[0.8, 0.2, 0.0], data={"key": "value2"}),
+            SimpleDataModel(id=3, embedding=[0.6, 0.0, 0.4], data={"key": "value3"}),
+            SimpleDataModel(id=4, embedding=[1.0, 1.0, 0.0], data={"key": "value4"}),
+            SimpleDataModel(id=5, embedding=[0.0, 1.0, 1.0], data={"key": "value5"}),
+            SimpleDataModel(id=6, embedding=[1.0, 0.0, 1.0], data={"key": "value6"}),
+        ]
+
+        await simple_collection.upsert(records)
+
+        try:
+            search_results = await simple_collection.search(vector=[1.0, 0.0, 0.0], top=3, include_total_count=True)
+            assert search_results is not None
+            assert search_results.total_count == 3
+            assert {result.record.id async for result in search_results.results} == {1, 2, 3}
+
+        finally:
+            await simple_collection.delete([r.id for r in records])