From 31264728854d266a541327e333f103469e14336d Mon Sep 17 00:00:00 2001 From: Rob Emanuele Date: Fri, 3 Jan 2025 17:43:48 -0500 Subject: [PATCH] Allow AZURE_DB_FOR_POSTGRES_CONNECTION_STRING --- python/.env.example | 1 + .../azure_db_for_postgres_settings.py | 13 ++++++++++++- .../memory/azure_db_for_postgres/constants.py | 8 ++++++++ 3 files changed, 21 insertions(+), 1 deletion(-) diff --git a/python/.env.example b/python/.env.example index 8e46ec2bb6de8..ef532e6dd80d3 100644 --- a/python/.env.example +++ b/python/.env.example @@ -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="" diff --git a/python/semantic_kernel/connectors/memory/azure_db_for_postgres/azure_db_for_postgres_settings.py b/python/semantic_kernel/connectors/memory/azure_db_for_postgres/azure_db_for_postgres_settings.py index ea5e61e646b6a..6fcc1c8d987dd 100644 --- a/python/semantic_kernel/connectors/memory/azure_db_for_postgres/azure_db_for_postgres_settings.py +++ b/python/semantic_kernel/connectors/memory/azure_db_for_postgres/azure_db_for_postgres_settings.py @@ -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 diff --git a/python/semantic_kernel/connectors/memory/azure_db_for_postgres/constants.py b/python/semantic_kernel/connectors/memory/azure_db_for_postgres/constants.py index 612b938173e10..55cdea88a3de8 100644 --- a/python/semantic_kernel/connectors/memory/azure_db_for_postgres/constants.py +++ b/python/semantic_kernel/connectors/memory/azure_db_for_postgres/constants.py @@ -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. +"""