Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Python: Add Azure DB for PostgreSQL vector store #9326

Draft
wants to merge 9 commits into
base: main
Choose a base branch
from
Draft
Prev Previous commit
Next Next commit
Allow AZURE_DB_FOR_POSTGRES_CONNECTION_STRING
lossyrob committed Jan 7, 2025
commit 1a8ad95805b1c4150597f327ae0cb31b8669d26a
1 change: 1 addition & 0 deletions python/.env.example
Original file line number Diff line number Diff line change
@@ -17,6 +17,7 @@ MONGODB_ATLAS_CONNECTION_STRING=""
PINECONE_API_KEY=""
PINECONE_ENVIRONMENT=""
POSTGRES_CONNECTION_STRING=""
AZURE_DB_FOR_POSTGRES_CONNECTION_STRING=""
WEAVIATE_URL=""
WEAVIATE_API_KEY=""
GOOGLE_SEARCH_ENGINE_ID=""
Original file line number Diff line number Diff line change
@@ -4,7 +4,11 @@

from psycopg.conninfo import conninfo_to_dict
from psycopg_pool import AsyncConnectionPool
from pydantic import Field, SecretStr

from semantic_kernel.connectors.memory.azure_db_for_postgres.constants import (
AZURE_DB_FOR_POSTGRES_CONNECTION_STRING_ENV_VAR,
)
from semantic_kernel.connectors.memory.azure_db_for_postgres.entra_connection import AsyncEntraConnection
from semantic_kernel.exceptions.memory_connector_exceptions import MemoryConnectorInitializationError

@@ -30,6 +34,12 @@ class AzureDBForPostgresSettings(PostgresSettings):

credential: AsyncTokenCredential | TokenCredential | None = None

azure_db_connection_string: SecretStr | None = Field(None, alias=AZURE_DB_FOR_POSTGRES_CONNECTION_STRING_ENV_VAR)
"""A azure db specific connection string. Can be supplied instead of POSTGRES_CONNECTION_STRING
This is useful if settings for both an Azure DB and a regular Postgres database are needed.
"""

def get_connection_args(self, **kwargs) -> dict[str, Any]:
"""Get connection arguments.
@@ -40,7 +50,8 @@ def get_connection_args(self, **kwargs) -> dict[str, Any]:
Returns:
dict[str, Any]: Connection arguments that can be passed to psycopg.connect
"""
result = conninfo_to_dict(self.connection_string.get_secret_value()) if self.connection_string else {}
conn_string_setting = self.azure_db_connection_string or self.connection_string
result = conninfo_to_dict(conn_string_setting.get_secret_value()) if conn_string_setting else {}

if self.host:
result["host"] = self.host
Original file line number Diff line number Diff line change
@@ -1,3 +1,11 @@
# Copyright (c) Microsoft. All rights reserved.

AZURE_DB_FOR_POSTGRES_SCOPE = "https://ossrdbms-aad.database.windows.net/.default"

AZURE_DB_FOR_POSTGRES_CONNECTION_STRING_ENV_VAR = "AZURE_DB_FOR_POSTGRES_CONNECTION_STRING"
"""Azure DB for Postgres specific environment variable for the connection string.
This is useful if settings for both an Azure DB and a regular Postgres database are needed.
If not set, the regular POSTGRES_CONNECTION_STRING environment variable or other standard
Postgres environment variables will be used.
"""