rails-pg-managed-identity 0.1.0

checksums.yaml

@@ -0,0 +1,7 @@
+---
+SHA256:
+  metadata.gz: c9fc9ac2aab390952a05fe1ee9473b6b15474216ea9684a7cca163f1201193f2
+  data.tar.gz: 89cbc58e3bbd6c1d6351cf8231fa4b6410f9c09494fc7a71d89a44a975e4efcb
+SHA512:
+  metadata.gz: 70c882a93e9b9cdaf6020c9830e67d9bc8124e01d8e3a93622e9a5192d9b7698831bd648779400ba605c7eafcbf0e69e35f322c54890a46cf1b11585023a77ed
+  data.tar.gz: 6e34664a2acfc61f760f13de130cefa4a8a0667f52ba76ca3e0b8b6f38bf3f4371c43061025989c5e946606b0c2654e9ced9f2303485a60019f580b8b81cc26b

data/CHANGELOG.md

@@ -0,0 +1,14 @@
+## [Unreleased]
+
+## [0.1.0] - 2026-10-15
+
+- Initial release. Supports Rails 7.2, 8.0 and 8.1 on Ruby 3.3+.
+- `azure_managed_identity` option in `config/database.yml` for the PostgreSQL adapter: the password is a
+  Microsoft Entra access token fetched when a connection is opened, cached, and refreshed 5 minutes before
+  it expires.
+- Credentials: managed identity (App Service, Functions, Container Apps, IMDS on VMs, scale sets, AKS
+  nodes), AKS workload identity, service principal with a client secret, and the Azure CLI, chained like
+  azure-identity's `DefaultAzureCredential`.
+- Reconnects open a new connection with the current token instead of `PG::Connection#reset`.
+- `db:schema:dump`/`db:schema:load` with `structure.sql` and `bin/rails dbconsole` get the token too.
+- Tokens never show up in logs, `inspect` or `pp`/IRB output.

data/CODE_OF_CONDUCT.md

@@ -0,0 +1,10 @@
+# Code of Conduct
+
+"rails-pg-managed-identity" follows [The Ruby Community Conduct Guideline](https://www.ruby-lang.org/en/conduct) in all "collaborative space", which is defined as community communications channels (such as mailing lists, submitted patches, commit comments, etc.):
+
+* Participants will be tolerant of opposing views.
+* Participants must ensure that their language and actions are free of personal attacks and disparaging personal remarks.
+* When interpreting the words and actions of others, participants should always assume good intentions.
+* Behaviour which can be reasonably considered harassment will not be tolerated.
+
+If you have any concerns about behaviour within this project, please contact us at ["fst@last.dk"](mailto:"fst@last.dk").

data/LICENSE.txt

@@ -0,0 +1,21 @@
+The MIT License (MIT)
+
+Copyright (c) 2026 Frederik Erbs Spang Thomsen
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in
+all copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
+THE SOFTWARE.

data/README.md

@@ -0,0 +1,104 @@
+# rails-pg-managed-identity
+
+> [!WARNING]
+> Not yet published. Not yet tested in real life world.
+
+Passwordless connections from Rails (7.2+, Ruby 3.3+) to **Azure Database for PostgreSQL flexible server** using Microsoft Entra access tokens: 
+the app's managed identity, AKS workload identity, a service principal, or your `az login` during development.
+
+## Installation
+
+```ruby
+# Gemfile
+gem "rails-pg-managed-identity"
+```
+
+## Usage
+
+Add `azure_managed_identity` to the database configuration and leave out the password:
+
+```yaml
+# config/database.yml
+production:
+  adapter: postgresql
+  host: my-server.postgres.database.azure.com
+  database: my_app
+  username: my-app-identity   # the identity's PostgreSQL role, see "Azure setup"
+  sslmode: require
+  azure_managed_identity: true
+```
+
+That's it. Booting the app requests nothing; the first query opens a connection, which fetches a token.
+Every connection in every pool shares the cached token.
+
+Options, instead of `true`:
+
+```yaml
+  azure_managed_identity:
+    client_id: <%= ENV["AZURE_CLIENT_ID"] %>  # user-assigned managed identity (default: AZURE_CLIENT_ID)
+    credential: managed_identity              # use only this credential (default: "default", the chain below)
+    scope: https://ossrdbms-aad.database.usgovcloudapi.net/.default  # other clouds
+```
+
+`DATABASE_URL` works too: `postgres://my-app-identity@my-server.postgres.database.azure.com/my_app?sslmode=require&azure_managed_identity=true`.
+
+### Where tokens come from
+
+`credential: default` tries these in order, like azure-identity's `DefaultAzureCredential`, and sticks
+with the first one that works:
+
+| `credential:` | Where | Configured by |
+| --- | --- | --- |
+| `environment` | Service principal with a secret, anywhere | `AZURE_TENANT_ID`, `AZURE_CLIENT_ID`, `AZURE_CLIENT_SECRET` |
+| `workload_identity` | AKS with Microsoft Entra Workload ID | `AZURE_TENANT_ID`, `AZURE_CLIENT_ID`, `AZURE_FEDERATED_TOKEN_FILE` (injected by the webhook) |
+| `managed_identity` | App Service, Functions, Container Apps | `IDENTITY_ENDPOINT`, `IDENTITY_HEADER` (set by the platform) |
+| `managed_identity` | VMs, scale sets, AKS nodes | Instance Metadata Service at `169.254.169.254` |
+| `azure_cli` | Your machine | `az login` |
+
+Azure Arc and Service Fabric managed identities aren't supported.
+
+### Azure setup
+
+1. Enable Microsoft Entra authentication on the server and add yourself as Microsoft Entra administrator.
+2. Give the app an identity (system- or user-assigned managed identity, or workload identity).
+3. As the Entra administrator, connected to the `postgres` database, create a role for the identity, named after it:
+
+   ```sql
+   select * from pgaadauth_create_principal('my-app-identity', false, false);
+   ```
+
+4. Grant that role what the app needs, and use its name as `username` in `database.yml`.
+
+For local development against Azure, `az login` and use your own user principal name (or an Entra group you're a member of) as `username`.
+
+### Token lifecycle
+
+- Tokens are fetched lazily, when a connection is opened, and reused until 5 minutes before expiry (`RailsPg::ManagedIdentity.refresh_offset = 300`).
+- While the cached token is still valid, a new one is requested at most every 30 seconds. 
+  If that request fails, the cached token is used and the failure is logged. A token is never handed out with less than 10 seconds left.
+
+`bin/rails dbconsole`, and `db:schema:dump`/`db:schema:load` with `structure.sql` (`pg_dump`/`psql`), get the token through `PGPASSWORD`. 
+The database tasks are only hooked when rake tasks are loaded, so app and job processes don't load that code.
+
+## Development
+
+```bash
+bin/setup
+bundle exec rake   # tests and RuboCop
+
+# Against another supported Rails version (CI runs Ruby 3.3, 3.4 and 4.0 with Rails 7.2, 8.0 and 8.1):
+BUNDLE_GEMFILE=gemfiles/rails_7.2.gemfile bundle install
+BUNDLE_GEMFILE=gemfiles/rails_7.2.gemfile bundle exec rake test
+```
+
+The tests need PostgreSQL's server binaries (`initdb`, `pg_ctl`); they're found through `pg_config`, or set `PG_BINDIR`.
+No running server or Azure account is needed:
+
+## Contributing
+
+Bug reports and pull requests are welcome. This project is intended to be a safe, welcoming space for
+collaboration, and contributors are expected to adhere to the [code of conduct](CODE_OF_CONDUCT.md).
+
+## License
+
+The gem is available as open source under the terms of the [MIT License](https://opensource.org/licenses/MIT).

data/Rakefile

@@ -0,0 +1,12 @@
+# frozen_string_literal: true
+
+require "bundler/gem_tasks"
+require "minitest/test_task"
+
+Minitest::TestTask.create
+
+require "rubocop/rake_task"
+
+RuboCop::RakeTask.new
+
+task default: %i[test rubocop]

data/gemfiles/rails_7.2.gemfile

@@ -0,0 +1,6 @@
+# frozen_string_literal: true
+
+# CI runs the tests against each supported Rails version: BUNDLE_GEMFILE=gemfiles/rails_7.2.gemfile
+eval_gemfile "../Gemfile"
+
+gem "activerecord", "~> 7.2.0"

data/gemfiles/rails_8.0.gemfile

@@ -0,0 +1,6 @@
+# frozen_string_literal: true
+
+# CI runs the tests against each supported Rails version: BUNDLE_GEMFILE=gemfiles/rails_8.0.gemfile
+eval_gemfile "../Gemfile"
+
+gem "activerecord", "~> 8.0.0"

data/gemfiles/rails_8.1.gemfile

@@ -0,0 +1,6 @@
+# frozen_string_literal: true
+
+# CI runs the tests against each supported Rails version: BUNDLE_GEMFILE=gemfiles/rails_8.1.gemfile
+eval_gemfile "../Gemfile"
+
+gem "activerecord", "~> 8.1.0"

data/lib/rails-pg-managed-identity.rb

@@ -0,0 +1,4 @@
+# frozen_string_literal: true
+
+# Bundler requires a gem by its name; the code lives in RailsPg::ManagedIdentity.
+require_relative "rails_pg/managed_identity"

data/lib/rails_pg/managed_identity/access_token.rb

@@ -0,0 +1,41 @@
+# frozen_string_literal: true
+
+require "time"
+
+module RailsPg
+  module ManagedIdentity
+    # An access token and its expiry in seconds since the epoch.
+    AccessToken = Data.define(:token, :expires_on) do
+      # Builds a token from a JSON token response. Managed identity endpoints send expires_on,
+      # Microsoft Entra ID sends expires_in.
+      def self.from_response(response, source)
+        unless response.status == 200
+          raise AuthenticationError, "#{source} returned HTTP #{response.status}: #{response.error_message}"
+        end
+
+        token = response.json["access_token"]
+        expires_on = Integer(response.json["expires_on"], exception: false)
+        expires_in = Integer(response.json["expires_in"], exception: false)
+        expires_on ||= Time.now.to_i + expires_in if expires_in
+        raise AuthenticationError, "#{source} returned an incomplete token response" unless token && expires_on
+
+        new(token:, expires_on:)
+      end
+
+      def expires_in
+        expires_on - Time.now.to_i
+      end
+
+      # Never show the token itself in logs, consoles or exception messages.
+      def inspect
+        "#<#{self.class.name} expires_on=#{Time.at(expires_on).utc.iso8601}>"
+      end
+      alias_method :to_s, :inspect
+
+      # pp and IRB use this rather than #inspect.
+      def pretty_print(printer)
+        printer.text(inspect)
+      end
+    end
+  end
+end

data/lib/rails_pg/managed_identity/active_record_extensions.rb

@@ -0,0 +1,53 @@
+# frozen_string_literal: true
+
+module RailsPg
+  module ManagedIdentity
+    # Prepended to ActiveRecord::ConnectionAdapters::PostgreSQLAdapter.
+    module PostgreSQLAdapterExtension
+      private
+
+      # Every physical connection, first or re-established, is opened here, so this is where
+      # the password is resolved.
+      def connect
+        if (provider = ManagedIdentity.token_provider_for(@config))
+          @connection_parameters = @connection_parameters.merge(password: provider.token)
+        end
+        super
+      rescue Error => e
+        raise e.set_pool(@pool)
+      end
+
+      # The stock implementation tries PG::Connection#reset first, which re-authenticates with the
+      # password the connection was first opened with, possibly an expired token by now. Without a
+      # raw connection it goes straight to #connect.
+      def reconnect
+        if ManagedIdentity.token_provider_for(@config)
+          begin
+            @raw_connection&.close
+          rescue ::PG::Error
+            nil
+          end
+          @raw_connection = nil
+        end
+        super
+      end
+    end
+
+    # Extends PostgreSQLAdapter's class methods.
+    module PostgreSQLAdapterClassExtension
+      # `bin/rails dbconsole`: psql can't prompt for a password as long as an access token, so it
+      # always gets the token through PGPASSWORD.
+      def dbconsole(config, options = {})
+        if (provider = ManagedIdentity.token_provider_for(config.configuration_hash))
+          ENV["PGPASSWORD"] = provider.token
+        end
+        super
+      end
+    end
+  end
+end
+
+ActiveSupport.on_load(:active_record_postgresqladapter) do
+  prepend RailsPg::ManagedIdentity::PostgreSQLAdapterExtension
+  singleton_class.prepend RailsPg::ManagedIdentity::PostgreSQLAdapterClassExtension
+end

data/lib/rails_pg/managed_identity/azure_cli_credential.rb

@@ -0,0 +1,68 @@
+# frozen_string_literal: true
+
+require "open3"
+
+module RailsPg
+  module ManagedIdentity
+    # The account signed in to the Azure CLI (`az login`), for local development.
+    class AzureCliCredential
+      # `az` is usually a shell script running Python. In its own process group, a timeout can kill both.
+      SPAWN_OPTIONS = Gem.win_platform? ? {}.freeze : { pgroup: true }.freeze
+
+      def initialize(process_timeout: 10)
+        @process_timeout = process_timeout
+      end
+
+      def get_token(scope)
+        stdout, stderr, status = run("az", "account", "get-access-token", "--output", "json", "--scope", scope)
+        unless status.success?
+          raise CredentialUnavailableError, "Azure CLI isn't signed in; run `az login`" if stderr.include?("az login")
+
+          raise AuthenticationError, "Azure CLI failed: #{stderr.strip.lines.first}"
+        end
+
+        parse(stdout)
+      end
+
+      private
+
+      def run(*command)
+        Open3.popen3(*command, **SPAWN_OPTIONS) do |stdin, stdout, stderr, wait_thread|
+          stdin.close
+          out = reader(stdout)
+          err = reader(stderr)
+          unless wait_thread.join(@process_timeout)
+            kill(wait_thread.pid)
+            raise AuthenticationError, "Azure CLI didn't respond within #{@process_timeout} seconds"
+          end
+          [out.value, err.value, wait_thread.value]
+        end
+      rescue Errno::ENOENT
+        raise CredentialUnavailableError, "Azure CLI (az) not found on PATH"
+      end
+
+      # On a timeout the pipes are closed under the reader; that's expected, not worth a report.
+      def reader(io)
+        Thread.new do
+          Thread.current.report_on_exception = false
+          io.read
+        end
+      end
+
+      def kill(pid)
+        Process.kill("KILL", SPAWN_OPTIONS[:pgroup] ? -pid : pid)
+      rescue Errno::ESRCH
+        nil # exited in the meantime
+      end
+
+      # Azure CLI 2.54+ sends expires_on (epoch seconds); older versions only expiresOn in local time.
+      def parse(stdout)
+        json = JSON.parse(stdout)
+        expires_on = Integer(json["expires_on"], exception: false) || Time.parse(json.fetch("expiresOn")).to_i
+        AccessToken.new(token: json.fetch("accessToken"), expires_on:)
+      rescue JSON::ParserError, KeyError, ArgumentError => e
+        raise AuthenticationError, "Unexpected output from Azure CLI: #{e.message}"
+      end
+    end
+  end
+end

data/lib/rails_pg/managed_identity/default_azure_credential.rb

@@ -0,0 +1,44 @@
+# frozen_string_literal: true
+
+module RailsPg
+  module ManagedIdentity
+    # Tries credentials in the same order as azure-identity's DefaultAzureCredential (minus the
+    # ones that need a desktop or an IDE): environment, workload identity, managed identity,
+    # Azure CLI. The first credential that returns a token is used from then on.
+    class DefaultAzureCredential
+      def initialize(client_id: nil)
+        @credentials = [
+          EnvironmentCredential.new,
+          WorkloadIdentityCredential.new(client_id:),
+          ManagedIdentityCredential.new(client_id:),
+          AzureCliCredential.new
+        ]
+        @selected = nil
+      end
+
+      def get_token(scope)
+        return @selected.get_token(scope) if @selected
+
+        attempts = []
+        @credentials.each do |credential|
+          token = credential.get_token(scope)
+          @selected = credential
+          ManagedIdentity.logger&.info { "[#{self.class.name}] Using #{credential.class.name.demodulize}" }
+          return token
+        rescue Error => e
+          attempts << "#{credential.class.name.demodulize}: #{e.message}"
+          # Unavailable means not applicable here, so try the next one; any other failure stops the
+          # chain, like azure-identity does.
+          raise AuthenticationError, failure_message(attempts) unless e.is_a?(CredentialUnavailableError)
+        end
+        raise CredentialUnavailableError, failure_message(attempts)
+      end
+
+      private
+
+      def failure_message(attempts)
+        "DefaultAzureCredential failed to get an access token:\n  #{attempts.join("\n  ")}"
+      end
+    end
+  end
+end

data/lib/rails_pg/managed_identity/entra_id.rb

@@ -0,0 +1,24 @@
+# frozen_string_literal: true
+
+module RailsPg
+  module ManagedIdentity
+    # OAuth 2.0 client credentials requests against Microsoft Entra ID's v2.0 token endpoint.
+    module EntraID
+      DEFAULT_AUTHORITY_HOST = "https://login.microsoftonline.com"
+
+      # params authenticate the client: client_secret, or client_assertion and client_assertion_type.
+      def self.client_credentials_token(tenant_id:, client_id:, scope:, **params)
+        uri = URI("#{authority_host}/#{tenant_id}/oauth2/v2.0/token")
+        form = { grant_type: "client_credentials", client_id:, scope:, **params }
+        AccessToken.from_response(HTTP.post_form(uri, form), "Microsoft Entra ID")
+      end
+
+      # AZURE_AUTHORITY_HOST selects a sovereign cloud; the AKS workload identity webhook sets it too.
+      def self.authority_host
+        host = ENV["AZURE_AUTHORITY_HOST"].to_s.strip.presence || DEFAULT_AUTHORITY_HOST
+        host = "https://#{host}" unless host.include?("://")
+        host.chomp("/")
+      end
+    end
+  end
+end

data/lib/rails_pg/managed_identity/environment_credential.rb

@@ -0,0 +1,19 @@
+# frozen_string_literal: true
+
+module RailsPg
+  module ManagedIdentity
+    # A service principal with a client secret, configured through AZURE_TENANT_ID,
+    # AZURE_CLIENT_ID and AZURE_CLIENT_SECRET.
+    class EnvironmentCredential
+      def get_token(scope)
+        tenant_id, client_id, client_secret = ENV.values_at("AZURE_TENANT_ID", "AZURE_CLIENT_ID", "AZURE_CLIENT_SECRET")
+        if [tenant_id, client_id, client_secret].any?(&:blank?)
+          raise CredentialUnavailableError,
+                "EnvironmentCredential requires AZURE_TENANT_ID, AZURE_CLIENT_ID and AZURE_CLIENT_SECRET"
+        end
+
+        EntraID.client_credentials_token(tenant_id:, client_id:, scope:, client_secret:)
+      end
+    end
+  end
+end

data/lib/rails_pg/managed_identity/http.rb

@@ -0,0 +1,74 @@
+# frozen_string_literal: true
+
+require "json"
+require "net/http"
+require "openssl"
+
+module RailsPg
+  module ManagedIdentity
+    # Minimal JSON-over-HTTP client for the token endpoints, with retries for transient failures.
+    module HTTP
+      NETWORK_ERRORS = [SystemCallError, IOError, SocketError, Timeout::Error, OpenSSL::SSL::SSLError].freeze
+      # The endpoint couldn't be reached at all.
+      class UnreachableError < AuthenticationError; end
+      RETRY_STATUSES = [429, 500, 502, 503, 504].freeze
+      MAX_RETRIES = 2
+
+      Response = Data.define(:status, :json, :body) do
+        def error_message
+          json["error_description"] || json["message"] || json["error"] || body.to_s[0, 200]
+        end
+      end
+
+      @retry_delay = 0.5
+
+      class << self
+        # Base delay in seconds before retrying; doubled for every further attempt.
+        attr_accessor :retry_delay
+
+        # Pass proxy: false for managed identity endpoints: they're local to the host, and Microsoft
+        # doesn't support reaching IMDS through a proxy.
+        def get(uri, headers: {}, open_timeout: 5, retry_statuses: RETRY_STATUSES, proxy: true)
+          request(Net::HTTP::Get.new(uri, headers), open_timeout:, retry_statuses:, proxy:)
+        end
+
+        def post_form(uri, form)
+          post = Net::HTTP::Post.new(uri)
+          post.set_form_data(form)
+          request(post, open_timeout: 5, retry_statuses: RETRY_STATUSES, proxy: true)
+        end
+
+        private
+
+        def request(request, open_timeout:, retry_statuses:, proxy:)
+          attempt = 0
+          loop do
+            response = send_request(request, open_timeout:, proxy:)
+            status = response.code.to_i
+            if retry_statuses.include?(status) && attempt < MAX_RETRIES
+              sleep(retry_delay * (2**attempt))
+              attempt += 1
+              next
+            end
+            return Response.new(status:, json: parse(response.body), body: response.body)
+          end
+        end
+
+        def send_request(request, open_timeout:, proxy:)
+          uri = request.uri
+          options = { use_ssl: uri.scheme == "https", open_timeout:, read_timeout: 10 }
+          Net::HTTP.start(uri.host, uri.port, (proxy ? :ENV : nil), options) { |http| http.request(request) }
+        rescue *NETWORK_ERRORS => e
+          raise UnreachableError, "#{uri.host}:#{uri.port} unreachable (#{e.class}: #{e.message})"
+        end
+
+        def parse(body)
+          parsed = JSON.parse(body.to_s)
+          parsed.is_a?(Hash) ? parsed : {}
+        rescue JSON::ParserError
+          {}
+        end
+      end
+    end
+  end
+end

data/lib/rails_pg/managed_identity/managed_identity_credential.rb

@@ -0,0 +1,70 @@
+# frozen_string_literal: true
+
+module RailsPg
+  module ManagedIdentity
+    # The managed identity of the Azure resource the app runs on: system-assigned by default, or
+    # the user-assigned identity with the given client ID (default: AZURE_CLIENT_ID).
+    #
+    # App Service, Azure Functions and Container Apps expose a local token endpoint through
+    # IDENTITY_ENDPOINT and IDENTITY_HEADER. Everywhere else (VMs, scale sets, AKS nodes) tokens
+    # come from the Azure Instance Metadata Service (IMDS).
+    class ManagedIdentityCredential
+      IMDS_HOST = "http://169.254.169.254"
+      IMDS_PATH = "/metadata/identity/oauth2/token"
+      # Microsoft documents 404 and 410 as transient for managed identity endpoints, too.
+      RETRY_STATUSES = (HTTP::RETRY_STATUSES + [404, 410]).freeze
+
+      def initialize(client_id: nil)
+        @client_id = client_id
+      end
+
+      def get_token(scope)
+        resource = scope.delete_suffix("/.default")
+        if ENV["IDENTITY_ENDPOINT"].blank? || ENV["IDENTITY_HEADER"].blank?
+          imds_token(resource)
+        else
+          app_service_token(resource)
+        end
+      end
+
+      private
+
+      def client_id
+        @client_id.presence || ENV.fetch("AZURE_CLIENT_ID", nil).presence
+      end
+
+      def token_uri(base, api_version, resource)
+        uri = URI(base)
+        uri.query = URI.encode_www_form({ "api-version" => api_version, "resource" => resource,
+                                          "client_id" => client_id }.compact)
+        uri
+      end
+
+      def app_service_token(resource)
+        uri = token_uri(ENV.fetch("IDENTITY_ENDPOINT"), "2019-08-01", resource)
+        response = HTTP.get(uri, headers: { "X-IDENTITY-HEADER" => ENV.fetch("IDENTITY_HEADER") },
+                                 retry_statuses: RETRY_STATUSES, proxy: false)
+        AccessToken.from_response(response, "Managed identity endpoint (IDENTITY_ENDPOINT)")
+      end
+
+      def imds_token(resource)
+        uri = token_uri("#{imds_host}#{IMDS_PATH}", "2018-02-01", resource)
+        # A short connect timeout: off Azure, IMDS is simply absent and the chain should move on.
+        response = HTTP.get(uri, headers: { "Metadata" => "true" }, open_timeout: 1,
+                                 retry_statuses: RETRY_STATUSES, proxy: false)
+        if response.status == 400
+          raise CredentialUnavailableError, "#{client_id ? "The requested identity has" : "No identity has"} " \
+                                            "been assigned to this resource (IMDS: #{response.error_message})"
+        end
+        AccessToken.from_response(response, "Azure Instance Metadata Service")
+      rescue HTTP::UnreachableError => e
+        raise CredentialUnavailableError, "No managed identity endpoint found (IMDS at #{e.message})"
+      end
+
+      # Same override as azure-identity, which is handy for tests.
+      def imds_host
+        (ENV["AZURE_POD_IDENTITY_AUTHORITY_HOST"].presence || IMDS_HOST).chomp("/")
+      end
+    end
+  end
+end

data/lib/rails_pg/managed_identity/postgresql_database_tasks.rb

@@ -0,0 +1,19 @@
+# frozen_string_literal: true
+
+module RailsPg
+  module ManagedIdentity
+    # Runs psql and pg_dump (db:schema:load and db:schema:dump with structure.sql) with the access
+    # token as PGPASSWORD. Registered for PostgreSQL databases by the Railtie.
+    class PostgreSQLDatabaseTasks < ActiveRecord::Tasks::PostgreSQLDatabaseTasks
+      private
+
+      def psql_env
+        env = super
+        if (provider = ManagedIdentity.token_provider_for(configuration_hash))
+          env["PGPASSWORD"] = provider.token
+        end
+        env
+      end
+    end
+  end
+end

data/lib/rails_pg/managed_identity/railtie.rb

@@ -0,0 +1,14 @@
+# frozen_string_literal: true
+
+module RailsPg
+  module ManagedIdentity
+    # Hooks into the database tasks only when rake tasks are loaded, so app and job processes never
+    # load the task code. Rails runs schema maintenance through rake too (bin/rails db:test:prepare).
+    class Railtie < ::Rails::Railtie
+      rake_tasks do
+        require_relative "postgresql_database_tasks"
+        ActiveRecord::Tasks::DatabaseTasks.register_task(/postgresql/, PostgreSQLDatabaseTasks)
+      end
+    end
+  end
+end

data/lib/rails_pg/managed_identity/token_provider.rb

@@ -0,0 +1,65 @@
+# frozen_string_literal: true
+
+module RailsPg
+  module ManagedIdentity
+    # Hands out access tokens for one credential and scope. A token is reused until it is
+    # ManagedIdentity.refresh_offset seconds from expiry, then replaced. If replacing it fails while it
+    # is still valid, the old token keeps being used. Shared by every connection using the same
+    # database configuration, so it is thread-safe and fetches at most one token at a time.
+    class TokenProvider
+      # While the cached token is usable, ask for a new one at most this often, whether the last
+      # request succeeded or failed (azure-identity does the same).
+      REFRESH_RETRY_DELAY = 30
+      # Never hand out a token that's about to expire mid-handshake.
+      MINIMUM_VALIDITY = 10
+
+      attr_reader :credential, :scope
+
+      def initialize(credential, scope:)
+        @credential = credential
+        @scope = scope
+        @mutex = Mutex.new
+        @access_token = nil
+        @requested_at = nil
+      end
+
+      def token
+        @mutex.synchronize do
+          refresh if refresh_due?
+          @access_token.token
+        end
+      end
+
+      private
+
+      def refresh_due?
+        return true unless usable?
+        return false if @access_token.expires_in > ManagedIdentity.refresh_offset
+
+        Time.now.to_i - @requested_at >= REFRESH_RETRY_DELAY
+      end
+
+      def refresh
+        @requested_at = Time.now.to_i
+        @access_token = credential.get_token(scope)
+        ManagedIdentity.logger&.info do
+          "[#{self.class.name}] Acquired access token for #{scope}, #{@access_token.inspect}"
+        end
+      rescue StandardError => e
+        # Whatever went wrong (a garbled response, a malformed endpoint URL, a failed `az` spawn),
+        # it's a failed token request: fall back to the cached token, or fail the connection attempt.
+        error = e.is_a?(Error) ? e : AuthenticationError.new("#{e.class}: #{e.message}")
+        raise error unless usable?
+
+        ManagedIdentity.logger&.warn do
+          "[#{self.class.name}] Couldn't refresh access token, using the cached one " \
+            "(#{@access_token.inspect}): #{error.message}"
+        end
+      end
+
+      def usable?
+        @access_token && @access_token.expires_in > MINIMUM_VALIDITY
+      end
+    end
+  end
+end

data/lib/rails_pg/managed_identity/version.rb

@@ -0,0 +1,7 @@
+# frozen_string_literal: true
+
+module RailsPg
+  module ManagedIdentity
+    VERSION = "0.1.0"
+  end
+end

data/lib/rails_pg/managed_identity/workload_identity_credential.rb

@@ -0,0 +1,38 @@
+# frozen_string_literal: true
+
+module RailsPg
+  module ManagedIdentity
+    # Microsoft Entra Workload ID on Kubernetes (AKS). The workload identity webhook injects
+    # AZURE_TENANT_ID, AZURE_CLIENT_ID and AZURE_FEDERATED_TOKEN_FILE; the projected service account
+    # token in that file is exchanged for an access token.
+    class WorkloadIdentityCredential
+      JWT_BEARER = "urn:ietf:params:oauth:client-assertion-type:jwt-bearer"
+
+      def initialize(client_id: nil)
+        @client_id = client_id
+      end
+
+      def get_token(scope)
+        tenant_id, token_file = ENV.values_at("AZURE_TENANT_ID", "AZURE_FEDERATED_TOKEN_FILE")
+        client_id = @client_id.presence || ENV.fetch("AZURE_CLIENT_ID", nil)
+        if [tenant_id, client_id, token_file].any?(&:blank?)
+          raise CredentialUnavailableError, "WorkloadIdentityCredential requires AZURE_TENANT_ID, " \
+                                            "AZURE_CLIENT_ID and AZURE_FEDERATED_TOKEN_FILE"
+        end
+
+        EntraID.client_credentials_token(tenant_id:, client_id:, scope:,
+                                         client_assertion_type: JWT_BEARER,
+                                         client_assertion: read_assertion(token_file))
+      end
+
+      private
+
+      # Kubernetes rotates the projected token in place, so it's read again for every exchange.
+      def read_assertion(token_file)
+        File.read(token_file).strip
+      rescue SystemCallError => e
+        raise AuthenticationError, "Can't read the federated token file #{token_file}: #{e.message}"
+      end
+    end
+  end
+end

data/lib/rails_pg/managed_identity.rb

@@ -0,0 +1,121 @@
+# frozen_string_literal: true
+
+require "active_record"
+require "active_support/core_ext/enumerable"
+require "active_support/core_ext/object/blank"
+require "active_support/core_ext/string/inflections"
+
+require_relative "managed_identity/version"
+
+module RailsPg
+  # Microsoft Entra authentication (managed identity, workload identity, ...) for Azure
+  # Database for PostgreSQL flexible server. Enable it per database in config/database.yml:
+  #
+  #   production:
+  #     adapter: postgresql
+  #     host: my-server.postgres.database.azure.com
+  #     database: my_app
+  #     username: my-app-identity   # the Entra principal's PostgreSQL role
+  #     sslmode: require
+  #     azure_managed_identity: true
+  #
+  # The password is then an access token resolved each time a connection is opened, rather than
+  # a value fixed when database.yml is loaded. Tokens are cached and replaced shortly before they
+  # expire.
+  module ManagedIdentity
+    # Raised when no access token can be obtained. A ConnectionNotEstablished, so Active Record
+    # handles it like any other failed connection attempt.
+    class Error < ActiveRecord::ConnectionNotEstablished; end
+
+    # The credential doesn't apply to this environment (e.g. no managed identity endpoint), so
+    # DefaultAzureCredential moves on to the next one.
+    class CredentialUnavailableError < Error; end
+
+    # The credential applies to this environment, but the token request failed.
+    class AuthenticationError < Error; end
+
+    CONFIG_KEY = :azure_managed_identity
+    # Azure Database for PostgreSQL's Microsoft Entra resource (public cloud).
+    DEFAULT_SCOPE = "https://ossrdbms-aad.database.windows.net/.default"
+    OPTIONS = %i[client_id credential scope].freeze
+
+    @providers = {}
+    @providers_mutex = Mutex.new
+    @refresh_offset = 300
+
+    class << self
+      attr_writer :logger
+
+      # Seconds before expiry at which a cached token is replaced. azure-identity uses 5 minutes.
+      attr_accessor :refresh_offset
+
+      def logger
+        @logger || ActiveRecord::Base.logger
+      end
+
+      # Returns the TokenProvider for a database configuration hash, or nil when the
+      # configuration doesn't enable Microsoft Entra authentication. Configurations with the same
+      # options share a provider, and therefore a cached token.
+      def token_provider_for(config)
+        options = parse_options(config[CONFIG_KEY])
+        return unless options
+
+        @providers_mutex.synchronize do
+          @providers[options] ||= TokenProvider.new(
+            build_credential(options.fetch(:credential, "default"), client_id: options[:client_id]),
+            scope: options.fetch(:scope, DEFAULT_SCOPE)
+          )
+        end
+      end
+
+      # Forgets all cached tokens.
+      def reset!
+        @providers_mutex.synchronize { @providers.clear }
+      end
+
+      private
+
+      def parse_options(value)
+        case value
+        when nil, false, "false", "" then nil
+        when true, "true" then {}
+        when Hash
+          options = value.to_h.transform_keys(&:to_sym).compact_blank
+          unknown = options.keys - OPTIONS
+          raise ArgumentError, "Unknown #{CONFIG_KEY} option(s): #{unknown.join(", ")}" if unknown.any?
+
+          options
+        else
+          raise ArgumentError, "#{CONFIG_KEY} must be true, false or a hash of options, got #{value.inspect}"
+        end
+      end
+
+      def build_credential(name, client_id:)
+        case name.to_s
+        when "default" then DefaultAzureCredential.new(client_id:)
+        when "managed_identity" then ManagedIdentityCredential.new(client_id:)
+        when "workload_identity" then WorkloadIdentityCredential.new(client_id:)
+        when "environment" then EnvironmentCredential.new
+        when "azure_cli" then AzureCliCredential.new
+        else
+          raise ArgumentError, "Unknown #{CONFIG_KEY} credential #{name.inspect}; expected one of " \
+                               "default, managed_identity, workload_identity, environment, azure_cli"
+        end
+      end
+    end
+  end
+end
+
+require_relative "managed_identity/access_token"
+require_relative "managed_identity/http"
+require_relative "managed_identity/entra_id"
+require_relative "managed_identity/environment_credential"
+require_relative "managed_identity/workload_identity_credential"
+require_relative "managed_identity/managed_identity_credential"
+require_relative "managed_identity/azure_cli_credential"
+require_relative "managed_identity/default_azure_credential"
+require_relative "managed_identity/token_provider"
+require_relative "managed_identity/active_record_extensions"
+
+# Runs when the Rails application class is defined, whether this gem was loaded before or after Rails.
+ActiveSupport.on_load(:before_configuration) { require_relative "managed_identity/railtie" }

data/sig/rails_pg/managed_identity.rbs

@@ -0,0 +1,80 @@
+module RailsPg
+  module ManagedIdentity
+    VERSION: String
+    CONFIG_KEY: Symbol
+    DEFAULT_SCOPE: String
+
+    class Error < ActiveRecord::ConnectionNotEstablished
+    end
+
+    class CredentialUnavailableError < Error
+    end
+
+    class AuthenticationError < Error
+    end
+
+    interface _Credential
+      def get_token: (String scope) -> AccessToken
+    end
+
+    attr_accessor self.refresh_offset: Integer
+    attr_writer self.logger: Logger?
+
+    def self.logger: () -> Logger?
+    def self.token_provider_for: (Hash[Symbol, untyped] config) -> TokenProvider?
+    def self.reset!: () -> void
+
+    class AccessToken
+      attr_reader token: String
+      attr_reader expires_on: Integer
+
+      def initialize: (token: String, expires_on: Integer) -> void
+      def expires_in: () -> Integer
+    end
+
+    class TokenProvider
+      REFRESH_RETRY_DELAY: Integer
+      MINIMUM_VALIDITY: Integer
+
+      attr_reader credential: _Credential
+      attr_reader scope: String
+
+      def initialize: (_Credential credential, scope: String) -> void
+      def token: () -> String
+    end
+
+    class DefaultAzureCredential
+      include _Credential
+
+      def initialize: (?client_id: String?) -> void
+    end
+
+    class ManagedIdentityCredential
+      include _Credential
+
+      def initialize: (?client_id: String?) -> void
+    end
+
+    class WorkloadIdentityCredential
+      include _Credential
+
+      def initialize: (?client_id: String?) -> void
+    end
+
+    class EnvironmentCredential
+      include _Credential
+    end
+
+    class AzureCliCredential
+      include _Credential
+
+      def initialize: (?process_timeout: Numeric) -> void
+    end
+
+    class PostgreSQLDatabaseTasks < ActiveRecord::Tasks::PostgreSQLDatabaseTasks
+    end
+
+    class Railtie < ::Rails::Railtie
+    end
+  end
+end

metadata

@@ -0,0 +1,102 @@
+--- !ruby/object:Gem::Specification
+name: rails-pg-managed-identity
+version: !ruby/object:Gem::Version
+  version: 0.1.0
+platform: ruby
+authors:
+- Frederik Erbs Spang Thomsen
+bindir: exe
+cert_chain: []
+date: 1980-01-02 00:00:00.000000000 Z
+dependencies:
+- !ruby/object:Gem::Dependency
+  name: activerecord
+  requirement: !ruby/object:Gem::Requirement
+    requirements:
+    - - ">="
+      - !ruby/object:Gem::Version
+        version: '7.2'
+    - - "<"
+      - !ruby/object:Gem::Version
+        version: '9'
+  type: :runtime
+  prerelease: false
+  version_requirements: !ruby/object:Gem::Requirement
+    requirements:
+    - - ">="
+      - !ruby/object:Gem::Version
+        version: '7.2'
+    - - "<"
+      - !ruby/object:Gem::Version
+        version: '9'
+- !ruby/object:Gem::Dependency
+  name: pg
+  requirement: !ruby/object:Gem::Requirement
+    requirements:
+    - - ">="
+      - !ruby/object:Gem::Version
+        version: '1.5'
+  type: :runtime
+  prerelease: false
+  version_requirements: !ruby/object:Gem::Requirement
+    requirements:
+    - - ">="
+      - !ruby/object:Gem::Version
+        version: '1.5'
+description: Connects Active Record's PostgreSQL adapter to Azure Database for PostgreSQL
+  flexible server with Microsoft Entra access tokens from managed identity, workload
+  identity, a service principal or the Azure CLI. Tokens are fetched when connections
+  are opened and refreshed before they expire, instead of a static password in database.yml.
+email:
+- fst@last.dk
+executables: []
+extensions: []
+extra_rdoc_files: []
+files:
+- CHANGELOG.md
+- CODE_OF_CONDUCT.md
+- LICENSE.txt
+- README.md
+- Rakefile
+- gemfiles/rails_7.2.gemfile
+- gemfiles/rails_8.0.gemfile
+- gemfiles/rails_8.1.gemfile
+- lib/rails-pg-managed-identity.rb
+- lib/rails_pg/managed_identity.rb
+- lib/rails_pg/managed_identity/access_token.rb
+- lib/rails_pg/managed_identity/active_record_extensions.rb
+- lib/rails_pg/managed_identity/azure_cli_credential.rb
+- lib/rails_pg/managed_identity/default_azure_credential.rb
+- lib/rails_pg/managed_identity/entra_id.rb
+- lib/rails_pg/managed_identity/environment_credential.rb
+- lib/rails_pg/managed_identity/http.rb
+- lib/rails_pg/managed_identity/managed_identity_credential.rb
+- lib/rails_pg/managed_identity/postgresql_database_tasks.rb
+- lib/rails_pg/managed_identity/railtie.rb
+- lib/rails_pg/managed_identity/token_provider.rb
+- lib/rails_pg/managed_identity/version.rb
+- lib/rails_pg/managed_identity/workload_identity_credential.rb
+- sig/rails_pg/managed_identity.rbs
+licenses:
+- MIT
+metadata:
+  rubygems_mfa_required: 'true'
+rdoc_options: []
+require_paths:
+- lib
+required_ruby_version: !ruby/object:Gem::Requirement
+  requirements:
+  - - ">="
+    - !ruby/object:Gem::Version
+      version: 3.3.0
+required_rubygems_version: !ruby/object:Gem::Requirement
+  requirements:
+  - - ">="
+    - !ruby/object:Gem::Version
+      version: '0'
+requirements: []
+rubygems_version: 4.0.21
+specification_version: 4
+summary: Azure managed identity (Microsoft Entra) authentication for Rails and Azure
+  Database for PostgreSQL
+test_files: []