Usage of Workload Identity to connect to PostgresSQL

[papanito]

We are running netbox on AKS using a managed PostgresSQL database. Is there any way we could connect with managed identity to Azure Database for PostgreSQL - Flexible Server

[dattebayo-ua]

Hi,

If this issue is still relevant for you, I wanted to share my solution for running NetBox on an Azure VM. I've implemented a class in configuration.py that handles automatic token refresh before expiration occurs.
I based my approach on a reference implementation I found in another GH discussion:

This solution has been working well in my environment, I did not notice any broken/closed sessions
Now I'm trying to understand how can I deliver same configuration but on Azure Container Instance

import time
from typing import Optional
from azure.core.credentials import AccessToken
from azure.identity import DefaultAzureCredential


class DatabaseToken(str):
    cached_token: Optional[AccessToken] = None

    def __new__(cls) -> 'DatabaseToken':
        token_obj = cls.get_new_token()
        # Create the new instance with the actual token value
        instance = super().__new__(cls, token_obj.token)
        instance.cached_token = token_obj
        return instance

    @classmethod
    def get_new_token(cls) -> AccessToken:
        try:
            managed_identity_client_id = ""
            credential = DefaultAzureCredential(managed_identity_client_id=managed_identity_client_id)
            return credential.get_token("https://ossrdbms-aad.database.windows.net/.default")
        except Exception as e:
            print(f"Error getting token: {str(e)}")
            # Return a dummy token as fallback
            return AccessToken("dummy_token", int(time.time()) + 3600)

    def token_is_valid(self) -> bool:
        if self.cached_token is None:
            return False
        # Check if the cached token's expiry is still in the future
        return self.cached_token.expires_on > time.time()

    def refresh_token(self) -> None:
        # Refresh the token if expired
        if not self.token_is_valid():
            new_token = self.get_new_token()
            self.cached_token = new_token
    
    # Optionally override __str__ if we want to always ensure it returns the current token:
    def __str__(self) -> str:
        # Optionally, we could refresh the token here if expired:
        if not self.token_is_valid():
            # Note: Since the string is immutable, we'd need to recreate a new instance.
            new_instance = DatabaseToken()
            return new_instance
        return self.cached_token.token

DATABASE = {
    'ENGINE': 'django.db.backends.postgresql', 
    'NAME': 'netbox',
    'USER': 'my-managed-identity',  # PostgreSQL username -> Managed Identity resource name, not client_id, object_id
    'PASSWORD': DatabaseToken(),
    'HOST': 'mypgsql.postgres.database.azure.com',
    'PORT': '5432',
    'CONN_MAX_AGE': 30 * 60,
    'TIME_ZONE': 'UTC',
    'OPTIONS': {
        'sslmode': 'require',  # Required for Azure PostgreSQL
    }  
}