Azure Container App Managed Identity token authentication for Rails PostgreSQL adapter

database_token.rb

require "net/http"
require "json"

module AzureTokenAuthentication
  module Config
    # NOTE: https://learn.microsoft.com/en-us/azure/app-service/overview-managed-identity#rest-endpoint-reference

    class << self
      # endpoint example: "http://localhost:42356/msi/token"
      def endpoint = ENV.fetch("IDENTITY_ENDPOINT")

      # identity_header example: "12345678-1234-1234-1234-123456789012"
      def identity_header = ENV.fetch("IDENTITY_HEADER")

      # client_id example: "12345678-1234-1234-1234-123456789012"
      def client_id = ENV.fetch("AZURE_CLIENT_ID")
    end
  end

  def self.authentication_enabled? = ENV["IDENTITY_ENDPOINT"].present?

  def self.fetch_token
    uri = URI(Config.endpoint)
    params = {
      "api-version" => "2019-08-01",
      "resource" => "https://ossrdbms-aad.database.windows.net",
      "client_id" => Config.client_id
    }
    uri.query = URI.encode_www_form(params)

    request = Net::HTTP::Get.new(uri)
    request["X-IDENTITY-HEADER"] = Config.identity_header

    response = Net::HTTP.start(uri.host, uri.port) do |http|
      http.request(request)
    end

    if response.is_a?(Net::HTTPSuccess)
      JSON.parse(response.body)["access_token"]
    else
      Rails.logger.error "Failed to acquire Azure token: #{response.body}"
      raise "Failed to acquire Azure token: #{response.code} #{response.message}"
    end
  rescue => e
    Rails.logger.error "Failed to acquire Azure token: #{e.message}"
    raise
  end
end

module ActiveRecord
  module ConnectionAdapters
    class PostgreSQLAdapter
      class << self
        alias_method :original_new_client, :new_client

        def new_client(config)
          if AzureTokenAuthentication.authentication_enabled?
            Rails.logger.info "Fetching Azure token for database connection"
            password = AzureTokenAuthentication.fetch_token
            config[:password] = password
          end

          Rails.logger.info "Establishing database connection"
          original_new_client(config)
        end
      end
    end
  end
end