diff --git a/.github/workflows/pr-dotnet.yml b/.github/workflows/pr-dotnet.yml index 8285f99..725bd2b 100644 --- a/.github/workflows/pr-dotnet.yml +++ b/.github/workflows/pr-dotnet.yml @@ -36,11 +36,11 @@ jobs: - name: Test if: ${{ always() }} working-directory: dotnet - run: dotnet test tests/Azure/Data/Postgresql/Npgsql/Azure.Data.Postgresql.Npgsql.Tests.csproj --configuration Release --logger trx --results-directory TestResults + run: dotnet test tests/Microsoft/Azure/PostgreSQL/Auth/Microsoft.Azure.PostgreSQL.Auth.csproj --configuration Release --framework net${{ matrix.dotnet-version }} --logger trx --results-directory TestResults - name: Pack working-directory: dotnet - run: dotnet pack --no-build --configuration Release --output nupkgs + run: dotnet pack src/Microsoft/Azure/PostgreSQL/Auth/Microsoft.Azure.PostgreSQL.Auth.csproj --no-build --configuration Release --output nupkgs - name: Upload test results if: always() diff --git a/dotnet/README.md b/dotnet/README.md index 1e3ecb8..474942e 100644 --- a/dotnet/README.md +++ b/dotnet/README.md @@ -41,6 +41,8 @@ Use the extension methods as needed: ### Asynchronous Authentication (Recommended) ```csharp +using Azure.Identity; + // Fill in with connection information to Azure PostgreSQL server // Note: No username/password in connection string - authentication handled by Entra ID var connectionString = "Host=myserver.postgres.database.azure.com;Database=mydb;Port=5432;SSL Mode=Require;"; @@ -51,17 +53,21 @@ var dataSourceBuilder = new NpgsqlDataSourceBuilder(connectionString); // - Detect the current Azure identity (managed identity, service principal, or user) // - Acquire a PostgreSQL-scoped access token // - Configure the connection to use token-based authentication -await dataSourceBuilder.UseEntraAuthenticationAsync(); +var credential = new DefaultAzureCredential(); +await dataSourceBuilder.UseEntraAuthenticationAsync(credential); ``` ### Synchronous Authentication ```csharp +using Azure.Identity; + // Fill in with connection information to Azure PostgreSQL server var connectionString = "Host=myserver.postgres.database.azure.com;Database=mydb;Port=5432;SSL Mode=Require;"; var dataSourceBuilder = new NpgsqlDataSourceBuilder(connectionString); // Use the sync extension method for Entra authentication -dataSourceBuilder.UseEntraAuthentication(); +var credential = new DefaultAzureCredential(); +dataSourceBuilder.UseEntraAuthentication(credential); ``` ## Benefits diff --git a/dotnet/samples/GettingStarted/CreateDbConnectionNpgsql.cs b/dotnet/samples/GettingStarted/CreateDbConnectionNpgsql.cs index 94797fc..8bacc27 100644 --- a/dotnet/samples/GettingStarted/CreateDbConnectionNpgsql.cs +++ b/dotnet/samples/GettingStarted/CreateDbConnectionNpgsql.cs @@ -1,6 +1,7 @@ // Copyright (c) Microsoft Corporation. // Licensed under the MIT License. +using Azure.Identity; using Npgsql; using Microsoft.Azure.PostgreSQL.Auth; using Microsoft.Extensions.Configuration; @@ -49,15 +50,15 @@ private static async Task ExecuteQueriesWithEntraAuth(string connectionString, b var dataSourceBuilder = new NpgsqlDataSourceBuilder(connectionString); // Here, we use the appropriate extension method provided by NpgsqlDataSourceBuilderExtensions.cs - // to enable Entra Authentication. This will handle token acquisition, username extraction, and - // token refresh as needed. + // to enable Entra Authentication. This will handle username extraction and token refresh as needed. + var credential = new DefaultAzureCredential(); if (useAsync) { - await dataSourceBuilder.UseEntraAuthenticationAsync(); + await dataSourceBuilder.UseEntraAuthenticationAsync(credential); } else { - dataSourceBuilder.UseEntraAuthentication(); + dataSourceBuilder.UseEntraAuthentication(credential); } using var dataSource = dataSourceBuilder.Build(); diff --git a/dotnet/samples/GettingStarted/GettingStarted.csproj b/dotnet/samples/GettingStarted/GettingStarted.csproj index d931dee..5349e7a 100644 --- a/dotnet/samples/GettingStarted/GettingStarted.csproj +++ b/dotnet/samples/GettingStarted/GettingStarted.csproj @@ -14,6 +14,7 @@ + diff --git a/dotnet/src/Microsoft/Azure/PostgreSQL/Auth/EntraIdExtension.cs b/dotnet/src/Microsoft/Azure/PostgreSQL/Auth/EntraIdExtension.cs index 4711ba9..0f8e687 100644 --- a/dotnet/src/Microsoft/Azure/PostgreSQL/Auth/EntraIdExtension.cs +++ b/dotnet/src/Microsoft/Azure/PostgreSQL/Auth/EntraIdExtension.cs @@ -4,7 +4,6 @@ using System.Text; using System.Text.Json; using Azure.Core; -using Azure.Identity; using Npgsql; namespace Microsoft.Azure.PostgreSQL.Auth; @@ -25,12 +24,11 @@ public static class EntraIdExtension /// Configures the NpgsqlDataSourceBuilder to use Entra ID authentication synchronously. /// /// The NpgsqlDataSourceBuilder to configure. - /// The TokenCredential to use for authentication. If null, DefaultAzureCredential is used. + /// The TokenCredential to use for authentication. /// A cancellation token that can be used to cancel the operation. /// The configured NpgsqlDataSourceBuilder. - public static NpgsqlDataSourceBuilder UseEntraAuthentication(this NpgsqlDataSourceBuilder dataSourceBuilder, TokenCredential? credential = default, CancellationToken cancellationToken = default) + public static NpgsqlDataSourceBuilder UseEntraAuthentication(this NpgsqlDataSourceBuilder dataSourceBuilder, TokenCredential credential, CancellationToken cancellationToken = default) { - credential ??= new DefaultAzureCredential(); if (dataSourceBuilder.ConnectionStringBuilder.Username == null) { @@ -60,12 +58,11 @@ public static NpgsqlDataSourceBuilder UseEntraAuthentication(this NpgsqlDataSour /// Configures the NpgsqlDataSourceBuilder to use Entra ID authentication asynchronously. /// /// The NpgsqlDataSourceBuilder to configure. - /// The TokenCredential to use for authentication. If null, DefaultAzureCredential is used. + /// The TokenCredential to use for authentication. /// A cancellation token that can be used to cancel the operation. /// A task that represents the asynchronous operation. The task result contains the configured NpgsqlDataSourceBuilder. - public static async Task UseEntraAuthenticationAsync(this NpgsqlDataSourceBuilder dataSourceBuilder, TokenCredential? credential = default, CancellationToken cancellationToken = default) + public static async Task UseEntraAuthenticationAsync(this NpgsqlDataSourceBuilder dataSourceBuilder, TokenCredential credential, CancellationToken cancellationToken = default) { - credential ??= new DefaultAzureCredential(); if (dataSourceBuilder.ConnectionStringBuilder.Username == null) { diff --git a/dotnet/src/Microsoft/Azure/PostgreSQL/Auth/Microsoft.Azure.PostgreSQL.Auth.csproj b/dotnet/src/Microsoft/Azure/PostgreSQL/Auth/Microsoft.Azure.PostgreSQL.Auth.csproj index 8374df3..12f7335 100644 --- a/dotnet/src/Microsoft/Azure/PostgreSQL/Auth/Microsoft.Azure.PostgreSQL.Auth.csproj +++ b/dotnet/src/Microsoft/Azure/PostgreSQL/Auth/Microsoft.Azure.PostgreSQL.Auth.csproj @@ -26,7 +26,6 @@ - diff --git a/javascript/README.md b/javascript/README.md index 646ee1b..862213d 100644 --- a/javascript/README.md +++ b/javascript/README.md @@ -60,14 +60,17 @@ The pg driver integration provides connection support with Azure Entra ID authen ```javascript import { Pool } from 'pg'; -import { getPassword } from 'azure-postgresql-auth'; +import { DefaultAzureCredential } from '@azure/identity'; +import { getEntraTokenPassword } from 'azure-postgresql-auth'; + +const credential = new DefaultAzureCredential(); const pool = new Pool({ host: process.env.PGHOST, port: process.env.PGPORT, database: process.env.PGDATABASE, user: process.env.PGUSER, - password: getPassword, // Dynamic password function + password: () => getEntraTokenPassword(credential), // Dynamic password function ssl: { rejectUnauthorized: true }, connectionTimeoutMillis: 10000, idleTimeoutMillis: 30000, @@ -86,8 +89,11 @@ Sequelize integration uses pg as the backend driver with automatic Entra ID auth ```javascript import { Sequelize } from 'sequelize'; +import { DefaultAzureCredential } from '@azure/identity'; import { configureEntraIdAuth } from 'azure-postgresql-auth'; +const credential = new DefaultAzureCredential(); + const sequelize = new Sequelize({ dialect: 'postgres', host: process.env.PGHOST, @@ -104,7 +110,7 @@ const sequelize = new Sequelize({ }); // Configure Entra ID authentication -configureEntraIdAuth(sequelize, { +configureEntraIdAuth(sequelize, credential, { fallbackUsername: 'my-db-user' // Optional fallback username }); @@ -138,11 +144,16 @@ GRANT ALL PRIVILEGES ON DATABASE your_database TO "your-user@your-domain.com"; **Connection Timeouts** ```javascript +import { DefaultAzureCredential } from '@azure/identity'; +import { getEntraTokenPassword } from 'postgres-entra-auth'; + +const credential = new DefaultAzureCredential(); + // Increase connection timeout for slow networks const pool = new Pool({ host: process.env.PGHOST, database: process.env.PGDATABASE, - password: getPassword, + password: () => getEntraTokenPassword(credential), connectionTimeoutMillis: 30000 // 30 seconds instead of default }); ``` diff --git a/javascript/package.json b/javascript/package.json index 60fd36b..b5f7abf 100644 --- a/javascript/package.json +++ b/javascript/package.json @@ -42,10 +42,10 @@ "node": ">=18.0.0" }, "dependencies": { - "@azure/identity": "^4.13.0", "dotenv": "^17.2.3" }, "devDependencies": { + "@azure/identity": "^4.13.0", "@eslint/js": "^9.12.0", "chai": "^5.1.1", "chai-as-promised": "^8.0.2", @@ -58,6 +58,7 @@ "testcontainers": "^10.13.0" }, "peerDependencies": { + "@azure/identity": ">=4.0.0", "pg": ">=8.0.0", "sequelize": ">=6.0.0" }, diff --git a/javascript/samples/pg/getting-started/create-db-connection.js b/javascript/samples/pg/getting-started/create-db-connection.js index 0e492ee..1a4de91 100644 --- a/javascript/samples/pg/getting-started/create-db-connection.js +++ b/javascript/samples/pg/getting-started/create-db-connection.js @@ -2,6 +2,7 @@ // Licensed under the MIT License. import pg from "pg"; +import { DefaultAzureCredential } from '@azure/identity'; import { getEntraTokenPassword } from 'azure-postgresql-auth'; import dotenv from 'dotenv'; import { fileURLToPath } from 'url'; @@ -14,12 +15,14 @@ dotenv.config({ path: join(__dirname, '.env') }); const { Pool } = pg; +const credential = new DefaultAzureCredential(); + const pool = new Pool({ host: process.env.PGHOST, port: Number(process.env.PGPORT || 5432), database: process.env.PGDATABASE, user: process.env.PGUSER, - password: getEntraTokenPassword, + password: () => getEntraTokenPassword(credential), ssl: { rejectUnauthorized: false // or true with proper certificates }, diff --git a/javascript/samples/sequelize/getting-started/create-db-connection.js b/javascript/samples/sequelize/getting-started/create-db-connection.js index ae6419e..605c437 100644 --- a/javascript/samples/sequelize/getting-started/create-db-connection.js +++ b/javascript/samples/sequelize/getting-started/create-db-connection.js @@ -5,6 +5,7 @@ import dotenv from 'dotenv'; import { join, dirname } from 'path'; import { fileURLToPath } from 'url'; import { Sequelize } from 'sequelize'; +import { DefaultAzureCredential } from '@azure/identity'; import { configureEntraIdAuth } from 'azure-postgresql-auth'; // Load .env from the same directory as this script @@ -24,7 +25,8 @@ async function main() { }); // Configure Entra ID authentication - configureEntraIdAuth(sequelize); + const credential = new DefaultAzureCredential(); + configureEntraIdAuth(sequelize, credential); await sequelize.authenticate(); // triggers beforeConnect and opens a connection console.log('✅ Sequelize connection established successfully with Entra ID!'); diff --git a/javascript/src/entra-connection.js b/javascript/src/entra-connection.js index a7e9c9a..f4016b7 100644 --- a/javascript/src/entra-connection.js +++ b/javascript/src/entra-connection.js @@ -1,15 +1,17 @@ // Copyright (c) Microsoft Corporation. // Licensed under the MIT License. -import { DefaultAzureCredential } from '@azure/identity'; - /** * Configure Sequelize instance to use Entra ID authentication * @param {Sequelize} sequelizeInstance - The Sequelize instance to configure + * @param {Object} credential - The TokenCredential to use for authentication * @param {Object} options - Configuration options * @param {string} options.fallbackUsername - Fallback username if token doesn't contain upn/appid */ -export function configureEntraIdAuth(sequelizeInstance, credential = null, options = {}) { +export function configureEntraIdAuth(sequelizeInstance, credential, options = {}) { + if (!credential) { + throw new Error('credential is required'); + } const { fallbackUsername } = options; // Runs before every new connection is created by Sequelize @@ -31,10 +33,14 @@ export function configureEntraIdAuth(sequelizeInstance, credential = null, optio /** * Get cached Entra ID access token or fetch a new one + * @param {Object} credential - The TokenCredential to use for authentication + * @param {string} scope - The scope to request the token for * @returns {Promise} - The access token */ -export async function getEntraTokenPassword(credential = null, scope = "https://ossrdbms-aad.database.windows.net/.default") { - credential = credential || new DefaultAzureCredential(); +export async function getEntraTokenPassword(credential, scope = "https://ossrdbms-aad.database.windows.net/.default") { + if (!credential) { + throw new Error('credential is required'); + } try { const t = await credential.getToken(scope); if (!t?.token) {throw new Error('Failed to acquire Entra ID token');} @@ -69,4 +75,4 @@ function decodeJwtToken(token) { console.error('Error decoding JWT token:', error); return null; } -} \ No newline at end of file +} diff --git a/python/README.md b/python/README.md index aa8bdfc..00aaae8 100644 --- a/python/README.md +++ b/python/README.md @@ -119,25 +119,33 @@ pip install "azure-postgresql-auth[psycopg2]" ### Connection Pooling (Recommended) ```python -from azure_postgresql_auth.psycopg2 import EntraConnection # import library -from psycopg2 import pool # import to use pooling +from azure_postgresql_auth.psycopg2 import EntraConnection +from azure.identity import DefaultAzureCredential +from psycopg2 import pool +from functools import partial + +# Create a connection factory with the credential bound +credential = DefaultAzureCredential() +connection_factory = partial(EntraConnection, credential=credential) with pool.ThreadedConnectionPool( minconn=1, maxconn=5, host="your-server.postgres.database.azure.com", database="your_database", - connection_factory=EntraConnection + connection_factory=connection_factory # Pass the factory with credential bound ) as connection_pool: ``` ### Direct Connection ```python -from azure_postgresql_auth.psycopg2 import EntraConnection # import library +from azure_postgresql_auth.psycopg2 import EntraConnection +from azure.identity import DefaultAzureCredential with EntraConnection( - "postgresql://your-server.postgres.database.azure.com:5432/your_database" + "postgresql://your-server.postgres.database.azure.com:5432/your_database", + credential=DefaultAzureCredential() # required ) as conn ``` @@ -155,28 +163,32 @@ pip install "azure-postgresql-auth[psycopg3]" ### Synchronous Connection ```python -from azure_postgresql_auth.psycopg3 import EntraConnection # import library -from psycopg_pool import ConnectionPool # import to use pooling +from azure_postgresql_auth.psycopg3 import EntraConnection +from azure.identity import DefaultAzureCredential +from psycopg_pool import ConnectionPool with ConnectionPool( conninfo="postgresql://your-server.postgres.database.azure.com:5432/your_database", connection_class=EntraConnection, - min_size=1, # keep at least 1 connection always open - max_size=5, # allow up to 5 concurrent connections + kwargs={"credential": DefaultAzureCredential()}, # required + min_size=1, + max_size=5, ) as pool ``` ### Asynchronous Connection ```python -from azure_postgresql_auth.psycopg3 import AsyncEntraConnection # import library -from psycopg_pool import AsyncConnectionPool # import to use pooling +from azure_postgresql_auth.psycopg3 import AsyncEntraConnection +from azure.identity.aio import DefaultAzureCredential +from psycopg_pool import AsyncConnectionPool async with AsyncConnectionPool( conninfo="postgresql://your-server.postgres.database.azure.com:5432/your_database", connection_class=AsyncEntraConnection, - min_size=1, # keep at least 1 connection always open - max_size=5, # allow up to 5 concurrent connections + kwargs={"credential": DefaultAzureCredential()}, # required + min_size=1, + max_size=5, ) as pool ``` @@ -197,9 +209,13 @@ pip install "azure-postgresql-auth[sqlalchemy]" ```python from sqlalchemy import create_engine -from azure_postgresql_auth.sqlalchemy import enable_entra_authentication # import library +from azure_postgresql_auth.sqlalchemy import enable_entra_authentication +from azure.identity import DefaultAzureCredential -with create_engine("postgresql+psycopg://your-server.postgres.database.azure.com/your_database") as engine: +with create_engine( + "postgresql+psycopg://your-server.postgres.database.azure.com/your_database", + connect_args={"credential": DefaultAzureCredential()} # required +) as engine: # Enable Entra ID authentication enable_entra_authentication(engine) @@ -215,9 +231,13 @@ with create_engine("postgresql+psycopg://your-server.postgres.database.azure.com ```python from sqlalchemy.ext.asyncio import create_async_engine -from azure_postgresql_auth.sqlalchemy import enable_entra_authentication_async # import library +from azure_postgresql_auth.sqlalchemy import enable_entra_authentication_async +from azure.identity import DefaultAzureCredential -async with create_async_engine("postgresql+psycopg://your-server.postgres.database.azure.com/your_database") as engine: +async with create_async_engine( + "postgresql+psycopg://your-server.postgres.database.azure.com/your_database", + connect_args={"credential": DefaultAzureCredential()} # required +) as engine: # Enable Entra ID authentication for async enable_entra_authentication_async(engine) @@ -233,7 +253,7 @@ async with create_async_engine("postgresql+psycopg://your-server.postgres.databa ### Authentication Flow -1. **Token Acquisition**: Uses Azure Identity libraries (`DefaultAzureCredential` by default) to acquire access tokens from Azure Entra ID +1. **Token Acquisition**: Uses Azure Identity credentials (you must pass `DefaultAzureCredential()` or another credential explicitly) to acquire access tokens from Azure Entra ID 2. **Automatic Refresh**: Tokens are automatically refreshed before each new database connection 3. **Secure Transport**: Tokens are passed as passwords in PostgreSQL connection strings over SSL 4. **Server Validation**: Azure Database for PostgreSQL validates the token and establishes the authenticated connection diff --git a/python/pyproject.toml b/python/pyproject.toml index bde219d..fa012a6 100644 --- a/python/pyproject.toml +++ b/python/pyproject.toml @@ -18,8 +18,7 @@ classifiers = [ "Operating System :: OS Independent" ] dependencies = [ - "azure-identity>=1.13.0", - "azure-core>=1.24.0", + "azure-core>=1.24.0" ] [project.optional-dependencies] diff --git a/python/samples/psycopg2/getting_started/create_db_connection.py b/python/samples/psycopg2/getting_started/create_db_connection.py index 054d701..2030453 100644 --- a/python/samples/psycopg2/getting_started/create_db_connection.py +++ b/python/samples/psycopg2/getting_started/create_db_connection.py @@ -6,7 +6,9 @@ """ import os +from functools import partial +from azure.identity import DefaultAzureCredential from dotenv import load_dotenv from psycopg2 import pool @@ -24,12 +26,17 @@ def main() -> None: # authentication tokens are properly managed and refreshed so that each connection uses a valid token. # # For more details, see: https://www.psycopg.org/docs/advanced.html#subclassing-connection + + # Create a connection factory with the credential bound using functools.partial + credential = DefaultAzureCredential() + connection_factory = partial(EntraConnection, credential=credential) + connection_pool = pool.ThreadedConnectionPool( minconn=1, maxconn=5, host=SERVER, database=DATABASE, - connection_factory=EntraConnection, + connection_factory=connection_factory, ) conn = connection_pool.getconn() diff --git a/python/samples/psycopg3/getting_started/create_db_connection.py b/python/samples/psycopg3/getting_started/create_db_connection.py index 1f3e39b..c6e21a6 100644 --- a/python/samples/psycopg3/getting_started/create_db_connection.py +++ b/python/samples/psycopg3/getting_started/create_db_connection.py @@ -11,6 +11,8 @@ import os import sys +from azure.identity import DefaultAzureCredential +from azure.identity.aio import DefaultAzureCredential as AsyncDefaultAzureCredential from dotenv import load_dotenv from psycopg_pool import AsyncConnectionPool, ConnectionPool @@ -36,6 +38,7 @@ def main_sync() -> None: max_size=5, open=False, connection_class=EntraConnection, + kwargs={"credential": DefaultAzureCredential()}, ) with pool, pool.connection() as conn, conn.cursor() as cur: cur.execute("SELECT now()") @@ -57,6 +60,7 @@ async def main_async() -> None: max_size=5, open=False, connection_class=AsyncEntraConnection, + kwargs={"credential": AsyncDefaultAzureCredential()}, ) async with pool, pool.connection() as conn, conn.cursor() as cur: await cur.execute("SELECT now()") diff --git a/python/samples/sqlalchemy/getting_started/create_db_connection.py b/python/samples/sqlalchemy/getting_started/create_db_connection.py index 0fa6cc2..1cadf3e 100644 --- a/python/samples/sqlalchemy/getting_started/create_db_connection.py +++ b/python/samples/sqlalchemy/getting_started/create_db_connection.py @@ -11,6 +11,7 @@ import os import sys +from azure.identity import DefaultAzureCredential from dotenv import load_dotenv from sqlalchemy import create_engine, text from sqlalchemy.ext.asyncio import create_async_engine @@ -30,7 +31,10 @@ def main_sync() -> None: """Synchronous connection example using SQLAlchemy with Entra ID authentication.""" # Create a synchronous engine - engine = create_engine(f"postgresql+psycopg://{SERVER}/{DATABASE}") + engine = create_engine( + f"postgresql+psycopg://{SERVER}/{DATABASE}", + connect_args={"credential": DefaultAzureCredential()}, + ) # We add an event listener to the engine to enable synchronous Entra authentication # for database access. This event listener is triggered whenever the connection pool @@ -53,7 +57,10 @@ async def main_async() -> None: """Asynchronous connection example using SQLAlchemy with Entra ID authentication.""" # Create an asynchronous engine - engine = create_async_engine(f"postgresql+psycopg://{SERVER}/{DATABASE}") + engine = create_async_engine( + f"postgresql+psycopg://{SERVER}/{DATABASE}", + connect_args={"credential": DefaultAzureCredential()}, + ) # We add an event listener to the engine to enable asynchronous Entra authentication # for database access. This event listener is triggered whenever the connection pool diff --git a/python/src/azure_postgresql_auth/core.py b/python/src/azure_postgresql_auth/core.py index b9b66a0..eef831c 100644 --- a/python/src/azure_postgresql_auth/core.py +++ b/python/src/azure_postgresql_auth/core.py @@ -8,9 +8,6 @@ from azure.core.credentials import TokenCredential from azure.core.credentials_async import AsyncTokenCredential from azure.core.exceptions import ClientAuthenticationError -from azure.identity import CredentialUnavailableError -from azure.identity import DefaultAzureCredential as DefaultAzureCredential -from azure.identity.aio import DefaultAzureCredential as AsyncDefaultAzureCredential from azure_postgresql_auth.errors import ( ScopePermissionError, @@ -22,36 +19,32 @@ AZURE_MANAGEMENT_SCOPE = "https://management.azure.com/.default" -def get_entra_token(credential: TokenCredential | None, scope: str) -> str: +def get_entra_token(credential: TokenCredential, scope: str) -> str: """Acquires an Entra authentication token for Azure PostgreSQL synchronously. Parameters: - credential (TokenCredential or None): Credential object used to obtain the token. - If None, the default Azure credentials are used. + credential (TokenCredential): Credential object used to obtain the token. scope (str): The scope for the token request. Returns: str: The acquired authentication token to be used as the database password. """ - credential = credential or DefaultAzureCredential() cred = credential.get_token(scope) return cred.token async def get_entra_token_async( - credential: AsyncTokenCredential | None, scope: str + credential: AsyncTokenCredential, scope: str ) -> str: """Asynchronously acquires an Entra authentication token for Azure PostgreSQL. Parameters: - credential (AsyncTokenCredential or None): Asynchronous credential used to obtain the token. - If None, the default Azure credentials are used. + credential (AsyncTokenCredential): Asynchronous credential used to obtain the token. scope (str): The scope for the token request. Returns: str: The acquired authentication token to be used as the database password. """ - credential = credential or AsyncDefaultAzureCredential() async with credential: cred = await credential.get_token(scope) return cred.token @@ -107,15 +100,14 @@ def parse_principal_name(xms_mirid: str) -> str | None: return principal_name -def get_entra_conninfo(credential: TokenCredential | None) -> dict[str, str]: +def get_entra_conninfo(credential: TokenCredential) -> dict[str, str]: """Synchronously obtains connection information from Entra authentication for Azure PostgreSQL. This function acquires an access token from Azure Entra ID and extracts the username from the token claims. It tries multiple claim sources to determine the username. Parameters: - credential (TokenCredential or None): The credential used for token acquisition. - If None, DefaultAzureCredential() is used to automatically discover credentials. + credential (TokenCredential): The credential used for token acquisition. Returns: dict[str, str]: A dictionary with 'user' and 'password' keys, where: @@ -127,7 +119,6 @@ def get_entra_conninfo(credential: TokenCredential | None) -> dict[str, str]: UsernameExtractionError: If the username cannot be extracted from token claims. ScopePermissionError: The token could not be acquired from the management scope, possibly due to insufficient permissions. """ - credential = credential or DefaultAzureCredential() # Always get the DB-scope token for password db_token = get_entra_token(credential, AZURE_DB_FOR_POSTGRES_SCOPE) @@ -149,7 +140,7 @@ def get_entra_conninfo(credential: TokenCredential | None) -> dict[str, str]: # Fall back to management scope ONLY to discover username try: mgmt_token = get_entra_token(credential, AZURE_MANAGEMENT_SCOPE) - except (CredentialUnavailableError, ClientAuthenticationError) as e: + except ClientAuthenticationError as e: raise ScopePermissionError( "Failed to acquire token from management scope" ) from e @@ -177,7 +168,7 @@ def get_entra_conninfo(credential: TokenCredential | None) -> dict[str, str]: async def get_entra_conninfo_async( - credential: AsyncTokenCredential | None, + credential: AsyncTokenCredential, ) -> dict[str, str]: """Asynchronously obtains connection information from Entra authentication for Azure PostgreSQL. @@ -185,8 +176,7 @@ async def get_entra_conninfo_async( from the token claims. It tries multiple claim sources to determine the username. Parameters: - credential (AsyncTokenCredential or None): The async credential used for token acquisition. - If None, AsyncDefaultAzureCredential() is used to automatically discover credentials. + credential (AsyncTokenCredential): The async credential used for token acquisition. Returns: dict[str, str]: A dictionary with 'user' and 'password' keys, where: @@ -198,7 +188,6 @@ async def get_entra_conninfo_async( UsernameExtractionError: If the username cannot be extracted from token claims. ScopePermissionError: The token could not be acquired from the management scope, possibly due to insufficient permissions. """ - credential = credential or AsyncDefaultAzureCredential() db_token = await get_entra_token_async(credential, AZURE_DB_FOR_POSTGRES_SCOPE) try: @@ -218,7 +207,7 @@ async def get_entra_conninfo_async( if not username: try: mgmt_token = await get_entra_token_async(credential, AZURE_MANAGEMENT_SCOPE) - except (CredentialUnavailableError, ClientAuthenticationError) as e: + except ClientAuthenticationError as e: raise ScopePermissionError( "Failed to acquire token from management scope" ) from e diff --git a/python/src/azure_postgresql_auth/psycopg2/entra_connection.py b/python/src/azure_postgresql_auth/psycopg2/entra_connection.py index 7644b4c..6aee1f3 100644 --- a/python/src/azure_postgresql_auth/psycopg2/entra_connection.py +++ b/python/src/azure_postgresql_auth/psycopg2/entra_connection.py @@ -30,8 +30,7 @@ class EntraConnection(connection): Parameters: dsn (str): PostgreSQL connection string (Data Source Name). **kwargs: Additional keyword arguments including: - - credential (TokenCredential, optional): Azure credential for token acquisition. - If None, DefaultAzureCredential() is used. + - credential (TokenCredential, required): Azure credential for token acquisition. - user (str, optional): Database username. If not provided, extracted from Entra token. - password (str, optional): Database password. If not provided, uses Entra access token. @@ -45,9 +44,9 @@ def __init__(self, dsn: str, **kwargs: Any) -> None: dsn_params = parse_dsn(dsn) if dsn else {} credential = kwargs.pop("credential", None) - if credential and not isinstance(credential, (TokenCredential)): + if credential is None or not isinstance(credential, (TokenCredential)): raise CredentialValueError( - "credential must be a TokenCredential for sync connections" + "credential is required and must be a TokenCredential for sync connections" ) # Check if user and password are already provided diff --git a/python/src/azure_postgresql_auth/psycopg3/async_entra_connection.py b/python/src/azure_postgresql_auth/psycopg3/async_entra_connection.py index 093a827..6e5d75f 100644 --- a/python/src/azure_postgresql_auth/psycopg3/async_entra_connection.py +++ b/python/src/azure_postgresql_auth/psycopg3/async_entra_connection.py @@ -34,7 +34,7 @@ async def connect(cls, *args: Any, **kwargs: Any) -> "AsyncEntraConnection": Parameters: *args: Positional arguments to be forwarded to the parent connection method. **kwargs: Keyword arguments including: - - credential (AsyncTokenCredential, optional): Async Azure credential for token acquisition. + - credential (AsyncTokenCredential, required): Async Azure credential for token acquisition. - user (str, optional): Database username. If not provided, extracted from Entra token. - password (str, optional): Database password. If not provided, uses Entra access token. @@ -46,9 +46,9 @@ async def connect(cls, *args: Any, **kwargs: Any) -> "AsyncEntraConnection": EntraConnectionValueError: If Entra connection credentials are invalid. """ credential = kwargs.pop("credential", None) - if credential and not isinstance(credential, (AsyncTokenCredential)): + if credential is None or not isinstance(credential, (AsyncTokenCredential)): raise CredentialValueError( - "credential must be an AsyncTokenCredential for async connections" + "credential is required and must be an AsyncTokenCredential for async connections" ) # Check if we need to acquire Entra authentication info diff --git a/python/src/azure_postgresql_auth/psycopg3/entra_connection.py b/python/src/azure_postgresql_auth/psycopg3/entra_connection.py index 82e2188..c9d9c4a 100644 --- a/python/src/azure_postgresql_auth/psycopg3/entra_connection.py +++ b/python/src/azure_postgresql_auth/psycopg3/entra_connection.py @@ -34,7 +34,7 @@ def connect(cls, *args: Any, **kwargs: Any) -> "EntraConnection": Parameters: *args: Positional arguments to be forwarded to the parent connection method. **kwargs: Keyword arguments including: - - credential (TokenCredential, optional): Azure credential for token acquisition. + - credential (TokenCredential, required): Azure credential for token acquisition. - user (str, optional): Database username. If not provided, extracted from Entra token. - password (str, optional): Database password. If not provided, uses Entra access token. @@ -46,9 +46,9 @@ def connect(cls, *args: Any, **kwargs: Any) -> "EntraConnection": EntraConnectionValueError: If Entra connection credentials cannot be retrieved """ credential = kwargs.pop("credential", None) - if credential and not isinstance(credential, (TokenCredential)): + if credential is None or not isinstance(credential, (TokenCredential)): raise CredentialValueError( - "credential must be a TokenCredential for sync connections" + "credential is required and must be a TokenCredential for sync connections" ) # Check if we need to acquire Entra authentication info diff --git a/python/src/azure_postgresql_auth/sqlalchemy/async_entra_connection.py b/python/src/azure_postgresql_auth/sqlalchemy/async_entra_connection.py index 5f7709b..a0b1f60 100644 --- a/python/src/azure_postgresql_auth/sqlalchemy/async_entra_connection.py +++ b/python/src/azure_postgresql_auth/sqlalchemy/async_entra_connection.py @@ -28,12 +28,18 @@ def enable_entra_authentication_async(engine: AsyncEngine) -> None: Enable Azure Entra ID authentication for an async SQLAlchemy engine. This function registers an event listener that automatically provides - Entra ID credentials for each database connection if they are not already set. - Event handlers do not support async behavior so the token fetching will still - be synchronous. + Entra ID credentials for each database connection. A credential must be + provided via connect_args when creating the engine. Event handlers do not + support async behavior so the token fetching will still be synchronous. Args: engine: The async SQLAlchemy Engine to enable Entra authentication for + Example: + engine = create_async_engine( + "postgresql+psycopg://server/db", + connect_args={"credential": DefaultAzureCredential()} + ) + enable_entra_authentication_async(engine) """ @event.listens_for(engine.sync_engine, "do_connect") @@ -47,9 +53,10 @@ def provide_token( EntraConnectionValueError: If Entra connection credentials cannot be retrieved """ credential = cparams.get("credential", None) - if credential and not isinstance(credential, (TokenCredential)): + if credential is None or not isinstance(credential, (TokenCredential)): raise CredentialValueError( - "credential must be a TokenCredential for async connections" + "credential is required and must be a TokenCredential. " + "Pass it via connect_args={'credential': DefaultAzureCredential()}" ) # Check if credentials are already present has_user = "user" in cparams diff --git a/python/src/azure_postgresql_auth/sqlalchemy/entra_connection.py b/python/src/azure_postgresql_auth/sqlalchemy/entra_connection.py index e35eb18..7e3cd49 100644 --- a/python/src/azure_postgresql_auth/sqlalchemy/entra_connection.py +++ b/python/src/azure_postgresql_auth/sqlalchemy/entra_connection.py @@ -26,10 +26,17 @@ def enable_entra_authentication(engine: Engine) -> None: Enable Azure Entra ID authentication for a SQLAlchemy engine. This function registers an event listener that automatically provides - Entra ID credentials for each database connection if they are not already set. + Entra ID credentials for each database connection. A credential must be + provided via connect_args when creating the engine. Args: engine: The SQLAlchemy Engine to enable Entra authentication for + Example: + engine = create_engine( + "postgresql+psycopg://server/db", + connect_args={"credential": DefaultAzureCredential()} + ) + enable_entra_authentication(engine) """ @event.listens_for(engine, "do_connect") @@ -43,9 +50,10 @@ def provide_token( EntraConnectionValueError: If Entra connection credentials cannot be retrieved """ credential = cparams.get("credential", None) - if credential and not isinstance(credential, (TokenCredential)): + if credential is None or not isinstance(credential, (TokenCredential)): raise CredentialValueError( - "credential must be a TokenCredential for sync connections" + "credential is required and must be a TokenCredential. " + "Pass it via connect_args={'credential': DefaultAzureCredential()}" ) # Check if credentials are already present has_user = "user" in cparams